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

vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py

@@ -5,8 +5,10 @@ from contextlib import contextmanager
 
 from opensearchpy import OpenSearch
 
-from ..api import IndexType, VectorDB
-from .config import AWSOpenSearchConfig, AWSOpenSearchIndexConfig, AWSOS_Engine
+from vectordb_bench.backend.filter import Filter, FilterOp
+
+from ..api import VectorDB
+from .config import AWSOpenSearchIndexConfig, AWSOS_Engine
 
 log = logging.getLogger(__name__)
 
@@ -16,6 +18,12 @@ SECONDS_WAITING_FOR_REPLICAS_TO_BE_ENABLED_SEC = 30
 
 
 class AWSOpenSearch(VectorDB):
+    supported_filter_types: list[FilterOp] = [
+        FilterOp.NonFilter,
+        FilterOp.NumGE,
+        FilterOp.StrEqual,
+    ]
+
     def __init__(
         self,
         dim: int,
@@ -23,8 +31,10 @@ class AWSOpenSearch(VectorDB):
         db_case_config: AWSOpenSearchIndexConfig,
         index_name: str = "vdb_bench_index",  # must be lowercase
         id_col_name: str = "_id",
+        label_col_name: str = "label",
         vector_col_name: str = "embedding",
         drop_old: bool = False,
+        with_scalar_labels: bool = False,
         **kwargs,
     ):
         self.dim = dim
@@ -32,8 +42,9 @@ class AWSOpenSearch(VectorDB):
         self.case_config = db_case_config
         self.index_name = index_name
         self.id_col_name = id_col_name
-        self.category_col_names = [f"scalar-{categoryCount}" for categoryCount in [2, 5, 10, 100, 1000]]
+        self.label_col_name = label_col_name
         self.vector_col_name = vector_col_name
+        self.with_scalar_labels = with_scalar_labels
 
         log.info(f"AWS_OpenSearch client config: {self.db_config}")
         log.info(f"AWS_OpenSearch db case config : {self.case_config}")
@@ -53,14 +64,6 @@ class AWSOpenSearch(VectorDB):
             self._update_ef_search_before_search(client)
             self._load_graphs_to_memory(client)
 
