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

vectordb_bench/backend/clients/pgvectorscale/pgvectorscale.py

@@ -1,9 +1,9 @@
 """Wrapper around the Pgvectorscale vector database over VectorDB"""
 
 import logging
-import pprint
+from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Any, Generator, Optional, Tuple
+from typing import Any
 
 import numpy as np
 import psycopg
@@ -44,20 +44,21 @@ class PgVectorScale(VectorDB):
         self._primary_field = "id"
         self._vector_field = "embedding"
 
-        self.conn, self.cursor = self._create_connection(**self.db_config)        
+        self.conn, self.cursor = self._create_connection(**self.db_config)
 
         log.info(f"{self.name} config values: {self.db_config}\n{self.case_config}")
         if not any(
             (
                 self.case_config.create_index_before_load,
                 self.case_config.create_index_after_load,
-            )
+            ),
         ):
-            err = f"{self.name} config must create an index using create_index_before_load or create_index_after_load"
-            log.error(err)
-            raise RuntimeError(
-                f"{err}\n{pprint.pformat(self.db_config)}\n{pprint.pformat(self.case_config)}"
+            msg = (
+                f"{self.name} config must create an index using create_index_before_load or create_index_after_load"
+                f"{self.name} config values: {self.db_config}\n{self.case_config}"
             )
+            log.error(msg)
+            raise RuntimeError(msg)
 
         if drop_old:
             self._drop_index()
@@ -72,7 +73,7 @@ class PgVectorScale(VectorDB):
         self.conn = None
 
     @staticmethod
-    def _create_connection(**kwargs) -> Tuple[Connection, Cursor]:
+    def _create_connection(**kwargs) -> tuple[Connection, Cursor]:
         conn = psycopg.connect(**kwargs)
         conn.cursor().execute("CREATE EXTENSION IF NOT EXISTS vectorscale CASCADE")
         conn.commit()
@@ -101,25 +102,25 @@ class PgVectorScale(VectorDB):
                 log.debug(command.as_string(self.cursor))
                 self.cursor.execute(command)
             self.conn.commit()
-        
+
         self._filtered_search = sql.Composed(
             [
                 sql.SQL("SELECT id FROM public.{} WHERE id >= %s ORDER BY embedding ").format(
                     sql.Identifier(self.table_name),
                 ),
                 sql.SQL(self.case_config.search_param()["metric_fun_op"]),
-                sql.SQL(" %s::vector LIMIT %s::int")
-            ]
+                sql.SQL(" %s::vector LIMIT %s::int"),
+            ],
         )
-        
+
         self._unfiltered_search = sql.Composed(
             [
                 sql.SQL("SELECT id FROM public.{} ORDER BY embedding ").format(
-                    sql.Identifier(self.table_name)
+                    sql.Identifier(self.table_name),
                 ),
                 sql.SQL(self.case_config.search_param()["metric_fun_op"]),
                 sql.SQL(" %s::vector LIMIT %s::int"),
-            ]
+            ],
         )
 
         try:
@@ -137,15 +138,12 @@ class PgVectorScale(VectorDB):
 
         self.cursor.execute(
             sql.SQL("DROP TABLE IF EXISTS public.{table_name}").format(
-                table_name=sql.Identifier(self.table_name)
-            )
+                table_name=sql.Identifier(self.table_name),
+            ),
         )
         self.conn.commit()
 
-    def ready_to_load(self):
-        pass
-
-    def optimize(self):
+    def optimize(self, data_size: int | None = None):
         self._post_insert()
 
     def _post_insert(self):
@@ -160,7 +158,7 @@ class PgVectorScale(VectorDB):
         log.info(f"{self.name} client drop index : {self._index_name}")
 
         drop_index_sql = sql.SQL("DROP INDEX IF EXISTS {index_name}").format(
-            index_name=sql.Identifier(self._index_name)
+            index_name=sql.Identifier(self._index_name),
         )
         log.debug(drop_index_sql.as_string(self.cursor))
         self.cursor.execute(drop_index_sql)
@@ -180,36 +178,31 @@ class PgVectorScale(VectorDB):
                     sql.SQL("{option_name} = {val}").format(
                         option_name=sql.Identifier(option_name),
                         val=sql.Identifier(str(option_val)),
-                    )
+                    ),
                 )
