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

vectordb_bench/backend/clients/redis/redis.py

@@ -1,37 +1,40 @@
 import logging
 from contextlib import contextmanager
-from typing import Any, Type
-from ..api import VectorDB, DBConfig, DBCaseConfig, EmptyDBCaseConfig, IndexType
-from .config import RedisConfig
+from typing import Any
+
+import numpy as np
 import redis
-from redis.commands.search.field import TagField, VectorField, NumericField
+from redis.commands.search.field import NumericField, TagField, VectorField
 from redis.commands.search.indexDefinition import IndexDefinition, IndexType
 from redis.commands.search.query import Query
-import numpy as np
 
+from ..api import DBCaseConfig, VectorDB
 
 log = logging.getLogger(__name__)
-INDEX_NAME = "index"                              # Vector Index Name
+INDEX_NAME = "index"  # Vector Index Name
+
 
 class Redis(VectorDB):
     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 = INDEX_NAME
 
         # Create a redis connection, if db has password configured, add it to the connection here and in init():
-        password=self.db_config["password"]
-        conn = redis.Redis(host=self.db_config["host"], port=self.db_config["port"], password=password, db=0)
-        
+        password = self.db_config["password"]
+        conn = redis.Redis(
+            host=self.db_config["host"],
+            port=self.db_config["port"],
+            password=password,
+            db=0,
+        )
 
         if drop_old:
             try:
@@ -40,7 +43,7 @@ class Redis(VectorDB):
             except redis.exceptions.ResponseError:
                 drop_old = False
                 log.info(f"Redis client drop_old collection: {self.collection_name}")
-        
+
         self.make_index(dim, conn)
         conn.close()
         conn = None
@@ -49,16 +52,20 @@ class Redis(VectorDB):
         try:
             # check to see if index exists
             conn.ft(INDEX_NAME).info()
