vectordb-bench 0.0.18__py3-none-any.whl → 0.0.20__py3-none-any.whl

vectordb_bench/backend/clients/milvus/cli.py

@@ -3,6 +3,7 @@ from typing import Annotated, TypedDict, Unpack
 import click
 from pydantic import SecretStr
 
+from vectordb_bench.backend.clients import DB
 from vectordb_bench.cli.cli import (
     CommonTypedDict,
     HNSWFlavor3,
@@ -10,33 +11,41 @@ from vectordb_bench.cli.cli import (
     cli,
     click_parameter_decorators_from_typed_dict,
     run,
-
 )
-from vectordb_bench.backend.clients import DB
 
 DBTYPE = DB.Milvus
 
 
 class MilvusTypedDict(TypedDict):
     uri: Annotated[
-        str, click.option("--uri", type=str, help="uri connection string", required=True)
+        str,
+        click.option("--uri", type=str, help="uri connection string", required=True),
+    ]
+    user_name: Annotated[
+        str | None,
+        click.option("--user-name", type=str, help="Db username", required=False),
+    ]
+    password: Annotated[
+        str | None,
+        click.option("--password", type=str, help="Db password", required=False),
     ]
 
 
-class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict):
-    ...
+class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict): ...
 
 
 @cli.command()
 @click_parameter_decorators_from_typed_dict(MilvusAutoIndexTypedDict)
 def MilvusAutoIndex(**parameters: Unpack[MilvusAutoIndexTypedDict]):
-    from .config import MilvusConfig, AutoIndexConfig
+    from .config import AutoIndexConfig, MilvusConfig
 
     run(
         db=DBTYPE,
         db_config=MilvusConfig(
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]),
         ),
         db_case_config=AutoIndexConfig(),
         **parameters,
@@ -46,33 +55,36 @@ def MilvusAutoIndex(**parameters: Unpack[MilvusAutoIndexTypedDict]):
 @cli.command()
 @click_parameter_decorators_from_typed_dict(MilvusAutoIndexTypedDict)
 def MilvusFlat(**parameters: Unpack[MilvusAutoIndexTypedDict]):
-    from .config import MilvusConfig, FLATConfig
+    from .config import FLATConfig, MilvusConfig
 
     run(
         db=DBTYPE,
         db_config=MilvusConfig(
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]),
         ),
         db_case_config=FLATConfig(),
         **parameters,
     )
 
 
-class MilvusHNSWTypedDict(CommonTypedDict, MilvusTypedDict, HNSWFlavor3):
-    ...
+class MilvusHNSWTypedDict(CommonTypedDict, MilvusTypedDict, HNSWFlavor3): ...
 
 
 @cli.command()
 @click_parameter_decorators_from_typed_dict(MilvusHNSWTypedDict)
 def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]):
-    from .config import MilvusConfig, HNSWConfig
+    from .config import HNSWConfig, MilvusConfig
 
     run(
         db=DBTYPE,
         db_config=MilvusConfig(
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
         ),
         db_case_config=HNSWConfig(
             M=parameters["m"],
@@ -83,20 +95,21 @@ def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]):
     )
 
 
-class MilvusIVFFlatTypedDict(CommonTypedDict, MilvusTypedDict, IVFFlatTypedDictN):
-    ...
+class MilvusIVFFlatTypedDict(CommonTypedDict, MilvusTypedDict, IVFFlatTypedDictN): ...
 
 
 @cli.command()
 @click_parameter_decorators_from_typed_dict(MilvusIVFFlatTypedDict)
 def MilvusIVFFlat(**parameters: Unpack[MilvusIVFFlatTypedDict]):
