vectordb-bench 0.0.28__py3-none-any.whl → 0.0.30__py3-none-any.whl

vectordb_bench/backend/clients/lancedb/cli.py

@@ -58,10 +58,46 @@ def LanceDBAutoIndex(**parameters: Unpack[LanceDBTypedDict]):
     )
 
 
+class LanceDBIVFPQTypedDict(CommonTypedDict, LanceDBTypedDict):
+    num_partitions: Annotated[
+        int,
+        click.option(
+            "--num-partitions",
+            type=int,
+            default=0,
+            help="Number of partitions for IVFPQ index, unset = use LanceDB default",
+        ),
+    ]
+    num_sub_vectors: Annotated[
+        int,
+        click.option(
+            "--num-sub-vectors",
+            type=int,
+            default=0,
+            help="Number of sub-vectors for IVFPQ index, unset = use LanceDB default",
+        ),
+    ]
+    nbits: Annotated[
+        int,
+        click.option(
+            "--nbits",
+            type=int,
+            default=8,
+            help="Number of bits for IVFPQ index (must be 4 or 8), unset = use LanceDB default",
+        ),
+    ]
+    nprobes: Annotated[
+        int,
+        click.option(
+            "--nprobes", type=int, default=0, help="Number of probes for IVFPQ search, unset = use LanceDB default"
+        ),
+    ]
+
+
 @cli.command()
-@click_parameter_decorators_from_typed_dict(LanceDBTypedDict)
-def LanceDBIVFPQ(**parameters: Unpack[LanceDBTypedDict]):
-    from .config import LanceDBConfig, _lancedb_case_config
+@click_parameter_decorators_from_typed_dict(LanceDBIVFPQTypedDict)
+def LanceDBIVFPQ(**parameters: Unpack[LanceDBIVFPQTypedDict]):
+    from .config import LanceDBConfig, LanceDBIndexConfig
 
     run(
         db=DB.LanceDB,
@@ -70,15 +106,29 @@ def LanceDBIVFPQ(**parameters: Unpack[LanceDBTypedDict]):
             uri=parameters["uri"],
             token=SecretStr(parameters["token"]) if parameters.get("token") else None,
         ),
-        db_case_config=_lancedb_case_config.get(IndexType.IVFPQ)(),
+        db_case_config=LanceDBIndexConfig(
+            index=IndexType.IVFPQ,
+            num_partitions=parameters["num_partitions"],
+            num_sub_vectors=parameters["num_sub_vectors"],
+            nbits=parameters["nbits"],
+            nprobes=parameters["nprobes"],
+        ),
         **parameters,
     )
 
 
+class LanceDBHNSWTypedDict(CommonTypedDict, LanceDBTypedDict):
+    m: Annotated[int, click.option("--m", type=int, default=0, help="HNSW parameter m")]
+    ef_construction: Annotated[
+        int, click.option("--ef-construction", type=int, default=0, help="HNSW parameter ef_construction")
+    ]
+    ef: Annotated[int, click.option("--ef", type=int, default=0, help="HNSW search parameter ef")]
+
+
 @cli.command()
-@click_parameter_decorators_from_typed_dict(LanceDBTypedDict)
-def LanceDBHNSW(**parameters: Unpack[LanceDBTypedDict]):
-    from .config import LanceDBConfig, _lancedb_case_config
+@click_parameter_decorators_from_typed_dict(LanceDBHNSWTypedDict)
+def LanceDBHNSW(**parameters: Unpack[LanceDBHNSWTypedDict]):
+    from .config import LanceDBConfig, LanceDBHNSWIndexConfig
 
     run(
         db=DB.LanceDB,
@@ -87,6 +137,10 @@ def LanceDBHNSW(**parameters: Unpack[LanceDBTypedDict]):
             uri=parameters["uri"],
             token=SecretStr(parameters["token"]) if parameters.get("token") else None,
         ),
-        db_case_config=_lancedb_case_config.get(IndexType.HNSW)(),
+        db_case_config=LanceDBHNSWIndexConfig(
+            m=parameters["m"],
+            ef_construction=parameters["ef_construction"],
+            ef=parameters["ef"],
+        ),
         **parameters,
     )

vectordb_bench/backend/clients/lancedb/config.py

@@ -25,6 +25,7 @@ class LanceDBIndexConfig(BaseModel, DBCaseConfig):
     nbits: int = 8  # Must be 4 or 8
     sample_rate: int = 256
     max_iterations: int = 50
+    nprobes: int = 0
 
     def index_param(self) -> dict:
         if self.index not in [
@@ -52,7 +53,11 @@ class LanceDBIndexConfig(BaseModel, DBCaseConfig):
         return params
 
     def search_param(self) -> dict:
-        pass
+        params = {}
+        if self.nprobes > 0:
+            params["nprobes"] = self.nprobes
+
+        return params
 
     def parse_metric(self) -> str:
         if self.metric_type in [MetricType.L2, MetricType.COSINE]:
@@ -81,6 +86,7 @@ class LanceDBHNSWIndexConfig(LanceDBIndexConfig):
     index: IndexType = IndexType.HNSW
     m: int = 0
     ef_construction: int = 0
+    ef: int = 0
 
     def index_param(self) -> dict:
         params = LanceDBIndexConfig.index_param(self)
@@ -94,6 +100,13 @@ class LanceDBHNSWIndexConfig(LanceDBIndexConfig):
 
         return params
 
+    def search_param(self) -> dict:
+        params = {}
+        if self.ef != 0:
+            params = {"ef": self.ef}
+
+        return params
+
 
 _lancedb_case_config = {
     IndexType.IVFPQ: LanceDBIndexConfig,

vectordb_bench/backend/clients/lancedb/lancedb.py

@@ -32,6 +32,10 @@ class LanceDB(VectorDB):
         self.table_name = collection_name
         self.dim = dim
         self.uri = db_config["uri"]
+        # avoid the search_param being called every time during the search process
+        self.search_config = db_case_config.search_param()
+
+        log.info(f"Search config: {self.search_config}")
 
         db = lancedb.connect(self.uri)
 
@@ -45,7 +49,7 @@ class LanceDB(VectorDB):
             db.open_table(self.table_name)
         except Exception:
             schema = pa.schema(
-                [pa.field("id", pa.int64()), pa.field("vector", pa.list_(pa.float64(), list_size=self.dim))]
+                [pa.field("id", pa.int64()), pa.field("vector", pa.list_(pa.float32(), list_size=self.dim))]
             )
             db.create_table(self.table_name, schema=schema, mode="overwrite")
 
@@ -77,14 +81,28 @@ class LanceDB(VectorDB):
         filters: dict | None = None,
     ) -> list[int]:
         if filters:
-            results = self.table.search(query).where(f"id >= {filters['id']}", prefilter=True).limit(k).to_list()
+            results = self.table.search(query).select(["id"]).where(f"id >= {filters['id']}", prefilter=True).limit(k)
+            if self.case_config.index == IndexType.IVFPQ and "nprobes" in self.search_config:
+                results = results.nprobes(self.search_config["nprobes"]).to_list()
+            elif self.case_config.index == IndexType.HNSW and "ef" in self.search_config:
+                results = results.ef(self.search_config["ef"]).to_list()
+            else:
+                results = results.to_list()
         else:
-            results = self.table.search(query).limit(k).to_list()
+            results = self.table.search(query).select(["id"]).limit(k)
+            if self.case_config.index == IndexType.IVFPQ and "nprobes" in self.search_config:
+                results = results.nprobes(self.search_config["nprobes"]).to_list()
+            elif self.case_config.index == IndexType.HNSW and "ef" in self.search_config:
+                results = results.ef(self.search_config["ef"]).to_list()
+            else:
+                results = results.to_list()
+
         return [int(result["id"]) for result in results]
 
     def optimize(self, data_size: int | None = None):
         if self.table and hasattr(self, "case_config") and self.case_config.index != IndexType.NONE:
             log.info(f"Creating index for LanceDB table ({self.table_name})")
+            log.info(f"Index parameters: {self.case_config.index_param()}")
             self.table.create_index(**self.case_config.index_param())
             # Better recall with IVF_PQ (though still bad) but breaks HNSW: https://github.com/lancedb/lancedb/issues/2369
             if self.case_config.index in (IndexType.IVFPQ, IndexType.AUTOINDEX):

vectordb_bench/backend/clients/memorydb/memorydb.py

@@ -9,10 +9,10 @@ import redis
 from redis import Redis
 from redis.cluster import RedisCluster
 from redis.commands.search.field import NumericField, TagField, VectorField
-from redis.commands.search.indexDefinition import IndexDefinition
+from redis.commands.search.indexDefinition import IndexDefinition, IndexType
 from redis.commands.search.query import Query
 
-from ..api import IndexType, VectorDB
+from ..api import VectorDB
 from .config import MemoryDBIndexConfig
 
 log = logging.getLogger(__name__)

vectordb_bench/backend/clients/milvus/cli.py

@@ -29,6 +29,17 @@ class MilvusTypedDict(TypedDict):
         str | None,
         click.option("--password", type=str, help="Db password", required=False),
     ]
+    num_shards: Annotated[
+        int,
+        click.option(
+            "--num-shards",
+            type=int,
+            help="Number of shards",
+            required=False,
+            default=1,
+            show_default=True,
+        ),
+    ]
 
 
 class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict): ...
