vectordb-bench 0.0.29__py3-none-any.whl → 1.0.0__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,10 +42,12 @@ 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}")
         client = OpenSearch(**self.db_config)
         if drop_old:
             log.info(f"AWS_OpenSearch client drop old index: {self.index_name}")
@@ -43,16 +55,26 @@ class AWSOpenSearch(VectorDB):
             if is_existed:
                 client.indices.delete(index=self.index_name)
             self._create_index(client)
+        else:
+            is_existed = client.indices.exists(index=self.index_name)
+            if not is_existed:
+                self._create_index(client)
+                log.info(f"AWS_OpenSearch client create index: {self.index_name}")
+
+            self._update_ef_search_before_search(client)
+            self._load_graphs_to_memory(client)
 
-    @classmethod
-    def config_cls(cls) -> AWSOpenSearchConfig:
-        return AWSOpenSearchConfig
+    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
+        )
+        log.info(f"Creating index with ef_search: {ef_search_value}")
+        log.info(f"Creating index with number_of_replicas: {self.case_config.number_of_replicas}")
 
-    @classmethod
-    def case_config_cls(cls, index_type: IndexType | None = None) -> AWSOpenSearchIndexConfig:
-        return AWSOpenSearchIndexConfig
+        log.info(f"Creating index with engine: {self.case_config.engine}")
+        log.info(f"Creating index with metric type: {self.case_config.metric_type_name}")
+        log.info(f"All case_config parameters: {self.case_config.__dict__}")
 
-    def _create_index(self, client: OpenSearch):
         cluster_settings_body = {
             "persistent": {
                 "knn.algo_param.index_thread_qty": self.case_config.index_thread_qty,
@@ -64,20 +86,18 @@ class AWSOpenSearch(VectorDB):
             "index": {
                 "knn": True,
                 "number_of_shards": self.case_config.number_of_shards,
-                "number_of_replicas": 0,
+                "number_of_replicas": self.case_config.number_of_replicas,
                 "translog.flush_threshold_size": self.case_config.flush_threshold_size,
-                # Setting trans log threshold to 5GB
-                **(
-                    {"knn.algo_param.ef_search": self.case_config.ef_search}
-                    if self.case_config.engine == AWSOS_Engine.nmslib
-                    else {}
-                ),
+                "knn.advanced.approximate_threshold": "-1",
             },
             "refresh_interval": self.case_config.refresh_interval,
         }
+        settings["index"]["knn.algo_param.ef_search"] = ef_search_value
         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,
@@ -86,6 +106,8 @@ class AWSOpenSearch(VectorDB):
             },
         }
         try:
+            log.info(f"Creating index with settings: {settings}")
+            log.info(f"Creating index with mappings: {mappings}")
             client.indices.create(
                 index=self.index_name,
                 body={"settings": settings, "mappings": mappings},
@@ -107,53 +129,193 @@ 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."""
         assert self.client is not None, "should self.init() first"
 
+        num_clients = self.case_config.number_of_indexing_clients or 1
+        log.info(f"Number of indexing clients from case_config: {num_clients}")
+
+        if num_clients <= 1:
+            log.info("Using single client for data insertion")
+            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, labels_data)
+
+    def _insert_with_single_client(
+        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_embeddings(embeddings, metadata)
+            return self._insert_with_single_client(embeddings, metadata)
+
+    def _insert_with_multiple_clients(
+        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
+
+        embeddings_list = list(embeddings)
+        chunk_size = max(1, len(embeddings_list) // num_clients)
+        chunks = []
+
+        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], labels_data[i:end]))
+
+        clients = []
+        for _ in range(min(num_clients, len(chunks))):
+            client = OpenSearch(**self.db_config)
+            clients.append(client)
+
+        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, chunk_labels_data = chunks[chunk_idx]
+            client = clients[client_idx]
+
+            insert_data = []
+            for i in range(len(chunk_embeddings)):
+                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)
+                log.info(f"Client {client_idx} added {len(resp['items'])} documents")
+                return len(chunk_embeddings), None
+            except Exception as e:
+                log.warning(f"Client {client_idx} failed to insert data: {e!s}")
+                return 0, e
+
+        results = []
+        with ThreadPoolExecutor(max_workers=len(clients)) as executor:
+            futures = []
+
+            for chunk_idx in range(len(chunks)):
+                client_idx = chunk_idx % len(clients)
+                futures.append(executor.submit(insert_chunk, client_idx, chunk_idx))
+
+            for future in concurrent.futures.as_completed(futures):
+                count, error = future.result()
+                results.append((count, error))
+
+        from contextlib import suppress
+
+        for client in clients:
+            with suppress(Exception):
+                client.close()
+
+        total_count = sum(count for count, _ in results)
+        errors = [error for _, error in results if error is not None]
+
+        if errors:
+            log.warning("Some clients failed to insert data, retrying with single client")
+            time.sleep(10)
+            return self._insert_with_single_client(embeddings, metadata)
+
+        resp = self.client.indices.stats(self.index_name)
+        log.info(
+            f"""Total document count in index after parallel insertion:
+                {resp['_all']['primaries']['indexing']['index_total']}""",
+        )
+
+        return (total_count, None)
+
+    def _update_ef_search_before_search(self, client: OpenSearch):
+        ef_search_value = (
+            self.case_config.ef_search if self.case_config.ef_search is not None else self.case_config.efSearch
+        )
+
+        try:
+            index_settings = client.indices.get_settings(index=self.index_name)
+            current_ef_search = (
+                index_settings.get(self.index_name, {})
+                .get("settings", {})
+                .get("index", {})
+                .get("knn.algo_param", {})
+                .get("ef_search")
+            )
+
+            if current_ef_search != str(ef_search_value):
+                log.info(f"Updating ef_search before search from {current_ef_search} to {ef_search_value}")
+                settings_body = {"index": {"knn.algo_param.ef_search": ef_search_value}}
+                client.indices.put_settings(index=self.index_name, body=settings_body)
+                log.info(f"Successfully updated ef_search to {ef_search_value} before search")
+
+            log.info(f"Current engine: {self.case_config.engine}")
+            log.info(f"Current metric_type: {self.case_config.metric_type_name}")
+
+        except Exception as e:
+            log.warning(f"Failed to update ef_search parameter before search: {e}")
 
     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.
+            list[int]: list of k most similar ids to the query embedding.
         """
         assert self.client is not None, "should self.init() first"
 
         body = {
             "size": k,
-            "query": {"knn": {self.vector_col_name: {"vector": query, "k": k}}},
-            **({"filter": {"range": {self.id_col_name: {"gt": filters["id"]}}}} if filters else {}),
+            "query": {
+                "knn": {
+                    self.vector_col_name: {
+                        "vector": query,
+                        "k": k,
+                        "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 {}
+                        ),
+                    }
+                }
+            },
         }