-    from .config import MilvusConfig, IVFFlatConfig
+    from .config import IVFFlatConfig, MilvusConfig
 
     run(
         db=DBTYPE,
         db_config=MilvusConfig(
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]),
         ),
         db_case_config=IVFFlatConfig(
             nlist=parameters["nlist"],
@@ -109,13 +122,15 @@ def MilvusIVFFlat(**parameters: Unpack[MilvusIVFFlatTypedDict]):
 @cli.command()
 @click_parameter_decorators_from_typed_dict(MilvusIVFFlatTypedDict)
 def MilvusIVFSQ8(**parameters: Unpack[MilvusIVFFlatTypedDict]):
-    from .config import MilvusConfig, IVFSQ8Config
+    from .config import IVFSQ8Config, MilvusConfig
 
     run(
         db=DBTYPE,
         db_config=MilvusConfig(
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]),
         ),
         db_case_config=IVFSQ8Config(
             nlist=parameters["nlist"],
@@ -126,23 +141,21 @@ def MilvusIVFSQ8(**parameters: Unpack[MilvusIVFFlatTypedDict]):
 
 
 class MilvusDISKANNTypedDict(CommonTypedDict, MilvusTypedDict):
-    search_list: Annotated[
-        str, click.option("--search-list",
-                          type=int,
-                          required=True)
-    ]
+    search_list: Annotated[str, click.option("--search-list", type=int, required=True)]
 
 
 @cli.command()
 @click_parameter_decorators_from_typed_dict(MilvusDISKANNTypedDict)
 def MilvusDISKANN(**parameters: Unpack[MilvusDISKANNTypedDict]):
-    from .config import MilvusConfig, DISKANNConfig
+    from .config import DISKANNConfig, MilvusConfig
 
     run(
         db=DBTYPE,
         db_config=MilvusConfig(
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]),
         ),
         db_case_config=DISKANNConfig(
             search_list=parameters["search_list"],
@@ -153,27 +166,24 @@ def MilvusDISKANN(**parameters: Unpack[MilvusDISKANNTypedDict]):
 
 class MilvusGPUIVFTypedDict(CommonTypedDict, MilvusTypedDict, MilvusIVFFlatTypedDict):
     cache_dataset_on_device: Annotated[
-        str, click.option("--cache-dataset-on-device",
-                          type=str,
-                          required=True)
-    ]
-    refine_ratio: Annotated[
-        str, click.option("--refine-ratio",
-                          type=float,
-                          required=True)
+        str,
+        click.option("--cache-dataset-on-device", type=str, required=True),
     ]
+    refine_ratio: Annotated[str, click.option("--refine-ratio", type=float, required=True)]
 
 
 @cli.command()
 @click_parameter_decorators_from_typed_dict(MilvusGPUIVFTypedDict)
 def MilvusGPUIVFFlat(**parameters: Unpack[MilvusGPUIVFTypedDict]):
-    from .config import MilvusConfig, GPUIVFFlatConfig
+    from .config import GPUIVFFlatConfig, MilvusConfig
 
     run(
         db=DBTYPE,
         db_config=MilvusConfig(
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]),
         ),
         db_case_config=GPUIVFFlatConfig(
             nlist=parameters["nlist"],
@@ -185,29 +195,28 @@ def MilvusGPUIVFFlat(**parameters: Unpack[MilvusGPUIVFTypedDict]):
     )
 
 
-class MilvusGPUIVFPQTypedDict(CommonTypedDict, MilvusTypedDict, MilvusIVFFlatTypedDict, MilvusGPUIVFTypedDict):
-    m: Annotated[
-        str, click.option("--m",
-                          type=int, help="hnsw m",
-                          required=True)
-    ]
-    nbits: Annotated[
-        str, click.option("--nbits",
-                          type=int,
-                          required=True)
-    ]
+class MilvusGPUIVFPQTypedDict(
+    CommonTypedDict,
+    MilvusTypedDict,
+    MilvusIVFFlatTypedDict,
+    MilvusGPUIVFTypedDict,
+):
+    m: Annotated[str, click.option("--m", type=int, help="hnsw m", required=True)]
+    nbits: Annotated[str, click.option("--nbits", type=int, required=True)]
 
 
 @cli.command()
 @click_parameter_decorators_from_typed_dict(MilvusGPUIVFPQTypedDict)
 def MilvusGPUIVFPQ(**parameters: Unpack[MilvusGPUIVFPQTypedDict]):