@@ -45,7 +56,8 @@ def MilvusAutoIndex(**parameters: Unpack[MilvusAutoIndexTypedDict]):
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
             user=parameters["user_name"],
-            password=SecretStr(parameters["password"]),
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
         ),
         db_case_config=AutoIndexConfig(),
         **parameters,
@@ -63,7 +75,8 @@ def MilvusFlat(**parameters: Unpack[MilvusAutoIndexTypedDict]):
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
             user=parameters["user_name"],
-            password=SecretStr(parameters["password"]),
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
         ),
         db_case_config=FLATConfig(),
         **parameters,
@@ -85,6 +98,7 @@ def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]):
             uri=SecretStr(parameters["uri"]),
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
         ),
         db_case_config=HNSWConfig(
             M=parameters["m"],
@@ -109,7 +123,8 @@ def MilvusIVFFlat(**parameters: Unpack[MilvusIVFFlatTypedDict]):
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
             user=parameters["user_name"],
-            password=SecretStr(parameters["password"]),
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
         ),
         db_case_config=IVFFlatConfig(
             nlist=parameters["nlist"],
@@ -130,7 +145,8 @@ def MilvusIVFSQ8(**parameters: Unpack[MilvusIVFFlatTypedDict]):
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
             user=parameters["user_name"],
-            password=SecretStr(parameters["password"]),
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
         ),
         db_case_config=IVFSQ8Config(
             nlist=parameters["nlist"],
@@ -155,7 +171,8 @@ def MilvusDISKANN(**parameters: Unpack[MilvusDISKANNTypedDict]):
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
             user=parameters["user_name"],
-            password=SecretStr(parameters["password"]),
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
         ),
         db_case_config=DISKANNConfig(
             search_list=parameters["search_list"],
@@ -183,7 +200,8 @@ def MilvusGPUIVFFlat(**parameters: Unpack[MilvusGPUIVFTypedDict]):
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
             user=parameters["user_name"],
-            password=SecretStr(parameters["password"]),
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
         ),
         db_case_config=GPUIVFFlatConfig(
             nlist=parameters["nlist"],
@@ -217,7 +235,8 @@ def MilvusGPUBruteForce(**parameters: Unpack[MilvusGPUBruteForceTypedDict]):
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
             user=parameters["user_name"],
-            password=SecretStr(parameters["password"]),
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
         ),
         db_case_config=GPUBruteForceConfig(
             metric_type=parameters["metric_type"],
@@ -248,7 +267,8 @@ def MilvusGPUIVFPQ(**parameters: Unpack[MilvusGPUIVFPQTypedDict]):
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
             user=parameters["user_name"],
-            password=SecretStr(parameters["password"]),
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
         ),
         db_case_config=GPUIVFPQConfig(
             nlist=parameters["nlist"],
@@ -287,7 +307,8 @@ def MilvusGPUCAGRA(**parameters: Unpack[MilvusGPUCAGRATypedDict]):
             db_label=parameters["db_label"],
             uri=SecretStr(parameters["uri"]),
             user=parameters["user_name"],
-            password=SecretStr(parameters["password"]),
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            num_shards=int(parameters["num_shards"]),
         ),
         db_case_config=GPUCAGRAConfig(
             intermediate_graph_degree=parameters["intermediate_graph_degree"],

vectordb_bench/backend/clients/milvus/config.py

@@ -7,12 +7,14 @@ class MilvusConfig(DBConfig):
     uri: SecretStr = "http://localhost:19530"
     user: str | None = None
     password: SecretStr | None = None
+    num_shards: int = 1
 
     def to_dict(self) -> dict:
         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,
+            "num_shards": self.num_shards,
         }
 
     @validator("*")

vectordb_bench/backend/clients/milvus/milvus.py

@@ -40,7 +40,12 @@ class Milvus(VectorDB):
 
         from pymilvus import connections
 
-        connections.connect(**self.db_config, timeout=30)
+        connections.connect(
+            uri=self.db_config.get("uri"),
+            user=self.db_config.get("user"),
+            password=self.db_config.get("password"),
+            timeout=30,
+        )
         if drop_old and utility.has_collection(self.collection_name):
             log.info(f"{self.name} client drop_old collection: {self.collection_name}")
             utility.drop_collection(self.collection_name)
@@ -59,6 +64,7 @@ class Milvus(VectorDB):
                 name=self.collection_name,
                 schema=CollectionSchema(fields),
                 consistency_level="Session",
+                num_shards=self.db_config.get("num_shards"),
             )
 
             log.info(f"{self.name} create index: index_params: {self.case_config.index_param()}")

vectordb_bench/backend/clients/qdrant_cloud/cli.py