+
         try:
             resp = self.client.search(
                 index=self.index_name,
@@ -162,17 +324,38 @@ class AWSOpenSearch(VectorDB):
                 _source=False,
                 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()
         # Call refresh first to ensure that all segments are created
         self._refresh_index()
         if self.case_config.force_merge_enabled:
@@ -182,7 +365,22 @@ class AWSOpenSearch(VectorDB):
         # Call refresh again to ensure that the index is ready after force merge.
         self._refresh_index()
         # ensure that all graphs are loaded in memory and ready for search
-        self._load_graphs_to_memory()
+        self._load_graphs_to_memory(self.client)
+
+    def _update_ef_search(self):
+        ef_search_value = (
+            self.case_config.ef_search if self.case_config.ef_search is not None else self.case_config.efSearch
+        )
+        log.info(f"Updating ef_search parameter to: {ef_search_value}")
+
+        settings_body = {"index": {"knn.algo_param.ef_search": ef_search_value}}
+        try:
+            self.client.indices.put_settings(index=self.index_name, body=settings_body)
+            log.info(f"Successfully updated ef_search to {ef_search_value}")
+            log.info(f"Current engine: {self.case_config.engine}")
+            log.info(f"Current metric_type: {self.case_config.metric_type}")
+        except Exception as e:
+            log.warning(f"Failed to update ef_search parameter: {e}")
 
     def _update_replicas(self):
         index_settings = self.client.indices.get_settings(index=self.index_name)
@@ -200,7 +398,7 @@ class AWSOpenSearch(VectorDB):
         while True:
             res = self.client.cat.indices(index=self.index_name, h="health", format="json")
             health = res[0]["health"]
-            if health != "green":
+            if health == "green":
                 break
             log.info(f"The index {self.index_name} has health : {health} and is not green. Retrying")
             time.sleep(SECONDS_WAITING_FOR_REPLICAS_TO_BE_ENABLED_SEC)
@@ -228,18 +426,26 @@ class AWSOpenSearch(VectorDB):
             "persistent": {"knn.algo_param.index_thread_qty": self.case_config.index_thread_qty_during_force_merge}
         }
         self.client.cluster.put_settings(cluster_settings_body)
-        log.debug(f"Starting force merge for index {self.index_name}")
-        force_merge_endpoint = f"/{self.index_name}/_forcemerge?max_num_segments=1&wait_for_completion=false"
+
+        log.info("Updating the graph threshold to ensure that during merge we can do graph creation.")
+        output = self.client.indices.put_settings(
+            index=self.index_name, body={"index.knn.advanced.approximate_threshold": "0"}
+        )
+        log.info(f"response of updating setting is: {output}")
+
+        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"]
         while True:
             time.sleep(WAITING_FOR_FORCE_MERGE_SEC)
             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):
+    def _load_graphs_to_memory(self, client: OpenSearch):
         if self.case_config.engine != AWSOS_Engine.lucene:
             log.info("Calling warmup API to load graphs into memory")
             warmup_endpoint = f"/_plugins/_knn/warmup/{self.index_name}"