-    from .config import MilvusConfig, GPUIVFPQConfig
+    from .config import GPUIVFPQConfig, MilvusConfig
 
     run(
         db=DBTYPE,
         db_config=MilvusConfig(
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]),
         ),
         db_case_config=GPUIVFPQConfig(
             nlist=parameters["nlist"],
@@ -223,57 +232,30 @@ def MilvusGPUIVFPQ(**parameters: Unpack[MilvusGPUIVFPQTypedDict]):
 
 class MilvusGPUCAGRATypedDict(CommonTypedDict, MilvusTypedDict, MilvusGPUIVFTypedDict):
     intermediate_graph_degree: Annotated[
-        str, click.option("--intermediate-graph-degree",
-                          type=int,
-                          required=True)
-    ]
-    graph_degree: Annotated[
-        str, click.option("--graph-degree",
-                          type=int,
-                          required=True)
-    ]
-    build_algo: Annotated[
-        str, click.option("--build_algo",
-                          type=str,
-                          required=True)
-    ]
-    team_size: Annotated[
-        str, click.option("--team-size",
-                          type=int,
-                          required=True)
-    ]
-    search_width: Annotated[
-        str, click.option("--search-width",
-                          type=int,
-                          required=True)
-    ]
-    itopk_size: Annotated[
-        str, click.option("--itopk-size",
-                          type=int,
-                          required=True)
-    ]
-    min_iterations: Annotated[
-        str, click.option("--min-iterations",
-                          type=int,
-                          required=True)
-    ]
-    max_iterations: Annotated[
-        str, click.option("--max-iterations",
-                          type=int,
-                          required=True)
+        str,
+        click.option("--intermediate-graph-degree", type=int, required=True),
     ]
+    graph_degree: Annotated[str, click.option("--graph-degree", type=int, required=True)]
+    build_algo: Annotated[str, click.option("--build_algo", type=str, required=True)]
+    team_size: Annotated[str, click.option("--team-size", type=int, required=True)]
+    search_width: Annotated[str, click.option("--search-width", type=int, required=True)]
+    itopk_size: Annotated[str, click.option("--itopk-size", type=int, required=True)]
+    min_iterations: Annotated[str, click.option("--min-iterations", type=int, required=True)]
+    max_iterations: Annotated[str, click.option("--max-iterations", type=int, required=True)]
 
 
 @cli.command()
 @click_parameter_decorators_from_typed_dict(MilvusGPUCAGRATypedDict)
 def MilvusGPUCAGRA(**parameters: Unpack[MilvusGPUCAGRATypedDict]):
-    from .config import MilvusConfig, GPUCAGRAConfig
+    from .config import GPUCAGRAConfig, MilvusConfig
 
     run(
         db=DBTYPE,
         db_config=MilvusConfig(
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]),
         ),
         db_case_config=GPUCAGRAConfig(
             intermediate_graph_degree=parameters["intermediate_graph_degree"],

vectordb_bench/backend/clients/milvus/config.py

@@ -1,12 +1,31 @@
-from pydantic import BaseModel, SecretStr
-from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
+from pydantic import BaseModel, SecretStr, validator
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
 
 
 class MilvusConfig(DBConfig):
     uri: SecretStr = "http://localhost:19530"
+    user: str | None = None
+    password: SecretStr | None = None
 
     def to_dict(self) -> dict:
-        return {"uri": self.uri.get_secret_value()}
+        return {
+            "uri": self.uri.get_secret_value(),
+            "user": self.user if self.user else None,
+            "password": self.password.get_secret_value() if self.password else None,
+        }
+
+    @validator("*")
+    def not_empty_field(cls, v: any, field: any):
+        if (
+            field.name in cls.common_short_configs()
+            or field.name in cls.common_long_configs()
+            or field.name in ["user", "password"]
+        ):
+            return v
+        if isinstance(v, str | SecretStr) and len(v) == 0:
+            raise ValueError("Empty string!")
+        return v
 
 
 class MilvusIndexConfig(BaseModel):
@@ -14,10 +33,14 @@ class MilvusIndexConfig(BaseModel):
 
     index: IndexType
     metric_type: MetricType | None = None
-    
+
     @property
     def is_gpu_index(self) -> bool:
-        return self.index in [IndexType.GPU_CAGRA, IndexType.GPU_IVF_FLAT, IndexType.GPU_IVF_PQ]
+        return self.index in [
+            IndexType.GPU_CAGRA,
+            IndexType.GPU_IVF_FLAT,
+            IndexType.GPU_IVF_PQ,
+        ]
 
     def parse_metric(self) -> str:
         if not self.metric_type:
@@ -99,7 +122,8 @@ class IVFFlatConfig(MilvusIndexConfig, DBCaseConfig):
             "metric_type": self.parse_metric(),
             "params": {"nprobe": self.nprobe},
         }
-        
+
+
 class IVFSQ8Config(MilvusIndexConfig, DBCaseConfig):
     nlist: int
     nprobe: int | None = None
@@ -196,7 +220,7 @@ class GPUCAGRAConfig(MilvusIndexConfig, DBCaseConfig):
     search_width: int = 4
     min_iterations: int = 0
     max_iterations: int = 0
-    build_algo: str = "IVF_PQ" # IVF_PQ; NN_DESCENT;
+    build_algo: str = "IVF_PQ"  # IVF_PQ; NN_DESCENT;
     cache_dataset_on_device: str
     refine_ratio: float | None = None
     index: IndexType = IndexType.GPU_CAGRA

vectordb_bench/backend/clients/milvus/milvus.py

