vectordb-bench 0.0.29__py3-none-any.whl → 1.0.0__py3-none-any.whl

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("*")
@@ -33,6 +35,7 @@ class MilvusIndexConfig(BaseModel):
 
     index: IndexType
     metric_type: MetricType | None = None
+    use_partition_key: bool = True  # for label-filter
 
     @property
     def is_gpu_index(self) -> bool:

vectordb_bench/backend/clients/milvus/milvus.py

@@ -7,6 +7,8 @@ from contextlib import contextmanager
 
 from pymilvus import Collection, CollectionSchema, DataType, FieldSchema, MilvusException, utility
 
+from vectordb_bench.backend.filter import Filter, FilterOp
+
 from ..api import VectorDB
 from .config import MilvusIndexConfig
 
@@ -16,6 +18,12 @@ MILVUS_LOAD_REQS_SIZE = 1.5 * 1024 * 1024
 
 
 class Milvus(VectorDB):
+    supported_filter_types: list[FilterOp] = [
+        FilterOp.NonFilter,
+        FilterOp.NumGE,
+        FilterOp.StrEqual,
+    ]
+
     def __init__(
         self,
         dim: int,
@@ -24,6 +32,7 @@ class Milvus(VectorDB):
         collection_name: str = "VectorDBBenchCollection",
         drop_old: bool = False,
         name: str = "Milvus",
+        with_scalar_labels: bool = False,
         **kwargs,
     ):
         """Initialize wrapper around the milvus vector database."""
@@ -32,15 +41,24 @@ class Milvus(VectorDB):
         self.case_config = db_case_config
         self.collection_name = collection_name
         self.batch_size = int(MILVUS_LOAD_REQS_SIZE / (dim * 4))
+        self.with_scalar_labels = with_scalar_labels
 
         self._primary_field = "pk"
-        self._scalar_field = "id"
+        self._scalar_id_field = "id"
+        self._scalar_label_field = "label"
         self._vector_field = "vector"
-        self._index_name = "vector_idx"
+        self._vector_index_name = "vector_idx"
+        self._scalar_id_index_name = "id_sort_idx"
+        self._scalar_labels_index_name = "labels_idx"
 
         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)
@@ -48,9 +66,20 @@ class Milvus(VectorDB):
         if not utility.has_collection(self.collection_name):
             fields = [
                 FieldSchema(self._primary_field, DataType.INT64, is_primary=True),
-                FieldSchema(self._scalar_field, DataType.INT64),
+                FieldSchema(self._scalar_id_field, DataType.INT64),
                 FieldSchema(self._vector_field, DataType.FLOAT_VECTOR, dim=dim),
             ]
+            if self.with_scalar_labels:
+                is_partition_key = db_case_config.use_partition_key
+                log.info(f"with_scalar_labels, add a new varchar field, as partition_key: {is_partition_key}")
+                fields.append(
+                    FieldSchema(
+                        self._scalar_label_field,
+                        DataType.VARCHAR,
+                        max_length=256,
+                        is_partition_key=is_partition_key,
+                    )
+                )
 
             log.info(f"{self.name} create collection: {self.collection_name}")
 
@@ -59,18 +88,40 @@ 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()}")
-            col.create_index(
-                self._vector_field,
-                self.case_config.index_param(),
-                index_name=self._index_name,
-            )
+            self.create_index()
             col.load()
 
         connections.disconnect("default")
 