@@ -0,0 +1,43 @@
+from typing import Annotated, Unpack
+
+import click
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from .. import DB
+
+
+class QdrantTypedDict(CommonTypedDict):
+    url: Annotated[
+        str,
+        click.option("--url", type=str, help="URL connection string", required=True),
+    ]
+    api_key: Annotated[
+        str | None,
+        click.option("--api-key", type=str, help="API key for authentication", required=False),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(QdrantTypedDict)
+def QdrantCloud(**parameters: Unpack[QdrantTypedDict]):
+    from .config import QdrantConfig, QdrantIndexConfig
+
+    config_params = {
+        "db_label": parameters["db_label"],
+        "url": SecretStr(parameters["url"]),
+    }
+
+    config_params["api_key"] = SecretStr(parameters["api_key"]) if parameters["api_key"] else None
+
+    run(
+        db=DB.QdrantCloud,
+        db_config=QdrantConfig(**config_params),
+        db_case_config=QdrantIndexConfig(),
+        **parameters,
+    )

vectordb_bench/backend/clients/qdrant_cloud/config.py

@@ -6,14 +6,14 @@ from ..api import DBCaseConfig, DBConfig, MetricType
 # Allowing `api_key` to be left empty, to ensure compatibility with the open-source Qdrant.
 class QdrantConfig(DBConfig):
     url: SecretStr
-    api_key: SecretStr
+    api_key: SecretStr | None = None
 
     def to_dict(self) -> dict:
-        api_key = self.api_key.get_secret_value()
-        if len(api_key) > 0:
+        api_key_value = self.api_key.get_secret_value() if self.api_key else None
+        if api_key_value:
             return {
                 "url": self.url.get_secret_value(),
-                "api_key": self.api_key.get_secret_value(),
+                "api_key": api_key_value,
                 "prefer_grpc": True,
             }
         return {

vectordb_bench/backend/clients/qdrant_local/cli.py

@@ -0,0 +1,60 @@
+from typing import Annotated, Unpack
+
+import click
+from pydantic import SecretStr
+
+from vectordb_bench.backend.clients import DB
+from vectordb_bench.cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+
+DBTYPE = DB.QdrantLocal
+
+
+class QdrantLocalTypedDict(CommonTypedDict):
+    url: Annotated[
+        str,
+        click.option("--url", type=str, help="Qdrant url", required=True),
+    ]
+    on_disk: Annotated[
+        bool,
+        click.option("--on-disk", type=bool, default=False, help="Store the vectors and the HNSW index on disk"),
+    ]
+    m: Annotated[
+        int,
+        click.option("--m", type=int, default=16, help="HNSW index parameter m, set 0 to disable the index"),
+    ]
+    ef_construct: Annotated[
+        int,
+        click.option("--ef-construct", type=int, default=200, help="HNSW index parameter ef_construct"),
+    ]
+    hnsw_ef: Annotated[
+        int,
+        click.option(
+            "--hnsw-ef",
+            type=int,
+            default=0,
+            help="HNSW index parameter hnsw_ef, set 0 to use ef_construct for search",
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(QdrantLocalTypedDict)
+def QdrantLocal(**parameters: Unpack[QdrantLocalTypedDict]):
+    from .config import QdrantLocalConfig, QdrantLocalIndexConfig
+
+    run(
+        db=DBTYPE,
+        db_config=QdrantLocalConfig(url=SecretStr(parameters["url"])),
+        db_case_config=QdrantLocalIndexConfig(
+            on_disk=parameters["on_disk"],
+            m=parameters["m"],
+            ef_construct=parameters["ef_construct"],
+            hnsw_ef=parameters["hnsw_ef"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/qdrant_local/config.py

@@ -0,0 +1,47 @@
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, MetricType
+
+
+class QdrantLocalConfig(DBConfig):
+    url: SecretStr
+
+    def to_dict(self) -> dict:
+        return {
+            "url": self.url.get_secret_value(),
+        }
+
+
+class QdrantLocalIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+    m: int
+    ef_construct: int
+    hnsw_ef: int | None = 0
+    on_disk: bool | None = False
+
+    def parse_metric(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "Euclid"
+
+        if self.metric_type == MetricType.IP:
+            return "Dot"
+
+        return "Cosine"
+
+    def index_param(self) -> dict:
+        return {
+            "distance": self.parse_metric(),
+            "m": self.m,
+            "ef_construct": self.ef_construct,
+            "on_disk": self.on_disk,
+        }
+
+    def search_param(self) -> dict:
+        search_params = {
+            "exact": False,  # Force to use ANNs
+        }
+
+        if self.hnsw_ef != 0:
+            search_params["hnsw_ef"] = self.hnsw_ef
+
+        return search_params