@@ -2,19 +2,18 @@
 
 import logging
 import time
+from collections.abc import Iterable
 from contextlib import contextmanager
-from typing import Iterable
 
-from pymilvus import Collection, utility
-from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusException
+from pymilvus import Collection, CollectionSchema, DataType, FieldSchema, MilvusException, utility
 
-from ..api import VectorDB, IndexType
+from ..api import VectorDB
 from .config import MilvusIndexConfig
 
-
 log = logging.getLogger(__name__)
 
-MILVUS_LOAD_REQS_SIZE = 1.5 * 1024 *1024
+MILVUS_LOAD_REQS_SIZE = 1.5 * 1024 * 1024
+
 
 class Milvus(VectorDB):
     def __init__(
@@ -32,7 +31,7 @@ class Milvus(VectorDB):
         self.db_config = db_config
         self.case_config = db_case_config
         self.collection_name = collection_name
-        self.batch_size = int(MILVUS_LOAD_REQS_SIZE / (dim *4))
+        self.batch_size = int(MILVUS_LOAD_REQS_SIZE / (dim * 4))
 
         self._primary_field = "pk"
         self._scalar_field = "id"
@@ -40,6 +39,7 @@ class Milvus(VectorDB):
         self._index_name = "vector_idx"
 
         from pymilvus import connections
+
         connections.connect(**self.db_config, timeout=30)
         if drop_old and utility.has_collection(self.collection_name):
             log.info(f"{self.name} client drop_old collection: {self.collection_name}")
@@ -49,7 +49,7 @@ class Milvus(VectorDB):
             fields = [
                 FieldSchema(self._primary_field, DataType.INT64, is_primary=True),
                 FieldSchema(self._scalar_field, DataType.INT64),
-                FieldSchema(self._vector_field, DataType.FLOAT_VECTOR, dim=dim)
+                FieldSchema(self._vector_field, DataType.FLOAT_VECTOR, dim=dim),
             ]
 
             log.info(f"{self.name} create collection: {self.collection_name}")
@@ -66,8 +66,7 @@ class Milvus(VectorDB):
                 self.case_config.index_param(),
                 index_name=self._index_name,
             )
-            if kwargs.get("pre_load") is True:
-                self._pre_load(col)
+            col.load()
 
         connections.disconnect("default")
 
