vectordb-bench 0.0.19__py3-none-any.whl → 0.0.21__py3-none-any.whl

vectordb_bench/backend/clients/chroma/chroma.py

@@ -1,55 +1,55 @@
-import chromadb
-import logging 
+import logging
 from contextlib import contextmanager
 from typing import Any
-from ..api import VectorDB, DBCaseConfig
+
+import chromadb
+
+from ..api import DBCaseConfig, VectorDB
 
 log = logging.getLogger(__name__)
+
+
 class ChromaClient(VectorDB):
-    """Chroma client for VectorDB. 
+    """Chroma client for VectorDB.
     To set up Chroma in docker, see https://docs.trychroma.com/usage-guide
     or the instructions in tests/test_chroma.py
 
     To change to running in process, modify the HttpClient() in __init__() and init().
-    """ 
+    """
 
     def __init__(
-            self,
-            dim: int,
-            db_config: dict,
-            db_case_config: DBCaseConfig,
-            drop_old: bool = False,
-
-            **kwargs
-        ):
-
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: DBCaseConfig,
+        drop_old: bool = False,
+        **kwargs,
+    ):
         self.db_config = db_config
         self.case_config = db_case_config
-        self.collection_name = 'example2'
+        self.collection_name = "example2"
 
-        client = chromadb.HttpClient(host=self.db_config["host"], 
-                                     port=self.db_config["port"])
+        client = chromadb.HttpClient(host=self.db_config["host"], port=self.db_config["port"])
         assert client.heartbeat() is not None
         if drop_old:
             try:
-                client.reset() # Reset the database
-            except:
+                client.reset()  # Reset the database
+            except Exception:
                 drop_old = False
                 log.info(f"Chroma client drop_old collection: {self.collection_name}")
 
     @contextmanager
     def init(self) -> None:
-        """ create and destory connections to database.
+        """create and destory connections to database.
 
         Examples:
             >>> with self.init():
             >>>     self.insert_embeddings()
         """
-        #create connection
-        self.client = chromadb.HttpClient(host=self.db_config["host"], 
-                                          port=self.db_config["port"])
-        
-        self.collection = self.client.get_or_create_collection('example2')
+        # create connection
+        self.client = chromadb.HttpClient(host=self.db_config["host"], port=self.db_config["port"])
+
+        self.collection = self.client.get_or_create_collection("example2")
         yield
         self.client = None
         self.collection = None
@@ -57,10 +57,7 @@ class ChromaClient(VectorDB):
     def ready_to_search(self) -> bool:
         pass
 
-    def ready_to_load(self) -> bool:
-        pass
-
-    def optimize(self) -> None:
+    def optimize(self, data_size: int | None = None):
         pass
 
     def insert_embeddings(
@@ -79,12 +76,12 @@ class ChromaClient(VectorDB):
         Returns:
             (int, Exception): number of embeddings inserted and exception if any
         """
-        ids=[str(i) for i in metadata]
-        metadata = [{"id": int(i)} for i in metadata] 
+        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)
         return len(embeddings), None
