vectordb-bench 0.0.19__py3-none-any.whl → 0.0.21__py3-none-any.whl

vectordb_bench/backend/clients/milvus/cli.py

@@ -1,8 +1,9 @@
-from typing import Annotated, TypedDict, Unpack, Optional
+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,33 @@ 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[
-        Optional[str], click.option("--user-name", type=str, help="Db username", required=False)
+        str | None,
+        click.option("--user-name", type=str, help="Db username", required=False),
     ]
     password: Annotated[
-        Optional[str], click.option("--password", type=str, help="Db password", required=False)
+        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,
@@ -54,7 +55,7 @@ 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,
@@ -69,14 +70,13 @@ def MilvusFlat(**parameters: Unpack[MilvusAutoIndexTypedDict]):
     )
 
 
-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,
@@ -95,14 +95,13 @@ 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,
@@ -123,7 +122,7 @@ 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,
@@ -142,17 +141,13 @@ 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,
@@ -171,21 +166,16 @@ 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,
@@ -205,23 +195,20 @@ 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,
@@ -245,51 +232,22 @@ 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,

vectordb_bench/backend/clients/milvus/config.py

@@ -1,5 +1,6 @@
 from pydantic import BaseModel, SecretStr, validator
-from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
 
 
 class MilvusConfig(DBConfig):
@@ -15,10 +16,14 @@ class MilvusConfig(DBConfig):
         }
 
     @validator("*")
-    def not_empty_field(cls, v, field):
-        if field.name in cls.common_short_configs() or field.name in cls.common_long_configs() or field.name in ["user", "password"]:
+    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:
+        if isinstance(v, str | SecretStr) and len(v) == 0:
             raise ValueError("Empty string!")
         return v
 
@@ -28,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:
@@ -113,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
@@ -210,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
 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}")
@@ -79,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)
@@ -108,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)
@@ -120,51 +122,30 @@ 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':
+                    if hasattr(e, "code"):
+                        if e.code().name == "PERMISSION_DENIED":
                             log.warning("Skip compact due to permission denied.")
-                            pass
                     else:
-                        raise e
+                        raise e from e
                 wait_index()
         except Exception as e:
             log.warning(f"{self.name} optimize error: {e}")
             raise e from None
 
-    def ready_to_load(self):
-        assert self.col, "Please call self.init() before"
-        self._pre_load(self.col)
-
-    def _pre_load(self, coll: Collection):
-        try:
-            if not coll.has_index(index_name=self._index_name):
-                log.info(f"{self.name} create index")
-                coll.create_index(
-                    self._vector_field,
-                    self.case_config.index_param(),
-                    index_name=self._index_name,
-                )
-            coll.load()
-            log.info(f"{self.name} load")
-        except Exception as e:
-            log.warning(f"{self.name} pre load error: {e}")
-            raise e from None
-
-
-    def optimize(self):
+    def optimize(self, data_size: int | None = None):
         assert self.col, "Please call self.init() before"
         self._optimize()
 
     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
@@ -184,9 +165,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)
@@ -217,5 +198,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,
-    )
+    )