+    def create_index(self):
+        col = Collection(self.collection_name)
+        # vector index
+        col.create_index(
+            self._vector_field,
+            self.case_config.index_param(),
+            index_name=self._vector_index_name,
+        )
+        # scalar index for range-expr (int-filter)
+        col.create_index(
+            self._scalar_id_field,
+            index_params={
+                "index_type": "STL_SORT",
+            },
+            index_name=self._scalar_id_index_name,
+        )
+        # scalar index for varchar (label-filter)
+        if self.with_scalar_labels:
+            col.create_index(
+                self._scalar_label_field,
+                index_params={
+                    "index_type": "BITMAP",
+                },
+                index_name=self._scalar_labels_index_name,
+            )
+
     @contextmanager
     def init(self):
         """
@@ -103,17 +154,13 @@ class Milvus(VectorDB):
         try:
             self.col.flush()
             # wait for index done and load refresh
-            self.col.create_index(
-                self._vector_field,
-                self.case_config.index_param(),
-                index_name=self._index_name,
-            )
+            self.create_index()
 
-            utility.wait_for_index_building_complete(self.collection_name)
+            utility.wait_for_index_building_complete(self.collection_name, index_name=self._vector_index_name)
 
             def wait_index():
                 while True:
-                    progress = utility.index_building_progress(self.collection_name)
+                    progress = utility.index_building_progress(self.collection_name, index_name=self._vector_index_name)
                     if progress.get("pending_index_rows", -1) == 0:
                         break
                     time.sleep(5)
@@ -156,6 +203,7 @@ class Milvus(VectorDB):
         self,
         embeddings: Iterable[list[float]],
         metadata: list[int],
+        labels_data: list[str] | None = None,
         **kwargs,
     ) -> tuple[int, Exception]:
         """Insert embeddings into Milvus. should call self.init() first"""
@@ -171,32 +219,42 @@ class Milvus(VectorDB):
                     metadata[batch_start_offset:batch_end_offset],
                     embeddings[batch_start_offset:batch_end_offset],
                 ]
+                if self.with_scalar_labels:
+                    insert_data.append(labels_data[batch_start_offset:batch_end_offset])
                 res = self.col.insert(insert_data)
                 insert_count += len(res.primary_keys)
         except MilvusException as e:
             log.info(f"Failed to insert data: {e}")
-            return (insert_count, e)
-        return (insert_count, None)
+            return insert_count, e
+        return insert_count, None
+
+    def prepare_filter(self, filters: Filter):
+        if filters.type == FilterOp.NonFilter:
+            self.expr = ""
+        elif filters.type == FilterOp.NumGE:
+            self.expr = f"{self._scalar_id_field} >= {filters.int_value}"
+        elif filters.type == FilterOp.StrEqual:
+            self.expr = f"{self._scalar_label_field} == '{filters.label_value}'"
+        else:
+            msg = f"Not support Filter for Milvus - {filters}"
+            raise ValueError(msg)
 
     def search_embedding(
         self,
         query: list[float],
         k: int = 100,
-        filters: dict | None = None,
         timeout: int | None = None,
     ) -> list[int]:
         """Perform a search on a query embedding and return results."""
         assert self.col is not None
 
-        expr = f"{self._scalar_field} {filters.get('metadata')}" if filters else ""
-
         # Perform the search.
         res = self.col.search(
             data=[query],
             anns_field=self._vector_field,
             param=self.case_config.search_param(),
             limit=k,
-            expr=expr,
+            expr=self.expr,
         )
 
         # Organize results.

vectordb_bench/backend/clients/oceanbase/cli.py

@@ -0,0 +1,100 @@
+import os
+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,
+    HNSWFlavor4,
+    OceanBaseIVFTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+
+from ..api import IndexType
+
+
+class OceanBaseTypedDict(CommonTypedDict):
+    host: Annotated[str, click.option("--host", type=str, help="OceanBase host", default="")]
+    user: Annotated[str, click.option("--user", type=str, help="OceanBase username", required=True)]
+    password: Annotated[
+        str,
+        click.option(
+            "--password",
+            type=str,
+            help="OceanBase database password",
+            default=lambda: os.environ.get("OB_PASSWORD", ""),
+        ),
+    ]
+    database: Annotated[str, click.option("--database", type=str, help="DataBase name", required=True)]
+    port: Annotated[int, click.option("--port", type=int, help="OceanBase port", required=True)]
+
+
+class OceanBaseHNSWTypedDict(CommonTypedDict, OceanBaseTypedDict, HNSWFlavor4): ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(OceanBaseHNSWTypedDict)
+def OceanBaseHNSW(**parameters: Unpack[OceanBaseHNSWTypedDict]):
+    from .config import OceanBaseConfig, OceanBaseHNSWConfig
+
+    run(
+        db=DB.OceanBase,
+        db_config=OceanBaseConfig(
+            db_label=parameters["db_label"],
+            user=SecretStr(parameters["user"]),
+            password=SecretStr(parameters["password"]),
+            host=parameters["host"],
+            port=parameters["port"],
+            database=parameters["database"],
+        ),
+        db_case_config=OceanBaseHNSWConfig(
+            m=parameters["m"],
+            efConstruction=parameters["ef_construction"],
+            ef_search=parameters["ef_search"],
+            index=parameters["index_type"],
+        ),
+        **parameters,
+    )
+
+
+class OceanBaseIVFTypedDict(CommonTypedDict, OceanBaseTypedDict, OceanBaseIVFTypedDict): ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(OceanBaseIVFTypedDict)
+def OceanBaseIVF(**parameters: Unpack[OceanBaseIVFTypedDict]):
+    from .config import OceanBaseConfig, OceanBaseIVFConfig
+
+    type_str = parameters["index_type"]
+    if type_str == "IVF_FLAT":
+        input_index_type = IndexType.IVFFlat
+    elif type_str == "IVF_PQ":
+        input_index_type = IndexType.IVFPQ
+    elif type_str == "IVF_SQ8":
+        input_index_type = IndexType.IVFSQ8
+
+    input_m = 0 if parameters["m"] is None else parameters["m"]
+
+    run(
+        db=DB.OceanBase,
+        db_config=OceanBaseConfig(
+            db_label=parameters["db_label"],
+            user=SecretStr(parameters["user"]),
+            password=SecretStr(parameters["password"]),
+            host=parameters["host"],
+            port=parameters["port"],
+            database=parameters["database"],
+        ),
+        db_case_config=OceanBaseIVFConfig(
+            m=input_m,
+            nlist=parameters["nlist"],
+            sample_per_nlist=parameters["sample_per_nlist"],
+            index=input_index_type,
+            ivf_nprobes=parameters["ivf_nprobes"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/oceanbase/config.py

@@ -0,0 +1,125 @@
+from typing import TypedDict
+
+from pydantic import BaseModel, SecretStr, validator
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+
+
+class OceanBaseConfigDict(TypedDict):
+    user: str
+    host: str
+    port: str
+    password: str
+    database: str
+
+
+class OceanBaseConfig(DBConfig):
+    user: SecretStr = SecretStr("root@perf")
+    password: SecretStr
+    host: str
+    port: int
+    database: str
+
+    def to_dict(self) -> OceanBaseConfigDict:
+        user_str = self.user.get_secret_value()
+        pwd_str = self.password.get_secret_value()
+        return {
+            "user": user_str,
+            "host": self.host,
+            "port": self.port,
+            "password": pwd_str,
+            "database": self.database,
+        }
+
+    @validator("*")
+    def not_empty_field(cls, v: any, field: any):
+        if field.name in ["password", "host", "db_label"]:
+            return v
+        if isinstance(v, str | SecretStr) and len(v) == 0:
+            raise ValueError("Empty string!")
+        return v
+
+
+class OceanBaseIndexConfig(BaseModel):
+    index: IndexType
+    metric_type: MetricType | None = None
+    lib: str = "vsag"
+
+    def parse_metric(self) -> str:
+        if self.metric_type == MetricType.L2 or (
+            self.index == IndexType.HNSW_BQ and self.metric_type == MetricType.COSINE
+        ):
+            return "l2"
+        if self.metric_type == MetricType.IP:
+            return "inner_product"
+        return "cosine"
+
+    def parse_metric_func_str(self) -> str:
+        if self.metric_type == MetricType.L2 or (
+            self.index == IndexType.HNSW_BQ and self.metric_type == MetricType.COSINE
+        ):
+            return "l2_distance"
+        if self.metric_type == MetricType.IP:
+            return "negative_inner_product"
+        return "cosine_distance"
+
+
+class OceanBaseHNSWConfig(OceanBaseIndexConfig, DBCaseConfig):
+    m: int
+    efConstruction: int
+    ef_search: int | None = None
+    index: IndexType
+
+    def index_param(self) -> dict:
+        return {
+            "lib": self.lib,
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {"m": self.m, "ef_construction": self.efConstruction},
+        }
+
+    def search_param(self) -> dict:
+        return {"metric_type": self.parse_metric_func_str(), "params": {"ef_search": self.ef_search}}
+
+
+class OceanBaseIVFConfig(OceanBaseIndexConfig, DBCaseConfig):
+    m: int
+    sample_per_nlist: int
+    nlist: int
+    index: IndexType
+    ivf_nprobes: int | None = None
+
+    def index_param(self) -> dict:
+        if self.index == IndexType.IVFPQ:
+            return {
+                "lib": "OB",
+                "metric_type": self.parse_metric(),
+                "index_type": self.index.value,
+                "params": {
+                    "m": self.M,
+                    "sample_per_nlist": self.sample_per_nlist,
+                    "nlist": self.nlist,
+                },
+            }
+        return {
+            "lib": "OB",
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {
+                "sample_per_nlist": self.sample_per_nlist,
+                "nlist": self.nlist,
+            },
+        }
+
+    def search_param(self) -> dict:
+        return {"metric_type": self.metric_type, "params": {"ivf_nprobes": self.ivf_nprobes}}
+
+
+_oceanbase_case_config = {
+    IndexType.HNSW_SQ: OceanBaseHNSWConfig,
+    IndexType.HNSW: OceanBaseHNSWConfig,
+    IndexType.HNSW_BQ: OceanBaseHNSWConfig,
+    IndexType.IVFFlat: OceanBaseIVFConfig,
+    IndexType.IVFPQ: OceanBaseIVFConfig,
+    IndexType.IVFSQ8: OceanBaseIVFConfig,
+}