-        except: 
+        except Exception:
             schema = (
-                TagField("id"),                   
-                NumericField("metadata"),              
-                VectorField("vector",                  # Vector Field Name
-                    "HNSW", {                          # Vector Index Type: FLAT or HNSW
-                        "TYPE": "FLOAT32",             # FLOAT32 or FLOAT64
-                        "DIM": vector_dimensions,      # Number of Vector Dimensions
-                        "DISTANCE_METRIC": "COSINE",   # Vector Search Distance Metric
-                    }
+                TagField("id"),
+                NumericField("metadata"),
+                VectorField(
+                    "vector",  # Vector Field Name
+                    "HNSW",  # Vector Index Type: FLAT or HNSW
+                    {
+                        "TYPE": "FLOAT32",  # FLOAT32 or FLOAT64
+                        "DIM": vector_dimensions,  # Number of Vector Dimensions
+                        "DISTANCE_METRIC": "COSINE",  # Vector Search Distance Metric
+                        "M": self.case_config.index_param()["params"]["M"],
+                        "EF_CONSTRUCTION": self.case_config.index_param()["params"]["efConstruction"],
+                    },
                 ),
             )
 
@@ -69,61 +76,62 @@ class Redis(VectorDB):
 
     @contextmanager
     def init(self) -> None:
-        """ create and destory connections to database.
+        """create and destory connections to database.
 
         Examples:
             >>> with self.init():
             >>>     self.insert_embeddings()
         """
-        self.conn = redis.Redis(host=self.db_config["host"], port=self.db_config["port"], password=self.db_config["password"], db=0)
+        self.conn = redis.Redis(
+            host=self.db_config["host"],
+            port=self.db_config["port"],
+            password=self.db_config["password"],
+            db=0,
+        )
         yield
         self.conn.close()
         self.conn = None
 
-
     def ready_to_search(self) -> bool:
         """Check if the database is ready to search."""
-        pass
-    
-
-    def ready_to_load(self) -> bool:
-        pass
 
-    def optimize(self) -> None:
+    def optimize(self, data_size: int | None = None):
         pass
 
-
     def insert_embeddings(
         self,
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs: Any,
-    ) -> (int, Exception):
+    ) -> tuple[int, Exception]:
         """Insert embeddings into the database.
         Should call self.init() first.
         """
 
-        batch_size = 1000 # Adjust this as needed, but don't make too big
+        batch_size = 1000  # Adjust this as needed, but don't make too big
         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 % batch_size == 0:
-                        res = pipe.execute()
+                        _ = pipe.execute()
 
-                res = pipe.execute()
+                _ = pipe.execute()
                 result_len = i + 1
         except Exception as e:
             return 0, e
-        
+
         return result_len, None
-    
+
     def search_embedding(
         self,
         query: list[float],
@@ -131,28 +139,53 @@ class Redis(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 as score]").sort_by("score").return_fields("id", "score").paging(0, k).dialect(2)
+        ef_runtime = self.case_config.search_param()["params"]["ef"]
+        query_obj = (
+            Query(f"*=>[KNN {k} @vector $vec EF_RUNTIME {ef_runtime} as score]")
+            .sort_by("score")
+            .return_fields("id", "score")
+            .paging(0, k)
+            .dialect(2)
+        )
         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
             id_value = filters.get("id")
             metadata_value = filters.get("metadata")
             if id_value and metadata_value:
-                query_obj = Query(f"(@metadata:[{metadata_value} +inf] @id:{ {id_value} })=>[KNN {k} @vector $vec as score]").sort_by("score").return_fields("id", "score").paging(0, k).dialect(2)
+                query_obj = (
+                    Query(
+                        f"(@metadata:[{metadata_value} +inf] @id:{ {id_value} })=>[KNN {k} ",
+                        f"@vector $vec EF_RUNTIME {ef_runtime} as score]",
+                    )
+                    .sort_by("score")
+                    .return_fields("id", "score")
+                    .paging(0, k)
+                    .dialect(2)
+                )
             elif id_value:
-                #gets exact match for id
-                query_obj = Query(f"@id:{ {id_value} }=>[KNN {k} @vector $vec as score]").sort_by("score").return_fields("id", "score").paging(0, k).dialect(2)
-            else: #metadata only case, greater than or equal to metadata value
-                query_obj = Query(f"@metadata:[{metadata_value} +inf]=>[KNN {k} @vector $vec as score]").sort_by("score").return_fields("id", "score").paging(0, k).dialect(2) 
+                # gets exact match for id
+                query_obj = (
+                    Query(f"@id:{ {id_value} }=>[KNN {k} @vector $vec EF_RUNTIME {ef_runtime} as score]")
+                    .sort_by("score")
+                    .return_fields("id", "score")
+                    .paging(0, k)
+                    .dialect(2)
+                )
+            else:  # metadata only case, greater than or equal to metadata value
+                query_obj = (
+                    Query(f"@metadata:[{metadata_value} +inf]=>[KNN {k} @vector $vec EF_RUNTIME {ef_runtime} as score]")
+                    .sort_by("score")
+                    .return_fields("id", "score")
+                    .paging(0, k)
+                    .dialect(2)
+                )
         res = self.conn.ft(INDEX_NAME).search(query_obj, query_params)
         # doc in res of format {'id': '9831', 'payload': None, 'score': '1.19209289551e-07'}
         return [int(doc["id"]) for doc in res.docs]
-
-    
-        

vectordb_bench/backend/clients/test/cli.py

@@ -10,15 +10,14 @@ from .. import DB
 from ..test.config import TestConfig, TestIndexConfig
 
 
-class TestTypedDict(CommonTypedDict):
-    ...
+class TestTypedDict(CommonTypedDict): ...
 
 
 @cli.command()
 @click_parameter_decorators_from_typed_dict(TestTypedDict)
 def Test(**parameters: Unpack[TestTypedDict]):
     run(
-        db=DB.NewClient,
+        db=DB.Test,
         db_config=TestConfig(db_label=parameters["db_label"]),
         db_case_config=TestIndexConfig(),
         **parameters,

vectordb_bench/backend/clients/test/config.py

@@ -1,6 +1,6 @@
-from pydantic import BaseModel, SecretStr
+from pydantic import BaseModel
 
-from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+from ..api import DBCaseConfig, DBConfig, MetricType
 
 
 class TestConfig(DBConfig):

vectordb_bench/backend/clients/test/test.py

@@ -1,6 +1,7 @@
 import logging
+from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Any, Generator, Optional, Tuple
+from typing import Any
 
 from ..api import DBCaseConfig, VectorDB
 
@@ -32,10 +33,7 @@ class Test(VectorDB):
 
         yield
 
-    def ready_to_load(self) -> bool:
-        return True
-
-    def optimize(self) -> None:
+    def optimize(self, data_size: int | None = None):
         pass
 
     def insert_embeddings(
@@ -43,11 +41,10 @@ class Test(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.
         """
-        raise RuntimeError("Not implemented")
         return len(metadata), None
 
     def search_embedding(
@@ -58,5 +55,4 @@ class Test(VectorDB):
         timeout: int | None = None,
         **kwargs: Any,
     ) -> list[int]:
-        raise NotImplementedError
-        return [i for i in range(k)]
+        return list(range(k))

vectordb_bench/backend/clients/weaviate_cloud/cli.py

@@ -14,7 +14,8 @@ from .. import DB
 
 class WeaviateTypedDict(CommonTypedDict):
     api_key: Annotated[
-        str, click.option("--api-key", type=str, help="Weaviate api key", required=True)
+        str,
+        click.option("--api-key", type=str, help="Weaviate api key", required=True),
     ]
     url: Annotated[
         str,
@@ -34,8 +35,6 @@ def Weaviate(**parameters: Unpack[WeaviateTypedDict]):
             api_key=SecretStr(parameters["api_key"]),
             url=SecretStr(parameters["url"]),
         ),
-        db_case_config=WeaviateIndexConfig(
-            ef=256, efConstruction=256, maxConnections=16
-        ),
+        db_case_config=WeaviateIndexConfig(ef=256, efConstruction=256, maxConnections=16),
         **parameters,
     )