-            self.client.transport.perform_request("GET", warmup_endpoint)
+            client.transport.perform_request("GET", warmup_endpoint)

vectordb_bench/backend/clients/aws_opensearch/cli.py

@@ -1,3 +1,4 @@
+import logging
 from typing import Annotated, TypedDict, Unpack
 
 import click
@@ -5,18 +6,21 @@ from pydantic import SecretStr
 
 from ....cli.cli import (
     CommonTypedDict,
-    HNSWFlavor2,
+    HNSWFlavor1,
     cli,
     click_parameter_decorators_from_typed_dict,
     run,
 )
 from .. import DB
+from .config import AWSOS_Engine, AWSOSQuantization
+
+log = logging.getLogger(__name__)
 
 
 class AWSOpenSearchTypedDict(TypedDict):
     host: Annotated[str, click.option("--host", type=str, help="Db host", required=True)]
-    port: Annotated[int, click.option("--port", type=int, default=443, help="Db Port")]
-    user: Annotated[str, click.option("--user", type=str, default="admin", help="Db User")]
+    port: Annotated[int, click.option("--port", type=int, default=80, help="Db Port")]
+    user: Annotated[str, click.option("--user", type=str, help="Db User")]
     password: Annotated[str, click.option("--password", type=str, help="Db password")]
     number_of_shards: Annotated[
         int,
@@ -38,23 +42,23 @@ class AWSOpenSearchTypedDict(TypedDict):
         ),
     ]
 
-    index_thread_qty_during_force_merge: Annotated[
-        int,
+    engine: Annotated[
+        str,
         click.option(
-            "--index-thread-qty-during-force-merge",
-            type=int,
-            help="Thread count during force merge operations",
-            default=4,
+            "--engine",
+            type=click.Choice(["nmslib", "faiss", "lucene"], case_sensitive=False),
+            help="HNSW algorithm implementation to use",
+            default="faiss",
         ),
     ]
 
-    number_of_indexing_clients: Annotated[
-        int,
+    metric_type: Annotated[
+        str,
         click.option(
-            "--number-of-indexing-clients",
-            type=int,
-            help="Number of concurrent indexing clients",
-            default=1,
+            "--metric-type",
+            type=click.Choice(["l2", "cosine", "ip"], case_sensitive=False),
+            help="Distance metric type for vector similarity",
+            default="l2",
         ),
     ]
 
@@ -64,26 +68,26 @@ class AWSOpenSearchTypedDict(TypedDict):
     ]
 
     refresh_interval: Annotated[
-        int,
+        str,
         click.option(
             "--refresh-interval", type=str, help="How often to make new data available for search", default="60s"
         ),
     ]
 
     force_merge_enabled: Annotated[
-        int,
+        bool,
         click.option("--force-merge-enabled", type=bool, help="Whether to perform force merge operation", default=True),
     ]
 
     flush_threshold_size: Annotated[
-        int,
+        str,
         click.option(
             "--flush-threshold-size", type=str, help="Size threshold for flushing the transaction log", default="5120mb"
         ),
     ]
 
     cb_threshold: Annotated[
-        int,
+        str,
         click.option(
             "--cb-threshold",
             type=str,
@@ -92,8 +96,30 @@ class AWSOpenSearchTypedDict(TypedDict):
         ),
     ]
 
+    quantization_type: Annotated[
+        str | None,
+        click.option(
+            "--quantization-type",
+            type=click.Choice(["fp32", "fp16"]),
+            help="quantization type for vectors (in index)",
+            default="fp32",
+            required=False,
+        ),
+    ]
+
+    engine: Annotated[
+        str | None,
+        click.option(
+            "--engine",
+            type=click.Choice(["faiss", "lucene"]),
+            help="quantization type for vectors (in index)",
+            default="faiss",
+            required=False,
+        ),
+    ]
+
 
-class AWSOpenSearchHNSWTypedDict(CommonTypedDict, AWSOpenSearchTypedDict, HNSWFlavor2): ...
+class AWSOpenSearchHNSWTypedDict(CommonTypedDict, AWSOpenSearchTypedDict, HNSWFlavor1): ...
 
 
 @cli.command()
@@ -117,9 +143,13 @@ def AWSOpenSearch(**parameters: Unpack[AWSOpenSearchHNSWTypedDict]):
             refresh_interval=parameters["refresh_interval"],
             force_merge_enabled=parameters["force_merge_enabled"],
             flush_threshold_size=parameters["flush_threshold_size"],
-            number_of_indexing_clients=parameters["number_of_indexing_clients"],
             index_thread_qty_during_force_merge=parameters["index_thread_qty_during_force_merge"],
             cb_threshold=parameters["cb_threshold"],
+            efConstruction=parameters["ef_construction"],
+            efSearch=parameters["ef_runtime"],
+            M=parameters["m"],
+            engine=AWSOS_Engine(parameters["engine"]),
+            quantization_type=AWSOSQuantization(parameters["quantization_type"]),
         ),
         **parameters,
     )