-    
+
     def search_embedding(
         self,
         query: list[float],
@@ -100,17 +97,19 @@ class ChromaClient(VectorDB):
             kwargs: other arguments
 
         Returns:
-            Dict {ids: list[list[int]], 
-                    embedding: list[list[float]] 
+            Dict {ids: list[list[int]],
+                    embedding: list[list[float]]
                     distance: list[list[float]]}
         """
         if filters:
             # assumes benchmark test filters of format: {'metadata': '>=10000', 'id': 10000}
             id_value = filters.get("id")
-            results = self.collection.query(query_embeddings=query, n_results=k, 
-                                                where={"id": {"$gt": id_value}})
-            #return list of id's in results
-            return [int(i) for i in results.get('ids')[0]]
+            results = self.collection.query(
+                query_embeddings=query,
+                n_results=k,
+                where={"id": {"$gt": id_value}},
+            )
+            # return list of id's in results
+            return [int(i) for i in results.get("ids")[0]]
         results = self.collection.query(query_embeddings=query, n_results=k)
-        return [int(i) for i in results.get('ids')[0]]
-    
+        return [int(i) for i in results.get("ids")[0]]

vectordb_bench/backend/clients/chroma/config.py

@@ -1,14 +1,16 @@
 from pydantic import SecretStr
+
 from ..api import DBConfig
 
+
 class ChromaConfig(DBConfig):
     password: SecretStr
     host: SecretStr
-    port: int 
+    port: int
 
     def to_dict(self) -> dict:
         return {
             "host": self.host.get_secret_value(),
             "port": self.port,
             "password": self.password.get_secret_value(),
-        }
+        }

vectordb_bench/backend/clients/elastic_cloud/config.py

@@ -1,7 +1,8 @@
 from enum import Enum
-from pydantic import SecretStr, BaseModel
 
-from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
 
 
 class ElasticCloudConfig(DBConfig, BaseModel):
@@ -32,12 +33,12 @@ class ElasticCloudIndexConfig(BaseModel, DBCaseConfig):
     def parse_metric(self) -> str:
         if self.metric_type == MetricType.L2:
             return "l2_norm"
-        elif self.metric_type == MetricType.IP:
+        if self.metric_type == MetricType.IP:
             return "dot_product"
         return "cosine"
 
     def index_param(self) -> dict:
-        params = {
+        return {
             "type": "dense_vector",
             "index": True,
             "element_type": self.element_type.value,
@@ -48,7 +49,6 @@ class ElasticCloudIndexConfig(BaseModel, DBCaseConfig):
                 "ef_construction": self.efConstruction,
             },
         }
-        return params
 
     def search_param(self) -> dict:
         return {

vectordb_bench/backend/clients/elastic_cloud/elastic_cloud.py

@@ -1,17 +1,22 @@
 import logging
 import time
+from collections.abc import Iterable
 from contextlib import contextmanager
-from typing import Iterable
-from ..api import VectorDB
-from .config import ElasticCloudIndexConfig
+
 from elasticsearch.helpers import bulk
 
+from ..api import VectorDB
+from .config import ElasticCloudIndexConfig
 
 for logger in ("elasticsearch", "elastic_transport"):
     logging.getLogger(logger).setLevel(logging.WARNING)
 
 log = logging.getLogger(__name__)
 
+
+SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC = 30
+
+
 class ElasticCloud(VectorDB):
     def __init__(
         self,
@@ -46,14 +51,14 @@ class ElasticCloud(VectorDB):
     def init(self) -> None:
         """connect to elasticsearch"""
         from elasticsearch import Elasticsearch
+
         self.client = Elasticsearch(**self.db_config, request_timeout=180)
 
         yield
-        # self.client.transport.close()
         self.client = None
-        del(self.client)
+        del self.client
 
-    def _create_indice(self, client) -> None:
+    def _create_indice(self, client: any) -> None:
         mappings = {
             "_source": {"excludes": [self.vector_col_name]},
             "properties": {
@@ -62,13 +67,13 @@ class ElasticCloud(VectorDB):
                     "dims": self.dim,
                     **self.case_config.index_param(),
                 },
-            }
+            },
         }
 
         try:
             client.indices.create(index=self.indice, mappings=mappings)
         except Exception as e:
-            log.warning(f"Failed to create indice: {self.indice} error: {str(e)}")
+            log.warning(f"Failed to create indice: {self.indice} error: {e!s}")
             raise e from None
 
     def insert_embeddings(
@@ -94,7 +99,7 @@ class ElasticCloud(VectorDB):
             bulk_insert_res = bulk(self.client, insert_data)
             return (bulk_insert_res[0], None)
         except Exception as e:
-            log.warning(f"Failed to insert data: {self.indice} error: {str(e)}")
+            log.warning(f"Failed to insert data: {self.indice} error: {e!s}")
             return (0, e)
 
     def search_embedding(
@@ -114,16 +119,12 @@ class ElasticCloud(VectorDB):
             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"
-        # is_existed_res = self.client.indices.exists(index=self.indice)
-        # assert is_existed_res.raw == True, "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 [],
+            "filter": [{"range": {self.id_col_name: {"gt": filters["id"]}}}] if filters else [],
             "query_vector": query,
         }
         size = k
@@ -137,26 +138,23 @@ class ElasticCloud(VectorDB):
                 stored_fields="_none_",
                 filter_path=[f"hits.hits.fields.{self.id_col_name}"],
             )
-            res = [h["fields"][self.id_col_name][0] for h in res["hits"]["hits"]]
-
-            return res
+            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: {str(e)}")
+            log.warning(f"Failed to search: {self.indice} error: {e!s}")
             raise e from None
 
-    def optimize(self):
+    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']
+        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}")
-        SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC = 30
         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']:
+            if task_status["completed"]:
                 return
-
-    def ready_to_load(self):
-        """ready_to_load will be called before load in load cases."""
-        pass

vectordb_bench/backend/clients/memorydb/cli.py

@@ -14,9 +14,7 @@ from .. import DB
 
 
 class MemoryDBTypedDict(TypedDict):
-    host: Annotated[
-        str, click.option("--host", type=str, help="Db host", required=True)
-    ]
+    host: Annotated[str, click.option("--host", type=str, help="Db host", required=True)]
     password: Annotated[str, click.option("--password", type=str, help="Db password")]
     port: Annotated[int, click.option("--port", type=int, default=6379, help="Db Port")]
     ssl: Annotated[
@@ -44,7 +42,10 @@ class MemoryDBTypedDict(TypedDict):
             is_flag=True,
             show_default=True,
             default=False,
-            help="Cluster Mode Disabled (CMD), use this flag when testing locally on a single node instance. In production, MemoryDB only supports cluster mode (CME)",
+            help=(
+                "Cluster Mode Disabled (CMD), use this flag when testing locally on a single node instance."
+                " In production, MemoryDB only supports cluster mode (CME)"
+            ),
         ),
     ]
     insert_batch_size: Annotated[
@@ -58,8 +59,7 @@ class MemoryDBTypedDict(TypedDict):
     ]
 
 
-class MemoryDBHNSWTypedDict(CommonTypedDict, MemoryDBTypedDict, HNSWFlavor2):
-    ...
+class MemoryDBHNSWTypedDict(CommonTypedDict, MemoryDBTypedDict, HNSWFlavor2): ...
 
 
 @cli.command()
@@ -82,7 +82,7 @@ def MemoryDB(**parameters: Unpack[MemoryDBHNSWTypedDict]):
             M=parameters["m"],
             ef_construction=parameters["ef_construction"],
             ef_runtime=parameters["ef_runtime"],
-            insert_batch_size=parameters["insert_batch_size"]
+            insert_batch_size=parameters["insert_batch_size"],
         ),
         **parameters,
-    )
+    )

vectordb_bench/backend/clients/memorydb/config.py

@@ -29,7 +29,7 @@ class MemoryDBIndexConfig(BaseModel, DBCaseConfig):
     def parse_metric(self) -> str:
         if self.metric_type == MetricType.L2:
             return "l2"
-        elif self.metric_type == MetricType.IP:
+        if self.metric_type == MetricType.IP:
             return "ip"
         return "cosine"
 
@@ -51,4 +51,4 @@ class MemoryDBHNSWConfig(MemoryDBIndexConfig):
     def search_param(self) -> dict:
         return {
             "ef_runtime": self.ef_runtime,
-        }
+        }

vectordb_bench/backend/clients/memorydb/memorydb.py

@@ -1,30 +1,33 @@
-import logging, time
+import logging
+import time
+from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Any, Generator, Optional, Tuple, Type
-from ..api import VectorDB, DBCaseConfig, IndexType
-from .config import MemoryDBIndexConfig
+from typing import Any
+
+import numpy as np
 import redis
 from redis import Redis
 from redis.cluster import RedisCluster
-from redis.commands.search.field import TagField, VectorField, NumericField
-from redis.commands.search.indexDefinition import IndexDefinition, IndexType
+from redis.commands.search.field import NumericField, TagField, VectorField
+from redis.commands.search.indexDefinition import IndexDefinition
 from redis.commands.search.query import Query
-import numpy as np
 
+from ..api import IndexType, VectorDB
+from .config import MemoryDBIndexConfig
 
 log = logging.getLogger(__name__)
-INDEX_NAME = "index"    # Vector Index Name
+INDEX_NAME = "index"  # Vector Index Name
+
 
 class MemoryDB(VectorDB):
     def __init__(
-            self,
-            dim: int,
-            db_config: dict,
-            db_case_config: MemoryDBIndexConfig,
-            drop_old: bool = False,
-            **kwargs
-        ):
-
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: MemoryDBIndexConfig,
+        drop_old: bool = False,
+        **kwargs,
+    ):
         self.db_config = db_config
         self.case_config = db_case_config
         self.collection_name = INDEX_NAME
@@ -44,10 +47,10 @@ class MemoryDB(VectorDB):
                 info = conn.ft(INDEX_NAME).info()
                 log.info(f"Index info: {info}")
             except redis.exceptions.ResponseError as e:
-                log.error(e)
+                log.warning(e)
                 drop_old = False
                 log.info(f"MemoryDB client drop_old collection: {self.collection_name}")
-            
+
             log.info("Executing FLUSHALL")
             conn.flushall()
 
@@ -59,7 +62,7 @@ class MemoryDB(VectorDB):
                     self.wait_until(self.wait_for_empty_db, 3, "", rc)
                     log.debug(f"Flushall done in the host: {host}")
                     rc.close()
-        
+
         self.make_index(dim, conn)
         conn.close()
         conn = None
@@ -69,7 +72,7 @@ class MemoryDB(VectorDB):
             # check to see if index exists
             conn.ft(INDEX_NAME).info()
         except Exception as e:
-            log.warn(f"Error getting info for index '{INDEX_NAME}': {e}")
+            log.warning(f"Error getting info for index '{INDEX_NAME}': {e}")
             index_param = self.case_config.index_param()
             search_param = self.case_config.search_param()
             vector_parameters = {  # Vector Index Type: FLAT or HNSW
@@ -85,17 +88,19 @@ class MemoryDB(VectorDB):
                 vector_parameters["EF_RUNTIME"] = search_param["ef_runtime"]
 
             schema = (
-                TagField("id"),                   
-                NumericField("metadata"),              
-                VectorField("vector",   # Vector Field Name
-                    "HNSW", vector_parameters
+                TagField("id"),
+                NumericField("metadata"),
+                VectorField(
+                    "vector",  # Vector Field Name
+                    "HNSW",
+                    vector_parameters,
                 ),
             )
 
             definition = IndexDefinition(index_type=IndexType.HASH)
             rs = conn.ft(INDEX_NAME)
             rs.create_index(schema, definition=definition)
-    
+
     def get_client(self, **kwargs):
         """
         Gets either cluster connection or normal connection based on `cmd` flag.
@@ -143,7 +148,7 @@ class MemoryDB(VectorDB):
 
     @contextmanager
     def init(self) -> Generator[None, None, None]:
-        """ create and destory connections to database.
+        """create and destory connections to database.
 
         Examples:
             >>> with self.init():
@@ -152,17 +157,14 @@ class MemoryDB(VectorDB):
         self.conn = self.get_client()
         search_param = self.case_config.search_param()
         if search_param["ef_runtime"]:
-            self.ef_runtime_str = f'EF_RUNTIME {search_param["ef_runtime"]}'
+            self.ef_runtime_str = f"EF_RUNTIME {search_param['ef_runtime']}"
         else:
             self.ef_runtime_str = ""
         yield
         self.conn.close()
         self.conn = None
 
-    def ready_to_load(self) -> bool:
-        pass
-
-    def optimize(self) -> None:
+    def optimize(self, data_size: int | None = None):
         self._post_insert()
 
     def insert_embeddings(
@@ -170,7 +172,7 @@ class MemoryDB(VectorDB):
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs: Any,
-    ) -> Tuple[int, Optional[Exception]]:
+    ) -> tuple[int, Exception | None]:
         """Insert embeddings into the database.
         Should call self.init() first.
         """
@@ -178,12 +180,15 @@ class MemoryDB(VectorDB):
         try:
             with self.conn.pipeline(transaction=False) as pipe:
                 for i, embedding in enumerate(embeddings):
-                    embedding = np.array(embedding).astype(np.float32)
-                    pipe.hset(metadata[i], mapping = {
-                        "id": str(metadata[i]),
-                        "metadata": metadata[i], 
-                        "vector": embedding.tobytes(),
-                    })
+                    ndarr_emb = np.array(embedding).astype(np.float32)
+                    pipe.hset(
+                        metadata[i],
+                        mapping={
+                            "id": str(metadata[i]),
+                            "metadata": metadata[i],
+                            "vector": ndarr_emb.tobytes(),
+                        },
+                    )
                     # Execute the pipe so we don't keep too much in memory at once
                     if (i + 1) % self.insert_batch_size == 0:
                         pipe.execute()
@@ -192,9 +197,9 @@ class MemoryDB(VectorDB):
                 result_len = i + 1
         except Exception as e:
             return 0, e
-        
+
         return result_len, None
-    
+
     def _post_insert(self):
         """Wait for indexing to finish"""
         client = self.get_client(primary=True)
@@ -208,21 +213,17 @@ class MemoryDB(VectorDB):
                 self.wait_until(*args)
                 log.debug(f"Background indexing completed in the host: {host_name}")
                 rc.close()
-    
-    def wait_until(
-        self, condition, interval=5, message="Operation took too long", *args
-    ):
+
+    def wait_until(self, condition: any, interval: int = 5, message: str = "Operation took too long", *args):
         while not condition(*args):
             time.sleep(interval)
-    
+
     def wait_for_no_activity(self, client: redis.RedisCluster | redis.Redis):
-        return (
-            client.info("search")["search_background_indexing_status"] == "NO_ACTIVITY"
-        )
-    
+        return client.info("search")["search_background_indexing_status"] == "NO_ACTIVITY"
+
     def wait_for_empty_db(self, client: redis.RedisCluster | redis.Redis):
         return client.execute_command("DBSIZE") == 0
-    
+
     def search_embedding(
         self,
         query: list[float],
@@ -230,13 +231,13 @@ class MemoryDB(VectorDB):
         filters: dict | None = None,
         timeout: int | None = None,
         **kwargs: Any,
-    ) -> (list[int]):
+    ) -> list[int]:
         assert self.conn is not None