vectordb_bench/backend/clients/weaviate_cloud/config.py

@@ -1,6 +1,6 @@
 from pydantic import BaseModel, SecretStr
 
-from ..api import DBConfig, DBCaseConfig, MetricType
+from ..api import DBCaseConfig, DBConfig, MetricType
 
 
 class WeaviateConfig(DBConfig):
@@ -23,7 +23,7 @@ class WeaviateIndexConfig(BaseModel, DBCaseConfig):
     def parse_metric(self) -> str:
         if self.metric_type == MetricType.L2:
             return "l2-squared"
-        elif self.metric_type == MetricType.IP:
+        if self.metric_type == MetricType.IP:
             return "dot"
         return "cosine"
 

vectordb_bench/backend/clients/weaviate_cloud/weaviate_cloud.py

@@ -1,13 +1,13 @@
 """Wrapper around the Weaviate vector database over VectorDB"""
 
 import logging
-from typing import Iterable
+from collections.abc import Iterable
 from contextlib import contextmanager
 
 import weaviate
 from weaviate.exceptions import WeaviateBaseError
 
-from ..api import VectorDB, DBCaseConfig
+from ..api import DBCaseConfig, VectorDB
 
 log = logging.getLogger(__name__)
 
@@ -23,7 +23,13 @@ class WeaviateCloud(VectorDB):
         **kwargs,
     ):
         """Initialize wrapper around the weaviate vector database."""
-        db_config.update({"auth_client_secret": weaviate.AuthApiKey(api_key=db_config.get("auth_client_secret"))})
+        db_config.update(
+            {
+                "auth_client_secret": weaviate.AuthApiKey(
+                    api_key=db_config.get("auth_client_secret"),
+                ),
+            },
+        )
         self.db_config = db_config
         self.case_config = db_case_config
         self.collection_name = collection_name
@@ -33,6 +39,7 @@ class WeaviateCloud(VectorDB):
         self._index_name = "vector_idx"
 
         from weaviate import Client
+
         client = Client(**db_config)
         if drop_old:
             try:
@@ -40,7 +47,7 @@ class WeaviateCloud(VectorDB):
                     log.info(f"weaviate client drop_old collection: {self.collection_name}")
                     client.schema.delete_class(self.collection_name)
             except WeaviateBaseError as e:
-                log.warning(f"Failed to drop collection: {self.collection_name} error: {str(e)}")
+                log.warning(f"Failed to drop collection: {self.collection_name} error: {e!s}")
                 raise e from None
         self._create_collection(client)
         client = None
@@ -54,20 +61,20 @@ class WeaviateCloud(VectorDB):
             >>>     self.search_embedding()
         """
         from weaviate import Client
+
         self.client = Client(**self.db_config)
         yield
         self.client = None
-        del(self.client)
-
-    def ready_to_load(self):
-        """Should call insert first, do nothing"""
-        pass
+        del self.client
 
-    def optimize(self):
+    def optimize(self, data_size: int | None = None):
         assert self.client.schema.exists(self.collection_name)
-        self.client.schema.update_config(self.collection_name, {"vectorIndexConfig": self.case_config.search_param() } )
+        self.client.schema.update_config(
+            self.collection_name,
+            {"vectorIndexConfig": self.case_config.search_param()},
+        )
 
-    def _create_collection(self, client):
+    def _create_collection(self, client: weaviate.Client) -> None:
         if not client.schema.exists(self.collection_name):
             log.info(f"Create collection: {self.collection_name}")
             class_obj = {
@@ -78,13 +85,13 @@ class WeaviateCloud(VectorDB):
                         "dataType": ["int"],
                         "name": self._scalar_field,
                     },
-                ]
+                ],
             }
             class_obj["vectorIndexConfig"] = self.case_config.index_param()
             try:
                 client.schema.create_class(class_obj)
             except WeaviateBaseError as e:
-                log.warning(f"Failed to create collection: {self.collection_name} error: {str(e)}")
+                log.warning(f"Failed to create collection: {self.collection_name} error: {e!s}")
                 raise e from None
 
     def insert_embeddings(
@@ -102,15 +109,17 @@ class WeaviateCloud(VectorDB):
                 batch.dynamic = True
                 res = []
                 for i in range(len(metadata)):
-                    res.append(batch.add_data_object(
-                        {self._scalar_field: metadata[i]},
-                        class_name=self.collection_name,
-                        vector=embeddings[i]
-                    ))
+                    res.append(
+                        batch.add_data_object(
+                            {self._scalar_field: metadata[i]},
+                            class_name=self.collection_name,
+                            vector=embeddings[i],
+                        ),
+                    )
                     insert_count += 1
                 return (len(res), None)
         except WeaviateBaseError as e:
-            log.warning(f"Failed to insert data, error: {str(e)}")
+            log.warning(f"Failed to insert data, error: {e!s}")
             return (insert_count, e)
 
     def search_embedding(
@@ -125,12 +134,17 @@ class WeaviateCloud(VectorDB):
         """
         assert self.client.schema.exists(self.collection_name)
 