-        
+
         num_bits_per_dimension = "2" if self.dim < 900 else "1"
         options.append(
             sql.SQL("{option_name} = {val}").format(
                 option_name=sql.Identifier("num_bits_per_dimension"),
                 val=sql.Identifier(num_bits_per_dimension),
-            )
+            ),
         )
 
-        if any(options):
-            with_clause = sql.SQL("WITH ({});").format(sql.SQL(", ").join(options))
-        else:
-            with_clause = sql.Composed(())
+        with_clause = sql.SQL("WITH ({});").format(sql.SQL(", ").join(options)) if any(options) else sql.Composed(())
 
         index_create_sql = sql.SQL(
             """
-            CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name} 
+            CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name}
             USING {index_type} (embedding {embedding_metric})
-            """
+            """,
         ).format(
             index_name=sql.Identifier(self._index_name),
             table_name=sql.Identifier(self.table_name),
             index_type=sql.Identifier(index_param["index_type"].lower()),
             embedding_metric=sql.Identifier(index_param["metric"]),
         )
-        index_create_sql_with_with_clause = (
-            index_create_sql + with_clause
-        ).join(" ")
+        index_create_sql_with_with_clause = (index_create_sql + with_clause).join(" ")
         log.debug(index_create_sql_with_with_clause.as_string(self.cursor))
         self.cursor.execute(index_create_sql_with_with_clause)
         self.conn.commit()
@@ -223,14 +216,12 @@ class PgVectorScale(VectorDB):
 
             self.cursor.execute(
                 sql.SQL(
-                    "CREATE TABLE IF NOT EXISTS public.{table_name} (id BIGINT PRIMARY KEY, embedding vector({dim}));"
-                ).format(table_name=sql.Identifier(self.table_name), dim=dim)
+                    "CREATE TABLE IF NOT EXISTS public.{table_name} (id BIGINT PRIMARY KEY, embedding vector({dim}));",
+                ).format(table_name=sql.Identifier(self.table_name), dim=dim),
             )
             self.conn.commit()
         except Exception as e:
-            log.warning(
-                f"Failed to create pgvectorscale table: {self.table_name} error: {e}"
-            )
+            log.warning(f"Failed to create pgvectorscale table: {self.table_name} error: {e}")
             raise e from None
 
     def insert_embeddings(
@@ -238,7 +229,7 @@ class PgVectorScale(VectorDB):
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs: Any,
-    ) -> Tuple[int, Optional[Exception]]:
+    ) -> tuple[int, Exception | None]:
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
 
@@ -248,8 +239,8 @@ class PgVectorScale(VectorDB):
 
             with self.cursor.copy(
                 sql.SQL("COPY public.{table_name} FROM STDIN (FORMAT BINARY)").format(
-                    table_name=sql.Identifier(self.table_name)
-                )
+                    table_name=sql.Identifier(self.table_name),
+                ),
             ) as copy:
                 copy.set_types(["bigint", "vector"])
                 for i, row in enumerate(metadata_arr):
@@ -261,9 +252,7 @@ class PgVectorScale(VectorDB):
 
             return len(metadata), None
         except Exception as e:
-            log.warning(
-                f"Failed to insert data into pgvector table ({self.table_name}), error: {e}"
-            )
+            log.warning(f"Failed to insert data into pgvector table ({self.table_name}), error: {e}")
             return 0, e
 
     def search_embedding(
@@ -280,11 +269,12 @@ class PgVectorScale(VectorDB):
         if filters:
             gt = filters.get("id")
             result = self.cursor.execute(
-                self._filtered_search, (gt, q, k), prepare=True, binary=True
+                self._filtered_search,
+                (gt, q, k),
+                prepare=True,
+                binary=True,
             )
         else:
-            result = self.cursor.execute(
-                self._unfiltered_search, (q, k), prepare=True, binary=True
-            )
+            result = self.cursor.execute(self._unfiltered_search, (q, k), prepare=True, binary=True)
 
         return [int(i[0]) for i in result.fetchall()]

vectordb_bench/backend/clients/pinecone/config.py

@@ -1,4 +1,5 @@
 from pydantic import SecretStr
+
 from ..api import DBConfig
 
 

vectordb_bench/backend/clients/pinecone/pinecone.py

@@ -2,11 +2,11 @@
 
 import logging
 from contextlib import contextmanager
-from typing import Type
+
 import pinecone
-from ..api import VectorDB, DBConfig, DBCaseConfig, EmptyDBCaseConfig, IndexType
-from .config import PineconeConfig
 
+from ..api import DBCaseConfig, DBConfig, EmptyDBCaseConfig, IndexType, VectorDB
+from .config import PineconeConfig
 
 log = logging.getLogger(__name__)
 
@@ -17,7 +17,7 @@ PINECONE_MAX_SIZE_PER_BATCH = 2 * 1024 * 1024  # 2MB
 class Pinecone(VectorDB):
     def __init__(
         self,
-        dim,
+        dim: int,
         db_config: dict,
         db_case_config: DBCaseConfig,
         drop_old: bool = False,
@@ -27,7 +27,7 @@ class Pinecone(VectorDB):
         self.index_name = db_config.get("index_name", "")
         self.api_key = db_config.get("api_key", "")
         self.batch_size = int(
-            min(PINECONE_MAX_SIZE_PER_BATCH / (dim * 5), PINECONE_MAX_NUM_PER_BATCH)
+            min(PINECONE_MAX_SIZE_PER_BATCH / (dim * 5), PINECONE_MAX_NUM_PER_BATCH),
         )
 
         pc = pinecone.Pinecone(api_key=self.api_key)
@@ -37,9 +37,8 @@ class Pinecone(VectorDB):
             index_stats = index.describe_index_stats()
             index_dim = index_stats["dimension"]
             if index_dim != dim:
-                raise ValueError(
-                    f"Pinecone index {self.index_name} dimension mismatch, expected {index_dim} got {dim}"
-                )
+                msg = f"Pinecone index {self.index_name} dimension mismatch, expected {index_dim} got {dim}"
+                raise ValueError(msg)
             for namespace in index_stats["namespaces"]:
                 log.info(f"Pinecone index delete namespace: {namespace}")
                 index.delete(delete_all=True, namespace=namespace)
@@ -47,11 +46,11 @@ class Pinecone(VectorDB):
         self._metadata_key = "meta"
 
     @classmethod
-    def config_cls(cls) -> Type[DBConfig]:
+    def config_cls(cls) -> type[DBConfig]:
         return PineconeConfig
 
     @classmethod
-    def case_config_cls(cls, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
+    def case_config_cls(cls, index_type: IndexType | None = None) -> type[DBCaseConfig]:
         return EmptyDBCaseConfig
 
     @contextmanager
@@ -60,10 +59,7 @@ class Pinecone(VectorDB):
         self.index = pc.Index(self.index_name)
         yield
 
-    def ready_to_load(self):
-        pass
-
-    def optimize(self):
+    def optimize(self, data_size: int | None = None):
         pass
 
     def insert_embeddings(
@@ -76,9 +72,7 @@ class Pinecone(VectorDB):
         insert_count = 0
         try:
             for batch_start_offset in range(0, len(embeddings), self.batch_size):
-                batch_end_offset = min(
-                    batch_start_offset + self.batch_size, len(embeddings)
-                )
+                batch_end_offset = min(batch_start_offset + self.batch_size, len(embeddings))
                 insert_datas = []
                 for i in range(batch_start_offset, batch_end_offset):
                     insert_data = (
@@ -100,10 +94,7 @@ class Pinecone(VectorDB):
         filters: dict | None = None,
         timeout: int | None = None,
     ) -> list[int]:
-        if filters is None:
-            pinecone_filters = {}
-        else:
-            pinecone_filters = {self._metadata_key: {"$gte": filters["id"]}}
+        pinecone_filters = {} if filters is None else {self._metadata_key: {"$gte": filters["id"]}}
         try:
             res = self.index.query(
                 top_k=k,
@@ -111,7 +102,6 @@ class Pinecone(VectorDB):
                 filter=pinecone_filters,
             )["matches"]
         except Exception as e:
-            print(f"Error querying index: {e}")
-            raise e
-        id_res = [int(one_res["id"]) for one_res in res]
-        return id_res
+            log.warning(f"Error querying index: {e}")
+            raise e from e
+        return [int(one_res["id"]) for one_res in res]

vectordb_bench/backend/clients/qdrant_cloud/config.py

@@ -1,7 +1,7 @@
-from pydantic import BaseModel, SecretStr
+from pydantic import BaseModel, SecretStr, validator
+
+from ..api import DBCaseConfig, DBConfig, MetricType
 
-from ..api import DBConfig, DBCaseConfig, MetricType
-from pydantic import validator
 
 # Allowing `api_key` to be left empty, to ensure compatibility with the open-source Qdrant.
 class QdrantConfig(DBConfig):
@@ -16,17 +16,19 @@ class QdrantConfig(DBConfig):
                 "api_key": self.api_key.get_secret_value(),
                 "prefer_grpc": True,
             }
-        else:
-            return {"url": self.url.get_secret_value(),}
-        
+        return {
+            "url": self.url.get_secret_value(),
+        }
+
     @validator("*")
-    def not_empty_field(cls, v, field):
+    def not_empty_field(cls, v: any, field: any):
         if field.name in ["api_key", "db_label"]:
             return v
-        if isinstance(v, (str, SecretStr)) and len(v) == 0:
+        if isinstance(v, str | SecretStr) and len(v) == 0:
             raise ValueError("Empty string!")
         return v
 
+
 class QdrantIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType | None = None
 
@@ -40,8 +42,7 @@ class QdrantIndexConfig(BaseModel, DBCaseConfig):
         return "Cosine"
 
     def index_param(self) -> dict:
-        params = {"distance": self.parse_metric()}
-        return params
+        return {"distance": self.parse_metric()}
 
     def search_param(self) -> dict:
         return {}

vectordb_bench/backend/clients/qdrant_cloud/qdrant_cloud.py

@@ -4,23 +4,26 @@ import logging
 import time
 from contextlib import contextmanager
 
-from ..api import VectorDB, DBCaseConfig
+from qdrant_client import QdrantClient
 from qdrant_client.http.models import (
-    CollectionStatus,
-    VectorParams,
-    PayloadSchemaType,
     Batch,
-    Filter,
+    CollectionStatus,
     FieldCondition,
+    Filter,
+    PayloadSchemaType,
     Range,
+    VectorParams,
 )
 
-from qdrant_client import QdrantClient
-
+from ..api import DBCaseConfig, VectorDB
 
 log = logging.getLogger(__name__)
 
 
+SECONDS_WAITING_FOR_INDEXING_API_CALL = 5
+QDRANT_BATCH_SIZE = 500
+
+
 class QdrantCloud(VectorDB):
     def __init__(
         self,
@@ -57,16 +60,11 @@ class QdrantCloud(VectorDB):
         self.qdrant_client = QdrantClient(**self.db_config)
         yield
         self.qdrant_client = None
-        del(self.qdrant_client)
-
-    def ready_to_load(self):
-        pass
+        del self.qdrant_client
 
-
-    def optimize(self):
+    def optimize(self, data_size: int | None = None):
         assert self.qdrant_client, "Please call self.init() before"
         # wait for vectors to be fully indexed
-        SECONDS_WAITING_FOR_INDEXING_API_CALL = 5
         try:
             while True:
                 info = self.qdrant_client.get_collection(self.collection_name)
@@ -74,19 +72,26 @@ class QdrantCloud(VectorDB):
                 if info.status != CollectionStatus.GREEN:
                     continue
                 if info.status == CollectionStatus.GREEN:
-                    log.info(f"Stored vectors: {info.vectors_count}, Indexed vectors: {info.indexed_vectors_count}, Collection status: {info.indexed_vectors_count}")
+                    msg = (
+                        f"Stored vectors: {info.vectors_count}, Indexed vectors: {info.indexed_vectors_count}, "
+                        f"Collection status: {info.indexed_vectors_count}"
+                    )
+                    log.info(msg)
                     return
         except Exception as e:
             log.warning(f"QdrantCloud ready to search error: {e}")
             raise e from None
 
-    def _create_collection(self, dim, qdrant_client: int):
+    def _create_collection(self, dim: int, qdrant_client: QdrantClient):
         log.info(f"Create collection: {self.collection_name}")
 
         try:
             qdrant_client.create_collection(
                 collection_name=self.collection_name,
-                vectors_config=VectorParams(size=dim, distance=self.case_config.index_param()["distance"])
+                vectors_config=VectorParams(
+                    size=dim,
+                    distance=self.case_config.index_param()["distance"],
+                ),
             )
 
             qdrant_client.create_payload_index(
@@ -109,13 +114,12 @@ class QdrantCloud(VectorDB):
     ) -> (int, Exception):
         """Insert embeddings into Milvus. should call self.init() first"""
         assert self.qdrant_client is not None
-        QDRANT_BATCH_SIZE = 500
         try:
             # TODO: counts
             for offset in range(0, len(embeddings), QDRANT_BATCH_SIZE):
-                vectors = embeddings[offset: offset + QDRANT_BATCH_SIZE]
-                ids = metadata[offset: offset + QDRANT_BATCH_SIZE]
-                payloads=[{self._primary_field: v} for v in ids]
+                vectors = embeddings[offset : offset + QDRANT_BATCH_SIZE]
+                ids = metadata[offset : offset + QDRANT_BATCH_SIZE]
+                payloads = [{self._primary_field: v} for v in ids]
                 _ = self.qdrant_client.upsert(
                     collection_name=self.collection_name,
                     wait=True,
@@ -142,21 +146,23 @@ class QdrantCloud(VectorDB):
         f = None
         if filters:
             f = Filter(
-                must=[FieldCondition(
-                    key = self._primary_field,
-                    range = Range(
-                        gt=filters.get('id'),
+                must=[
+                    FieldCondition(
+                        key=self._primary_field,
+                        range=Range(
+                            gt=filters.get("id"),
+                        ),
                     ),
-                )]
+                ],
             )
 
-        res = self.qdrant_client.search(
-            collection_name=self.collection_name,
-            query_vector=query,
-            limit=k,
-            query_filter=f,
-            #  with_payload=True,
-        ),
+        res = (
+            self.qdrant_client.search(
+                collection_name=self.collection_name,
+                query_vector=query,
+                limit=k,
+                query_filter=f,
+            ),
+        )
 
-        ret = [result.id for result in res[0]]
-        return ret
+        return [result.id for result in res[0]]

vectordb_bench/backend/clients/redis/cli.py

@@ -3,9 +3,6 @@ from typing import Annotated, TypedDict, Unpack
 import click
 from pydantic import SecretStr
 
-from .config import RedisHNSWConfig
-
-
 from ....cli.cli import (
     CommonTypedDict,
     HNSWFlavor2,
@@ -14,12 +11,11 @@ from ....cli.cli import (
     run,
 )
 from .. import DB
+from .config import RedisHNSWConfig
 
 
 class RedisTypedDict(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[
@@ -52,27 +48,25 @@ class RedisTypedDict(TypedDict):
     ]
 
 
-class RedisHNSWTypedDict(CommonTypedDict, RedisTypedDict, HNSWFlavor2):
-    ...
+class RedisHNSWTypedDict(CommonTypedDict, RedisTypedDict, HNSWFlavor2): ...
 
 
 @cli.command()
 @click_parameter_decorators_from_typed_dict(RedisHNSWTypedDict)
 def Redis(**parameters: Unpack[RedisHNSWTypedDict]):
     from .config import RedisConfig
+
     run(
         db=DB.Redis,
         db_config=RedisConfig(
             db_label=parameters["db_label"],
-            password=SecretStr(parameters["password"])
-            if parameters["password"]
-            else None,
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
             host=SecretStr(parameters["host"]),
             port=parameters["port"],
             ssl=parameters["ssl"],
             ssl_ca_certs=parameters["ssl_ca_certs"],
             cmd=parameters["cmd"],
-        ),        
+        ),
         db_case_config=RedisHNSWConfig(
             M=parameters["m"],
             efConstruction=parameters["ef_construction"],

vectordb_bench/backend/clients/redis/config.py

@@ -1,10 +1,12 @@
-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 RedisConfig(DBConfig):
     password: SecretStr | None = None
     host: SecretStr
-    port: int | None = None 
+    port: int | None = None
 
     def to_dict(self) -> dict:
         return {
@@ -12,7 +14,6 @@ class RedisConfig(DBConfig):
             "port": self.port,
             "password": self.password.get_secret_value() if self.password is not None else None,
         }
-        
 
 
 class RedisIndexConfig(BaseModel):
@@ -24,7 +25,8 @@ class RedisIndexConfig(BaseModel):
         if not self.metric_type:
             return ""
         return self.metric_type.value
-    
+
+
 class RedisHNSWConfig(RedisIndexConfig, DBCaseConfig):
     M: int
     efConstruction: int