-        
+
         query_vector = np.array(query).astype(np.float32).tobytes()
         query_obj = Query(f"*=>[KNN {k} @vector $vec]").return_fields("id").paging(0, k)
         query_params = {"vec": query_vector}
-        
+
         if filters:
             # benchmark test filters of format: {'metadata': '>=10000', 'id': 10000}
             # gets exact match for id, and range for metadata if they exist in filters
@@ -244,11 +245,19 @@ class MemoryDB(VectorDB):
             # Removing '>=' from the id_value: '>=10000'
             metadata_value = filters.get("metadata")[2:]
             if id_value and metadata_value:
-                query_obj = Query(f"(@metadata:[{metadata_value} +inf] @id:{ {id_value} })=>[KNN {k} @vector $vec]").return_fields("id").paging(0, k)
+                query_obj = (
+                    Query(
+                        f"(@metadata:[{metadata_value} +inf] @id:{ {id_value} })=>[KNN {k} @vector $vec]",
+                    )
+                    .return_fields("id")
+                    .paging(0, k)
+                )
             elif id_value:
-                #gets exact match for id
+                # gets exact match for id
                 query_obj = Query(f"@id:{ {id_value} }=>[KNN {k} @vector $vec]").return_fields("id").paging(0, k)
-            else: #metadata only case, greater than or equal to metadata value
-                query_obj = Query(f"@metadata:[{metadata_value} +inf]=>[KNN {k} @vector $vec]").return_fields("id").paging(0, k)
+            else:  # metadata only case, greater than or equal to metadata value
+                query_obj = (
+                    Query(f"@metadata:[{metadata_value} +inf]=>[KNN {k} @vector $vec]").return_fields("id").paging(0, k)
+                )
         res = self.conn.ft(INDEX_NAME).search(query_obj, query_params)
-        return [int(doc["id"]) for doc in res.docs]
+        return [int(doc["id"]) for doc in res.docs]