-        query_obj = self.client.query.get(self.collection_name, [self._scalar_field]).with_additional("distance").with_near_vector({"vector": query}).with_limit(k)
+        query_obj = (
+            self.client.query.get(self.collection_name, [self._scalar_field])
+            .with_additional("distance")
+            .with_near_vector({"vector": query})
+            .with_limit(k)
+        )
         if filters:
             where_filter = {
                 "path": "key",
                 "operator": "GreaterThanEqual",
-                "valueInt": filters.get('id')
+                "valueInt": filters.get("id"),
             }
             query_obj = query_obj.with_where(where_filter)
 
@@ -138,7 +152,4 @@ class WeaviateCloud(VectorDB):
         res = query_obj.do()
 
         # Organize results.
-        ret = [result[self._scalar_field] for result in res["data"]["Get"][self.collection_name]]
-
-        return ret
-
+        return [result[self._scalar_field] for result in res["data"]["Get"][self.collection_name]]

vectordb_bench/backend/clients/zilliz_cloud/cli.py

@@ -1,33 +1,36 @@
+import os
 from typing import Annotated, Unpack
 
 import click
-import os
 from pydantic import SecretStr
 
+from vectordb_bench.backend.clients import DB
 from vectordb_bench.cli.cli import (
     CommonTypedDict,
     cli,
     click_parameter_decorators_from_typed_dict,
     run,
 )
-from vectordb_bench.backend.clients import DB
 
 
 class ZillizTypedDict(CommonTypedDict):
     uri: Annotated[
-        str, click.option("--uri", type=str, help="uri connection string", required=True)
+        str,
+        click.option("--uri", type=str, help="uri connection string", required=True),
     ]
     user_name: Annotated[
-        str, click.option("--user-name", type=str, help="Db username", required=True)
+        str,
+        click.option("--user-name", type=str, help="Db username", required=True),
     ]
     password: Annotated[
         str,
-        click.option("--password",
-                     type=str,
-                     help="Zilliz password",
-                     default=lambda: os.environ.get("ZILLIZ_PASSWORD", ""),
-                     show_default="$ZILLIZ_PASSWORD",
-                     ),
+        click.option(
+            "--password",
+            type=str,
+            help="Zilliz password",
+            default=lambda: os.environ.get("ZILLIZ_PASSWORD", ""),
+            show_default="$ZILLIZ_PASSWORD",
+        ),
     ]
     level: Annotated[
         str,
@@ -38,7 +41,7 @@ class ZillizTypedDict(CommonTypedDict):
 @cli.command()
 @click_parameter_decorators_from_typed_dict(ZillizTypedDict)
 def ZillizAutoIndex(**parameters: Unpack[ZillizTypedDict]):
-    from .config import ZillizCloudConfig, AutoIndexConfig
+    from .config import AutoIndexConfig, ZillizCloudConfig
 
     run(
         db=DB.ZillizCloud,

vectordb_bench/backend/clients/zilliz_cloud/config.py

@@ -1,7 +1,7 @@
 from pydantic import SecretStr
 
 from ..api import DBCaseConfig, DBConfig
-from ..milvus.config import MilvusIndexConfig, IndexType
+from ..milvus.config import IndexType, MilvusIndexConfig
 
 
 class ZillizCloudConfig(DBConfig):
@@ -33,7 +33,5 @@ class AutoIndexConfig(MilvusIndexConfig, DBCaseConfig):
             "metric_type": self.parse_metric(),
             "params": {
                 "level": self.level,
-            }
+            },
         }
-
-

vectordb_bench/backend/clients/zilliz_cloud/zilliz_cloud.py

@@ -1,7 +1,7 @@
 """Wrapper around the ZillizCloud vector database over VectorDB"""
 
-from ..milvus.milvus import Milvus
 from ..api import DBCaseConfig
+from ..milvus.milvus import Milvus
 
 
 class ZillizCloud(Milvus):