@@ -80,6 +79,7 @@ class Milvus(VectorDB):
             >>>     self.search_embedding()
         """
         from pymilvus import connections
+
         self.col: Collection | None = None
 
         connections.connect(**self.db_config, timeout=60)
@@ -90,16 +90,15 @@ class Milvus(VectorDB):
         connections.disconnect("default")
 
     def _optimize(self):
-        self._post_insert()
         log.info(f"{self.name} optimizing before search")
+        self._post_insert()
         try:
-            self.col.load()
+            self.col.load(refresh=True)
         except Exception as e:
             log.warning(f"{self.name} optimize error: {e}")
             raise e from None
 
     def _post_insert(self):
-        log.info(f"{self.name} post insert before optimize")
         try:
             self.col.flush()
             # wait for index done and load refresh
@@ -110,6 +109,7 @@ class Milvus(VectorDB):
             )
 
             utility.wait_for_index_building_complete(self.collection_name)
+
             def wait_index():
                 while True:
                     progress = utility.index_building_progress(self.collection_name)
@@ -122,18 +122,17 @@ class Milvus(VectorDB):
             # Skip compaction if use GPU indexType
             if self.case_config.is_gpu_index:
                 log.debug("skip compaction for gpu index type.")
-            else :
+            else:
                 try:
                     self.col.compact()
                     self.col.wait_for_compaction_completed()
                 except Exception as e:
                     log.warning(f"{self.name} compact error: {e}")
-                    if hasattr(e, 'code'):
-                        if e.code().name == 'PERMISSION_DENIED':
-                            log.warning(f"Skip compact due to permission denied.")
-                            pass
+                    if hasattr(e, "code"):
+                        if e.code().name == "PERMISSION_DENIED":
+                            log.warning("Skip compact due to permission denied.")
                     else:
-                        raise e
+                        raise e from e
                 wait_index()
         except Exception as e:
             log.warning(f"{self.name} optimize error: {e}")
@@ -158,7 +157,6 @@ class Milvus(VectorDB):
             log.warning(f"{self.name} pre load error: {e}")
             raise e from None
 
-
     def optimize(self):
         assert self.col, "Please call self.init() before"
         self._optimize()
@@ -166,7 +164,7 @@ class Milvus(VectorDB):
     def need_normalize_cosine(self) -> bool:
         """Wheather this database need to normalize dataset to support COSINE"""
         if self.case_config.is_gpu_index:
-            log.info(f"current gpu_index only supports IP / L2, cosine dataset need normalize.")
+            log.info("current gpu_index only supports IP / L2, cosine dataset need normalize.")
             return True
 
         return False
@@ -186,9 +184,9 @@ class Milvus(VectorDB):
             for batch_start_offset in range(0, len(embeddings), self.batch_size):
                 batch_end_offset = min(batch_start_offset + self.batch_size, len(embeddings))
                 insert_data = [
-                        metadata[batch_start_offset : batch_end_offset],
-                        metadata[batch_start_offset : batch_end_offset],
-                        embeddings[batch_start_offset : batch_end_offset],
+                    metadata[batch_start_offset:batch_end_offset],
+                    metadata[batch_start_offset:batch_end_offset],
+                    embeddings[batch_start_offset:batch_end_offset],
                 ]
                 res = self.col.insert(insert_data)
                 insert_count += len(res.primary_keys)
@@ -219,5 +217,4 @@ class Milvus(VectorDB):
         )
 
         # Organize results.
-        ret = [result.id for result in res[0]]
-        return ret
+        return [result.id for result in res[0]]

vectordb_bench/backend/clients/pgdiskann/cli.py

@@ -1,57 +1,63 @@
-import click
 import os
+from typing import Annotated, Unpack
+
+import click
 from pydantic import SecretStr
 
+from vectordb_bench.backend.clients import DB
+
 from ....cli.cli import (
     CommonTypedDict,
     cli,
     click_parameter_decorators_from_typed_dict,
     run,
 )
-from typing import Annotated, Optional, Unpack
-from vectordb_bench.backend.clients import DB
 
 
 class PgDiskAnnTypedDict(CommonTypedDict):
     user_name: Annotated[
-        str, click.option("--user-name", type=str, help="Db username", required=True)
+        str,
+        click.option("--user-name", type=str, help="Db username", required=True),
     ]
     password: Annotated[
         str,
-        click.option("--password",
-                     type=str,
-                     help="Postgres database password",
-                     default=lambda: os.environ.get("POSTGRES_PASSWORD", ""),
-                     show_default="$POSTGRES_PASSWORD",
-                     ),
+        click.option(
+            "--password",
+            type=str,
+            help="Postgres database password",
+            default=lambda: os.environ.get("POSTGRES_PASSWORD", ""),
+            show_default="$POSTGRES_PASSWORD",
+        ),
     ]
 
-    host: Annotated[
-        str, click.option("--host", type=str, help="Db host", required=True)
-    ]
-    db_name: Annotated[
-        str, click.option("--db-name", type=str, help="Db name", required=True)
-    ]
+    host: Annotated[str, click.option("--host", type=str, help="Db host", required=True)]
+    db_name: Annotated[str, click.option("--db-name", type=str, help="Db name", required=True)]
     max_neighbors: Annotated[
         int,
         click.option(
-            "--max-neighbors", type=int, help="PgDiskAnn max neighbors",
+            "--max-neighbors",
+            type=int,
+            help="PgDiskAnn max neighbors",
         ),
     ]
     l_value_ib: Annotated[
         int,
         click.option(
-            "--l-value-ib", type=int, help="PgDiskAnn l_value_ib",
+            "--l-value-ib",
+            type=int,
+            help="PgDiskAnn l_value_ib",
         ),
     ]
     l_value_is: Annotated[
         float,
         click.option(
-            "--l-value-is", type=float, help="PgDiskAnn l_value_is",
+            "--l-value-is",
+            type=float,
+            help="PgDiskAnn l_value_is",
         ),
     ]
     maintenance_work_mem: Annotated[
-        Optional[str],
+        str | None,
         click.option(
             "--maintenance-work-mem",
             type=str,
@@ -63,7 +69,7 @@ class PgDiskAnnTypedDict(CommonTypedDict):
         ),
     ]
     max_parallel_workers: Annotated[
-        Optional[int],
+        int | None,
         click.option(
             "--max-parallel-workers",
             type=int,
@@ -72,6 +78,7 @@ class PgDiskAnnTypedDict(CommonTypedDict):
         ),
     ]
 
+
 @cli.command()
 @click_parameter_decorators_from_typed_dict(PgDiskAnnTypedDict)
 def PgDiskAnn(
@@ -96,4 +103,4 @@ def PgDiskAnn(
             maintenance_work_mem=parameters["maintenance_work_mem"],
         ),
         **parameters,
-    )
+    )