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

vectordb_bench/backend/clients/aws_opensearch/config.py

@@ -10,17 +10,21 @@ log = logging.getLogger(__name__)
 
 class AWSOpenSearchConfig(DBConfig, BaseModel):
     host: str = ""
-    port: int = 443
+    port: int = 80
     user: str = ""
     password: SecretStr = ""
 
     def to_dict(self) -> dict:
+        use_ssl = self.port == 443
+        http_auth = (
+            (self.user, self.password.get_secret_value()) if len(self.user) != 0 and len(self.password) != 0 else ()
+        )
         return {
             "hosts": [{"host": self.host, "port": self.port}],
-            "http_auth": (self.user, self.password.get_secret_value()),
-            "use_ssl": True,
+            "http_auth": http_auth,
+            "use_ssl": use_ssl,
             "http_compress": True,
-            "verify_certs": True,
+            "verify_certs": use_ssl,
             "ssl_assert_hostname": False,
             "ssl_show_warn": False,
             "timeout": 600,
@@ -28,16 +32,22 @@ class AWSOpenSearchConfig(DBConfig, BaseModel):
 
 
 class AWSOS_Engine(Enum):
-    nmslib = "nmslib"
     faiss = "faiss"
-    lucene = "Lucene"
+    lucene = "lucene"
+
+
+class AWSOSQuantization(Enum):
+    fp32 = "fp32"
+    fp16 = "fp16"
 
 
 class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType = MetricType.L2
     engine: AWSOS_Engine = AWSOS_Engine.faiss
     efConstruction: int = 256
-    efSearch: int = 256
+    efSearch: int = 100
+    engine_name: str | None = None
+    metric_type_name: str | None = None
     M: int = 16
     index_thread_qty: int | None = 4
     number_of_shards: int | None = 1
@@ -46,33 +56,65 @@ class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
     refresh_interval: str | None = "60s"
     force_merge_enabled: bool | None = True
     flush_threshold_size: str | None = "5120mb"
-    number_of_indexing_clients: int | None = 1
-    index_thread_qty_during_force_merge: int
+    index_thread_qty_during_force_merge: int = 8
     cb_threshold: str | None = "50%"
+    number_of_indexing_clients: int | None = 1
+    use_routing: bool = False  # for label-filter cases
+    oversample_factor: float = 1.0
+    quantization_type: AWSOSQuantization = AWSOSQuantization.fp32
+
+    def __eq__(self, obj: any):
+        return (
+            self.engine == obj.engine
+            and self.M == obj.M
+            and self.efConstruction == obj.efConstruction
+            and self.number_of_shards == obj.number_of_shards
+            and self.number_of_replicas == obj.number_of_replicas
+            and self.number_of_segments == obj.number_of_segments
+            and self.use_routing == obj.use_routing
+            and self.quantization_type == obj.quantization_type
+        )
 
     def parse_metric(self) -> str:
+        log.info(f"User specified metric_type: {self.metric_type_name}")
+        self.metric_type = MetricType[self.metric_type_name.upper()]
         if self.metric_type == MetricType.IP:
             return "innerproduct"
         if self.metric_type == MetricType.COSINE:
-            if self.engine == AWSOS_Engine.faiss:
-                log.info(
-                    "Using innerproduct because faiss doesn't support cosine as metric type for Opensearch",
-                )
-                return "innerproduct"
             return "cosinesimil"
+        if self.metric_type == MetricType.L2:
+            log.info("Using l2 as specified by user")
+            return "l2"
         return "l2"
 
+    @property
+    def use_quant(self) -> bool:
+        return self.quantization_type is not AWSOSQuantization.fp32
+
     def index_param(self) -> dict:
+        log.info(f"Using engine: {self.engine} for index creation")
+        log.info(f"Using metric_type: {self.metric_type_name} for index creation")
+        log.info(f"Resulting space_type: {self.parse_metric()} for index creation")
+
+        parameters = {"ef_construction": self.efConstruction, "m": self.M}
+
+        if self.engine == AWSOS_Engine.faiss and self.faiss_use_fp16:
+            parameters["encoder"] = {"name": "sq", "parameters": {"type": "fp16"}}
+
         return {
             "name": "hnsw",
-            "space_type": self.parse_metric(),
             "engine": self.engine.value,
             "parameters": {
                 "ef_construction": self.efConstruction,
                 "m": self.M,
                 "ef_search": self.efSearch,
+                **(
+                    {"encoder": {"name": "sq", "parameters": {"type": self.quantization_type.fp16.value}}}
+                    if self.use_quant
+                    else {}
+                ),
             },
         }
 
     def search_param(self) -> dict:
-        return {}
+        return {"ef_search": self.efSearch}

vectordb_bench/backend/clients/chroma/chroma.py

@@ -78,8 +78,12 @@ class ChromaClient(VectorDB):
         """
         ids = [str(i) for i in metadata]
         metadata = [{"id": int(i)} for i in metadata]
-        if len(embeddings) > 0:
-            self.collection.add(embeddings=embeddings, ids=ids, metadatas=metadata)
+        try:
+            if len(embeddings) > 0:
+                self.collection.add(embeddings=embeddings, ids=ids, metadatas=metadata)
+        except Exception as e:
+            log.warning(f"Failed to insert data: error: {e!s}")
+            return 0, e
         return len(embeddings), None
 
     def search_embedding(

vectordb_bench/backend/clients/elastic_cloud/config.py

@@ -23,13 +23,31 @@ class ESElementType(str, Enum):
 
 class ElasticCloudIndexConfig(BaseModel, DBCaseConfig):
     element_type: ESElementType = ESElementType.float
-    index: IndexType = IndexType.ES_HNSW  # ES only support 'hnsw'
+    index: IndexType = IndexType.ES_HNSW
+    number_of_shards: int = 1
+    number_of_replicas: int = 0
+    refresh_interval: str = "30s"
+    merge_max_thread_count: int = 8
+    use_rescore: bool = False
+    oversample_ratio: float = 2.0
+    use_routing: bool = False
+    use_force_merge: bool = True
 
     metric_type: MetricType | None = None
     efConstruction: int | None = None
     M: int | None = None
     num_candidates: int | None = None
 
+    def __eq__(self, obj: any):
+        return (
+            self.index == obj.index
+            and self.number_of_shards == obj.number_of_shards
+            and self.number_of_replicas == obj.number_of_replicas
+            and self.use_routing == obj.use_routing
+            and self.efConstruction == obj.efConstruction
+            and self.M == obj.M
+        )
+
     def parse_metric(self) -> str:
         if self.metric_type == MetricType.L2:
             return "l2_norm"

vectordb_bench/backend/clients/elastic_cloud/elastic_cloud.py

@@ -5,6 +5,8 @@ from contextlib import contextmanager
 
 from elasticsearch.helpers import bulk
 
+from vectordb_bench.backend.filter import Filter, FilterOp
+
 from ..api import VectorDB
 from .config import ElasticCloudIndexConfig
 
@@ -18,6 +20,12 @@ SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC = 30
 
 
 class ElasticCloud(VectorDB):
+    supported_filter_types: list[FilterOp] = [
+        FilterOp.NonFilter,
+        FilterOp.NumGE,
+        FilterOp.StrEqual,
+    ]
+
     def __init__(
         self,
         dim: int,
@@ -25,8 +33,10 @@ class ElasticCloud(VectorDB):
         db_case_config: ElasticCloudIndexConfig,
         indice: str = "vdb_bench_indice",  # must be lowercase
         id_col_name: str = "id",
+        label_col_name: str = "label",
         vector_col_name: str = "vector",
         drop_old: bool = False,
+        with_scalar_labels: bool = False,
         **kwargs,
     ):
         self.dim = dim
@@ -34,7 +44,9 @@ class ElasticCloud(VectorDB):
         self.case_config = db_case_config
         self.indice = indice
         self.id_col_name = id_col_name
+        self.label_col_name = label_col_name
         self.vector_col_name = vector_col_name
+        self.with_scalar_labels = with_scalar_labels
 
         from elasticsearch import Elasticsearch
 
@@ -69,9 +81,17 @@ class ElasticCloud(VectorDB):
                 },
             },
         }
+        settings = {
+            "index": {
+                "number_of_shards": self.case_config.number_of_shards,
+                "number_of_replicas": self.case_config.number_of_replicas,
+                "refresh_interval": self.case_config.refresh_interval,
+                "merge.scheduler.max_thread_count": self.case_config.merge_max_thread_count,
+            }
+        }
 
         try:
-            client.indices.create(index=self.indice, mappings=mappings)
+            client.indices.create(index=self.indice, mappings=mappings, settings=settings)
         except Exception as e:
             log.warning(f"Failed to create indice: {self.indice} error: {e!s}")
             raise e from None
@@ -80,21 +100,48 @@ class ElasticCloud(VectorDB):
         self,
         embeddings: Iterable[list[float]],
         metadata: list[int],
+        labels_data: list[str] | None = None,
         **kwargs,
     ) -> tuple[int, Exception]:
         """Insert the embeddings to the elasticsearch."""
         assert self.client is not None, "should self.init() first"
 
-        insert_data = [
-            {
-                "_index": self.indice,
-                "_source": {
-                    self.id_col_name: metadata[i],
-                    self.vector_col_name: embeddings[i],
-                },
-            }
-            for i in range(len(embeddings))
-        ]
+        insert_data = (
+            [
+                (
+                    {
+                        "_index": self.indice,
+                        "_source": {
+                            self.id_col_name: metadata[i],
+                            self.label_col_name: labels_data[i],
+                            self.vector_col_name: embeddings[i],
+                        },
+                        "_routing": labels_data[i],
+                    }
+                    if self.case_config.use_routing
+                    else {
+                        "_index": self.indice,
+                        "_source": {
+                            self.id_col_name: metadata[i],
+                            self.label_col_name: labels_data[i],
+                            self.vector_col_name: embeddings[i],
+                        },
+                    }
+                )
+                for i in range(len(embeddings))
+            ]
+            if self.with_scalar_labels
+            else [
+                {
+                    "_index": self.indice,
+                    "_source": {
+                        self.id_col_name: metadata[i],
+                        self.vector_col_name: embeddings[i],
+                    },
+                }
+                for i in range(len(embeddings))
+            ]
+        )
         try:
             bulk_insert_res = bulk(self.client, insert_data)
             return (bulk_insert_res[0], None)
@@ -102,59 +149,100 @@ class ElasticCloud(VectorDB):
             log.warning(f"Failed to insert data: {self.indice} error: {e!s}")
             return (0, e)
 
+    def prepare_filter(self, filters: Filter):
+        self.routing_key = None
+        if filters.type == FilterOp.NonFilter:
+            self.filter = []
+        elif filters.type == FilterOp.NumGE:
+            self.filter = {"range": {self.id_col_name: {"gt": filters.int_value}}}
+        elif filters.type == FilterOp.StrEqual:
+            self.filter = {"term": {self.label_col_name: filters.label_value}}
+            if self.case_config.use_routing:
+                self.routing_key = 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,
+        **kwargs,
     ) -> list[int]:
         """Get k most similar embeddings to query vector.
 
         Args:
             query(list[float]): query embedding to look up documents similar to.
             k(int): Number of most similar embeddings to return. Defaults to 100.
-            filters(dict, optional): filtering expression to filter the data while searching.
 
         Returns:
             list[tuple[int, float]]: list of k most similar embeddings in (id, score) tuple to the query embedding.
         """
         assert self.client is not None, "should self.init() first"
 
-        knn = {
-            "field": self.vector_col_name,
-            "k": k,
-            "num_candidates": self.case_config.num_candidates,
-            "filter": [{"range": {self.id_col_name: {"gt": filters["id"]}}}] if filters else [],
-            "query_vector": query,
-        }
+        if self.case_config.use_rescore:
+            oversample_k = int(k * self.case_config.oversample_ratio)
+            oversample_num_candidates = int(self.case_config.num_candidates * self.case_config.oversample_ratio)
+            knn = {
+                "field": self.vector_col_name,
+                "k": oversample_k,
+                "num_candidates": oversample_num_candidates,
+                "filter": self.filter,
+                "query_vector": query,
+            }
+            rescore = {
+                "window_size": oversample_k,
+                "query": {
+                    "rescore_query": {
+                        "script_score": {
+                            "query": {"match_all": {}},
+                            "script": {
+                                "source": f"cosineSimilarity(params.queryVector, '{self.vector_col_name}')",
+                                "params": {"queryVector": query},
+                            },
+                        }
+                    },
+                    "query_weight": 0,
+                    "rescore_query_weight": 1,
+                },
+            }
+        else:
+            knn = {
+                "field": self.vector_col_name,
+                "k": k,
+                "num_candidates": self.case_config.num_candidates,
+                "filter": self.filter,
+                "query_vector": query,
+            }
+            rescore = None
         size = k
-        try:
-            res = self.client.search(
-                index=self.indice,
-                knn=knn,
-                size=size,
-                _source=False,
-                docvalue_fields=[self.id_col_name],
-                stored_fields="_none_",
-                filter_path=[f"hits.hits.fields.{self.id_col_name}"],
-            )
-            return [h["fields"][self.id_col_name][0] for h in res["hits"]["hits"]]
-        except Exception as e:
-            log.warning(f"Failed to search: {self.indice} error: {e!s}")
-            raise e from None
+
+        res = self.client.search(
+            index=self.indice,
+            knn=knn,
+            routing=self.routing_key,
+            rescore=rescore,
+            size=size,
+            _source=False,
+            docvalue_fields=[self.id_col_name],
+            stored_fields="_none_",
+            filter_path=[f"hits.hits.fields.{self.id_col_name}"],
+        )
+        return [h["fields"][self.id_col_name][0] for h in res["hits"]["hits"]]
 
     def optimize(self, data_size: int | None = None):
         """optimize will be called between insertion and search in performance cases."""
         assert self.client is not None, "should self.init() first"
         self.client.indices.refresh(index=self.indice)
-        force_merge_task_id = self.client.indices.forcemerge(
-            index=self.indice,
-            max_num_segments=1,
-            wait_for_completion=False,
-        )["task"]
-        log.info(f"Elasticsearch force merge task id: {force_merge_task_id}")
-        while True:
-            time.sleep(SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC)
-            task_status = self.client.tasks.get(task_id=force_merge_task_id)
-            if task_status["completed"]:
-                return
+        if self.case_config.use_force_merge:
+            force_merge_task_id = self.client.indices.forcemerge(
+                index=self.indice,
+                max_num_segments=1,
+                wait_for_completion=False,
+            )["task"]
+            log.info(f"Elasticsearch force merge task id: {force_merge_task_id}")
+            while True:
+                time.sleep(SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC)
+                task_status = self.client.tasks.get(task_id=force_merge_task_id)
+                if task_status["completed"]:
+                    return

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,20 +81,28 @@ class LanceDB(VectorDB):
         filters: dict | None = None,
     ) -> list[int]:
         if filters:
-            results = (
-                self.table.search(query)
-                .select(["id"])
-                .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).select(["id"]).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__)