-    @classmethod
-    def config_cls(cls) -> AWSOpenSearchConfig:
-        return AWSOpenSearchConfig
-
-    @classmethod
-    def case_config_cls(cls, index_type: IndexType | None = None) -> AWSOpenSearchIndexConfig:
-        return AWSOpenSearchIndexConfig
-
     def _create_index(self, client: OpenSearch) -> None:
         ef_search_value = (
             self.case_config.ef_search if self.case_config.ef_search is not None else self.case_config.efSearch
@@ -93,7 +96,8 @@ class AWSOpenSearch(VectorDB):
         mappings = {
             "_source": {"excludes": [self.vector_col_name], "recovery_source_excludes": [self.vector_col_name]},
             "properties": {
-                **{categoryCol: {"type": "keyword"} for categoryCol in self.category_col_names},
+                self.id_col_name: {"type": "integer", "store": True},
+                self.label_col_name: {"type": "keyword"},
                 self.vector_col_name: {
                     "type": "knn_vector",
                     "dimension": self.dim,
@@ -125,6 +129,7 @@ class AWSOpenSearch(VectorDB):
         self,
         embeddings: Iterable[list[float]],
         metadata: list[int],
+        labels_data: list[str] | None = None,
         **kwargs,
     ) -> tuple[int, Exception]:
         """Insert the embeddings to the opensearch."""
@@ -135,34 +140,42 @@ class AWSOpenSearch(VectorDB):
 
         if num_clients <= 1:
             log.info("Using single client for data insertion")
-            return self._insert_with_single_client(embeddings, metadata)
+            return self._insert_with_single_client(embeddings, metadata, labels_data)
         log.info(f"Using {num_clients} parallel clients for data insertion")
-        return self._insert_with_multiple_clients(embeddings, metadata, num_clients)
+        return self._insert_with_multiple_clients(embeddings, metadata, num_clients, labels_data)
 
     def _insert_with_single_client(
-        self, embeddings: Iterable[list[float]], metadata: list[int]
+        self,
+        embeddings: Iterable[list[float]],
+        metadata: list[int],
+        labels_data: list[str] | None = None,
     ) -> tuple[int, Exception]:
         insert_data = []
         for i in range(len(embeddings)):
-            insert_data.append(
-                {"index": {"_index": self.index_name, self.id_col_name: metadata[i]}},
-            )
-            insert_data.append({self.vector_col_name: embeddings[i]})
+            index_data = {"index": {"_index": self.index_name, self.id_col_name: metadata[i]}}
+            if self.with_scalar_labels and self.case_config.use_routing:
+                index_data["routing"] = labels_data[i]
+            insert_data.append(index_data)
+
+            other_data = {self.vector_col_name: embeddings[i]}
+            if self.with_scalar_labels:
+                other_data[self.label_col_name] = labels_data[i]
+            insert_data.append(other_data)
+
         try:
-            resp = self.client.bulk(insert_data)
-            log.info(f"AWS_OpenSearch adding documents: {len(resp['items'])}")
-            resp = self.client.indices.stats(self.index_name)
-            log.info(
-                f"Total document count in index: {resp['_all']['primaries']['indexing']['index_total']}",
-            )
-            return (len(embeddings), None)
+            self.client.bulk(insert_data)
+            return len(embeddings), None
         except Exception as e:
             log.warning(f"Failed to insert data: {self.index_name} error: {e!s}")
             time.sleep(10)
             return self._insert_with_single_client(embeddings, metadata)
 
     def _insert_with_multiple_clients(
-        self, embeddings: Iterable[list[float]], metadata: list[int], num_clients: int
+        self,
+        embeddings: Iterable[list[float]],
+        metadata: list[int],
+        num_clients: int,
+        labels_data: list[str] | None = None,
     ) -> tuple[int, Exception]:
         import concurrent.futures
         from concurrent.futures import ThreadPoolExecutor
@@ -173,7 +186,7 @@ class AWSOpenSearch(VectorDB):
 
         for i in range(0, len(embeddings_list), chunk_size):
             end = min(i + chunk_size, len(embeddings_list))
-            chunks.append((embeddings_list[i:end], metadata[i:end]))
+            chunks.append((embeddings_list[i:end], metadata[i:end], labels_data[i:end]))
 
         clients = []
         for _ in range(min(num_clients, len(chunks))):
@@ -183,15 +196,20 @@ class AWSOpenSearch(VectorDB):
         log.info(f"AWS_OpenSearch using {len(clients)} parallel clients for data insertion")
 
         def insert_chunk(client_idx: int, chunk_idx: int):
-            chunk_embeddings, chunk_metadata = chunks[chunk_idx]
+            chunk_embeddings, chunk_metadata, chunk_labels_data = chunks[chunk_idx]
             client = clients[client_idx]
 
             insert_data = []
             for i in range(len(chunk_embeddings)):
-                insert_data.append(
-                    {"index": {"_index": self.index_name, self.id_col_name: chunk_metadata[i]}},
-                )
-                insert_data.append({self.vector_col_name: chunk_embeddings[i]})
+                index_data = {"index": {"_index": self.index_name, self.id_col_name: chunk_metadata[i]}}
+                if self.with_scalar_labels and self.case_config.use_routing:
+                    index_data["routing"] = chunk_labels_data[i]
+                insert_data.append(index_data)
+
+                other_data = {self.vector_col_name: chunk_embeddings[i]}
+                if self.with_scalar_labels:
+                    other_data[self.label_col_name] = chunk_labels_data[i]
+                insert_data.append(other_data)
 
             try:
                 resp = client.bulk(insert_data)
@@ -266,17 +284,16 @@ class AWSOpenSearch(VectorDB):
         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.
+            list[int]: list of k most similar ids to the query embedding.
         """
         assert self.client is not None, "should self.init() first"
 
@@ -287,11 +304,16 @@ class AWSOpenSearch(VectorDB):
                     self.vector_col_name: {
                         "vector": query,
                         "k": k,
-                        "method_parameters": {"ef_search": self.case_config.efSearch},
+                        "method_parameters": self.case_config.search_param(),
+                        **({"filter": self.filter} if self.filter else {}),
+                        **(
+                            {"rescore": {"oversample_factor": self.case_config.oversample_factor}}
+                            if self.case_config.use_quant
+                            else {}
+                        ),
                     }
                 }
             },
-            **({"filter": {"range": {self.id_col_name: {"gt": filters["id"]}}}} if filters else {}),
         }
 
         try:
@@ -303,15 +325,34 @@ class AWSOpenSearch(VectorDB):
                 docvalue_fields=[self.id_col_name],
                 stored_fields="_none_",
                 preference="_only_local" if self.case_config.number_of_shards == 1 else None,
+                routing=self.routing_key,
             )
             log.debug(f"Search took: {resp['took']}")
             log.debug(f"Search shards: {resp['_shards']}")
             log.debug(f"Search hits total: {resp['hits']['total']}")
-            return [int(h["fields"][self.id_col_name][0]) for h in resp["hits"]["hits"]]
+            try:
+                return [int(h["fields"][self.id_col_name][0]) for h in resp["hits"]["hits"]]
+            except Exception:
+                # empty results
+                return []
         except Exception as e:
             log.warning(f"Failed to search: {self.index_name} error: {e!s}")
             raise e from None
 
+    def prepare_filter(self, filters: Filter):
+        self.routing_key = None
+        if filters.type == FilterOp.NonFilter:
+            self.filter = None
+        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 OpenSearch - {filters}"
+            raise ValueError(msg)
+
     def optimize(self, data_size: int | None = None):
         """optimize will be called between insertion and search in performance cases."""
         self._update_ef_search()
@@ -392,7 +433,7 @@ class AWSOpenSearch(VectorDB):
         )
         log.info(f"response of updating setting is: {output}")
 
-        log.debug(f"Starting force merge for index {self.index_name}")
+        log.info(f"Starting force merge for index {self.index_name}")
         segments = self.case_config.number_of_segments
         force_merge_endpoint = f"/{self.index_name}/_forcemerge?max_num_segments={segments}&wait_for_completion=false"
         force_merge_task_id = self.client.transport.perform_request("POST", force_merge_endpoint)["task"]
@@ -401,7 +442,7 @@ class AWSOpenSearch(VectorDB):
             task_status = self.client.tasks.get(task_id=force_merge_task_id)
             if task_status["completed"]:
                 break
-        log.debug(f"Completed force merge for index {self.index_name}")
+        log.info(f"Completed force merge for index {self.index_name}")
 
     def _load_graphs_to_memory(self, client: OpenSearch):
         if self.case_config.engine != AWSOS_Engine.lucene:

vectordb_bench/backend/clients/aws_opensearch/config.py

@@ -45,7 +45,7 @@ class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType = MetricType.L2
     engine: AWSOS_Engine = AWSOS_Engine.faiss
     efConstruction: int = 256
-    ef_search: int = 200
+    efSearch: int = 100
     engine_name: str | None = None
     metric_type_name: str | None = None
     M: int = 16
@@ -56,10 +56,39 @@ class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
     refresh_interval: str | None = "60s"
     force_merge_enabled: bool | None = True
     flush_threshold_size: str | None = "5120mb"
-    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 __hash__(self) -> int:
+        return hash(
+            (
+                self.engine,
+                self.M,
+                self.efConstruction,
+                self.number_of_shards,
+                self.number_of_replicas,
+                self.number_of_segments,
+                self.use_routing,
+                self.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()]
@@ -72,6 +101,10 @@ class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
             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")
@@ -91,11 +124,11 @@ class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
                 "ef_search": self.efSearch,
                 **(
                     {"encoder": {"name": "sq", "parameters": {"type": self.quantization_type.fp16.value}}}
-                    if self.quantization_type is not AWSOSQuantization.fp32
+                    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,43 @@ 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 __hash__(self) -> int:
+        return hash(
+            (
+                self.index,
+                self.number_of_shards,
+                self.number_of_replicas,
+                self.use_routing,
+                self.efConstruction,
+                self.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/milvus/config.py

@@ -35,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: