vectordb-bench 0.0.23__py3-none-any.whl → 0.0.24__py3-none-any.whl

vectordb_bench/backend/clients/mariadb/mariadb.py

@@ -1,27 +1,25 @@
-from ..api import VectorDB
-
 import logging
 from contextlib import contextmanager
-from typing import Any, Optional, Tuple
-from ..api import VectorDB
-from .config import MariaDBConfigDict, MariaDBIndexConfig
-import numpy as np
 
 import mariadb
+import numpy as np
+
+from ..api import VectorDB
+from .config import MariaDBConfigDict, MariaDBIndexConfig
 
 log = logging.getLogger(__name__)
 
+
 class MariaDB(VectorDB):
     def __init__(
-            self,
-            dim: int,
-            db_config: MariaDBConfigDict,
-            db_case_config: MariaDBIndexConfig,
-            collection_name: str = "vec_collection",
-            drop_old: bool = False,
-            **kwargs,
-        ):
-
+        self,
+        dim: int,
+        db_config: MariaDBConfigDict,
+        db_case_config: MariaDBIndexConfig,
+        collection_name: str = "vec_collection",
+        drop_old: bool = False,
+        **kwargs,
+    ):
         self.name = "MariaDB"
         self.db_config = db_config
         self.case_config = db_case_config
@@ -31,7 +29,7 @@ class MariaDB(VectorDB):
 
         # construct basic units
         self.conn, self.cursor = self._create_connection(**self.db_config)
-        
+
         if drop_old:
             self._drop_db()
             self._create_db_table(dim)
@@ -41,9 +39,8 @@ class MariaDB(VectorDB):
         self.cursor = None
         self.conn = None
 
-
     @staticmethod
-    def _create_connection(**kwargs) -> Tuple[mariadb.Connection, mariadb.Cursor]:
+    def _create_connection(**kwargs) -> tuple[mariadb.Connection, mariadb.Cursor]:
         conn = mariadb.connect(**kwargs)
         cursor = conn.cursor()
 
@@ -52,7 +49,6 @@ class MariaDB(VectorDB):
 
         return conn, cursor
 
-
     def _drop_db(self):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
@@ -77,24 +73,23 @@ class MariaDB(VectorDB):
             log.info(f"{self.name} client create table : {self.table_name}")
             self.cursor.execute(f"USE {self.db_name}")
 
-            self.cursor.execute(f"""
+            self.cursor.execute(
+                f"""
               CREATE TABLE {self.table_name} (
                 id INT PRIMARY KEY,
                 v VECTOR({self.dim}) NOT NULL
               ) ENGINE={index_param["storage_engine"]}
-            """)
+            """
+            )
             self.cursor.execute("COMMIT")
 
         except Exception as e:
-            log.warning(
-                f"Failed to create table: {self.table_name} error: {e}"
-            )
+            log.warning(f"Failed to create table: {self.table_name} error: {e}")
             raise e from None
 
-
     @contextmanager
-    def init(self) -> None:
-        """ create and destory connections to database.
+    def init(self):
+        """create and destory connections to database.
 
         Examples:
             >>> with self.init():
@@ -109,15 +104,21 @@ class MariaDB(VectorDB):
         self.cursor.execute("SET GLOBAL max_allowed_packet = 1073741824")
 
         if index_param["index_type"] == "HNSW":
-            if index_param["max_cache_size"] != None:
-                self.cursor.execute(f"SET GLOBAL mhnsw_max_cache_size = {index_param["max_cache_size"]}")
-            if search_param["ef_search"] != None:
-                self.cursor.execute(f"SET mhnsw_ef_search = {search_param["ef_search"]}")
+            if index_param["max_cache_size"] is not None:
+                self.cursor.execute(f"SET GLOBAL mhnsw_max_cache_size = {index_param['max_cache_size']}")
+            if search_param["ef_search"] is not None:
+                self.cursor.execute(f"SET mhnsw_ef_search = {search_param['ef_search']}")
             self.cursor.execute("COMMIT")
 
-        self.insert_sql = f"INSERT INTO {self.db_name}.{self.table_name} (id, v) VALUES (%s, %s)"
-        self.select_sql = f"SELECT id FROM {self.db_name}.{self.table_name} ORDER by vec_distance_{search_param["metric_type"]}(v, %s) LIMIT %d"
-        self.select_sql_with_filter = f"SELECT id FROM {self.db_name}.{self.table_name} WHERE id >= %d ORDER by vec_distance_{search_param["metric_type"]}(v, %s) LIMIT %d"
+        self.insert_sql = f"INSERT INTO {self.db_name}.{self.table_name} (id, v) VALUES (%s, %s)"  # noqa: S608
+        self.select_sql = (
+            f"SELECT id FROM {self.db_name}.{self.table_name}"  # noqa: S608
+            f"ORDER by vec_distance_{search_param['metric_type']}(v, %s) LIMIT %d"
+        )
+        self.select_sql_with_filter = (
+            f"SELECT id FROM {self.db_name}.{self.table_name} WHERE id >= %d "  # noqa: S608
+            f"ORDER by vec_distance_{search_param['metric_type']}(v, %s) LIMIT %d"
+        )
 
         try:
             yield
@@ -126,7 +127,6 @@ class MariaDB(VectorDB):
             self.conn.close()
             self.cursor = None
             self.conn = None
-    
 
     def ready_to_load(self) -> bool:
         pass
@@ -139,33 +139,31 @@ class MariaDB(VectorDB):
 
         try:
             index_options = f"DISTANCE={index_param['metric_type']}"
-            if index_param["index_type"] == "HNSW" and index_param["M"] != None:
+            if index_param["index_type"] == "HNSW" and index_param["M"] is not None:
                 index_options += f" M={index_param['M']}"
 
-            self.cursor.execute(f"""
+            self.cursor.execute(
+                f"""
               ALTER TABLE {self.db_name}.{self.table_name}
               ADD VECTOR KEY v(v) {index_options}
-            """)
+            """
+            )
             self.cursor.execute("COMMIT")
 
         except Exception as e:
-            log.warning(
-                f"Failed to create index: {self.table_name} error: {e}"
-            )
+            log.warning(f"Failed to create index: {self.table_name} error: {e}")
             raise e from None
 
-        pass
-
     @staticmethod
-    def vector_to_hex(v):
-        return np.array(v, 'float32').tobytes()
+    def vector_to_hex(v):  # noqa: ANN001
+        return np.array(v, "float32").tobytes()
 
     def insert_embeddings(
         self,
         embeddings: list[list[float]],
         metadata: list[int],
-        **kwargs: Any,
-    ) -> Tuple[int, Optional[Exception]]:
+        **kwargs,
+    ) -> tuple[int, Exception]:
         """Insert embeddings into the database.
         Should call self.init() first.
         """
@@ -178,7 +176,7 @@ class MariaDB(VectorDB):
 
             batch_data = []
             for i, row in enumerate(metadata_arr):
-                batch_data.append((int(row), self.vector_to_hex(embeddings_arr[i])));
+                batch_data.append((int(row), self.vector_to_hex(embeddings_arr[i])))
 
             self.cursor.executemany(self.insert_sql, batch_data)
             self.cursor.execute("COMMIT")
@@ -186,11 +184,8 @@ class MariaDB(VectorDB):
 
             return len(metadata), None
         except Exception as e:
-            log.warning(
-                f"Failed to insert data into Vector table ({self.table_name}), error: {e}"
-            )
+            log.warning(f"Failed to insert data into Vector table ({self.table_name}), error: {e}")
             return 0, e
-    
 
     def search_embedding(
         self,
@@ -198,17 +193,16 @@ class MariaDB(VectorDB):
         k: int = 100,
         filters: dict | None = None,
         timeout: int | None = None,
-        **kwargs: Any,
-    ) -> (list[int]):
+        **kwargs,
+    ) -> list[int]:
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
 
-        search_param = self.case_config.search_param()
+        search_param = self.case_config.search_param()  # noqa: F841
 
         if filters:
-            self.cursor.execute(self.select_sql_with_filter, (filters.get('id'), self.vector_to_hex(query), k))
+            self.cursor.execute(self.select_sql_with_filter, (filters.get("id"), self.vector_to_hex(query), k))
         else:
             self.cursor.execute(self.select_sql, (self.vector_to_hex(query), k))
 
-        return [id for id, in self.cursor.fetchall()]
-
+        return [id for (id,) in self.cursor.fetchall()]  # noqa: A001

vectordb_bench/backend/clients/milvus/cli.py

@@ -194,25 +194,6 @@ def MilvusGPUIVFFlat(**parameters: Unpack[MilvusGPUIVFTypedDict]):
         **parameters,
     )
 
-@cli.command()
-@click_parameter_decorators_from_typed_dict(MilvusGPUBruteForceTypedDict)
-def MilvusGPUBruteForce(**parameters: Unpack[MilvusGPUBruteForceTypedDict]):
-    from .config import GPUBruteForceConfig, MilvusConfig
-
-    run(
-        db=DBTYPE,
-        db_config=MilvusConfig(
-            db_label=parameters["db_label"],
-            uri=SecretStr(parameters["uri"]),
-            user=parameters["user_name"],
-            password=SecretStr(parameters["password"]),
-        ),
-        db_case_config=GPUBruteForceConfig(
-            metric_type=parameters["metric_type"],
-            limit=parameters["limit"],  # top-k for search
-        ),
-        **parameters,
-    )
 
 class MilvusGPUBruteForceTypedDict(CommonTypedDict, MilvusTypedDict):
     metric_type: Annotated[
@@ -224,6 +205,7 @@ class MilvusGPUBruteForceTypedDict(CommonTypedDict, MilvusTypedDict):
         click.option("--limit", type=int, required=True, help="Top-k limit for search"),
     ]
 
+
 @cli.command()
 @click_parameter_decorators_from_typed_dict(MilvusGPUBruteForceTypedDict)
 def MilvusGPUBruteForce(**parameters: Unpack[MilvusGPUBruteForceTypedDict]):

vectordb_bench/backend/clients/milvus/config.py

@@ -215,7 +215,6 @@ class GPUBruteForceConfig(MilvusIndexConfig, DBCaseConfig):
         }
 
 
-
 class GPUIVFPQConfig(MilvusIndexConfig, DBCaseConfig):
     nlist: int = 1024
     m: int = 0

vectordb_bench/backend/clients/milvus/milvus.py

@@ -155,7 +155,7 @@ class Milvus(VectorDB):
         embeddings: Iterable[list[float]],
         metadata: list[int],
         **kwargs,
-    ) -> (int, Exception):
+    ) -> tuple[int, Exception]:
         """Insert embeddings into Milvus. should call self.init() first"""
         # use the first insert_embeddings to init collection
         assert self.col is not None

vectordb_bench/backend/clients/pgvector/cli.py

@@ -18,8 +18,7 @@ from ....cli.cli import (
 )
 
 
-# ruff: noqa
-def set_default_quantized_fetch_limit(ctx: any, param: any, value: any):
+def set_default_quantized_fetch_limit(ctx: any, param: any, value: any):  # noqa: ARG001
     if ctx.params.get("reranking") and value is None:
         # ef_search is the default value for quantized_fetch_limit as it's bound by ef_search.
         # 100 is default value for quantized_fetch_limit for IVFFlat.

vectordb_bench/backend/clients/pinecone/pinecone.py

@@ -67,7 +67,7 @@ class Pinecone(VectorDB):
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs,
-    ) -> (int, Exception):
+    ) -> tuple[int, Exception]:
         assert len(embeddings) == len(metadata)
         insert_count = 0
         try:

vectordb_bench/backend/clients/qdrant_cloud/config.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel, SecretStr, validator
+from pydantic import BaseModel, SecretStr
 
 from ..api import DBCaseConfig, DBConfig, MetricType
 
@@ -20,14 +20,6 @@ class QdrantConfig(DBConfig):
             "url": self.url.get_secret_value(),
         }
 
-    @validator("*")
-    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:
-            raise ValueError("Empty string!")
-        return v
-
 
 class QdrantIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType | None = None

vectordb_bench/backend/clients/qdrant_cloud/qdrant_cloud.py

@@ -111,7 +111,7 @@ class QdrantCloud(VectorDB):
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs,
-    ) -> (int, Exception):
+    ) -> tuple[int, Exception]:
         """Insert embeddings into Milvus. should call self.init() first"""
         assert self.qdrant_client is not None
         try:

vectordb_bench/backend/clients/tidb/config.py

@@ -1,5 +1,6 @@
-from pydantic import SecretStr, BaseModel, validator
-from ..api import DBConfig, DBCaseConfig, MetricType
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, MetricType
 
 
 class TiDBConfig(DBConfig):
@@ -10,10 +11,6 @@ class TiDBConfig(DBConfig):
     db_name: str = "test"
     ssl: bool = False
 
-    @validator("*")
-    def not_empty_field(cls, v: any, field: any):
-        return v
-
     def to_dict(self) -> dict:
         pwd_str = self.password.get_secret_value()
         return {
@@ -33,10 +30,10 @@ class TiDBIndexConfig(BaseModel, DBCaseConfig):
     def get_metric_fn(self) -> str:
         if self.metric_type == MetricType.L2:
             return "vec_l2_distance"
-        elif self.metric_type == MetricType.COSINE:
+        if self.metric_type == MetricType.COSINE:
             return "vec_cosine_distance"
-        else:
-            raise ValueError(f"Unsupported metric type: {self.metric_type}")
+        msg = f"Unsupported metric type: {self.metric_type}"
+        raise ValueError(msg)
 
     def index_param(self) -> dict:
         return {

vectordb_bench/backend/clients/tidb/tidb.py

@@ -3,7 +3,7 @@ import io
 import logging
 import time
 from contextlib import contextmanager
-from typing import Any, Optional, Tuple
+from typing import Any
 
 import pymysql
 
@@ -62,7 +62,7 @@ class TiDB(VectorDB):
                 conn.commit()
         except Exception as e:
             log.warning("Failed to drop table: %s error: %s", self.table_name, e)
-            raise e
+            raise
 
     def _create_table(self):
         try:
@@ -80,7 +80,7 @@ class TiDB(VectorDB):
                 conn.commit()
         except Exception as e:
             log.warning("Failed to create table: %s error: %s", self.table_name, e)
-            raise e
+            raise
 
     def ready_to_load(self) -> bool:
         pass
@@ -122,25 +122,25 @@ class TiDB(VectorDB):
                     f"""
                     SELECT PROGRESS FROM information_schema.tiflash_replica
                     WHERE TABLE_SCHEMA = "{database}" AND TABLE_NAME = "{self.table_name}"
-                    """
+                    """  # noqa: S608
                 )
                 result = cursor.fetchone()
                 return result[0]
         except Exception as e:
             log.warning("Failed to check TiFlash replica progress: %s", e)
-            raise e
+            raise
 
     def _optimize_wait_tiflash_catch_up(self):
         try:
             with self._get_connection() as (conn, cursor):
                 cursor.execute('SET @@TIDB_ISOLATION_READ_ENGINES="tidb,tiflash"')
                 conn.commit()
-                cursor.execute(f"SELECT COUNT(*) FROM {self.table_name}")
+                cursor.execute(f"SELECT COUNT(*) FROM {self.table_name}")  # noqa: S608
                 result = cursor.fetchone()
                 return result[0]
         except Exception as e:
             log.warning("Failed to wait TiFlash to catch up: %s", e)
-            raise e
+            raise
 
     def _optimize_compact_tiflash(self):
         try:
@@ -149,7 +149,7 @@ class TiDB(VectorDB):
                 conn.commit()
         except Exception as e:
             log.warning("Failed to compact table: %s", e)
-            raise e
+            raise
 
     def _optimize_get_tiflash_index_pending_rows(self):
         try:
@@ -160,13 +160,13 @@ class TiDB(VectorDB):
                     SELECT SUM(ROWS_STABLE_NOT_INDEXED)
                     FROM information_schema.tiflash_indexes
                     WHERE TIDB_DATABASE = "{database}" AND TIDB_TABLE = "{self.table_name}"
-                    """
+                    """  # noqa: S608
                 )
                 result = cursor.fetchone()
                 return result[0]
         except Exception as e:
             log.warning("Failed to read TiFlash index pending rows: %s", e)
-            raise e
+            raise
 
     def _insert_embeddings_serial(
         self,
@@ -178,29 +178,28 @@ class TiDB(VectorDB):
         try:
             with self._get_connection() as (conn, cursor):
                 buf = io.StringIO()
-                buf.write(f"INSERT INTO {self.table_name} (id, embedding) VALUES ")
+                buf.write(f"INSERT INTO {self.table_name} (id, embedding) VALUES ")  # noqa: S608
                 for i in range(offset, offset + size):
                     if i > offset:
                         buf.write(",")
-                    buf.write(f'({metadata[i]}, "{str(embeddings[i])}")')
+                    buf.write(f'({metadata[i]}, "{embeddings[i]!s}")')
                 cursor.execute(buf.getvalue())
                 conn.commit()
         except Exception as e:
             log.warning("Failed to insert data into table: %s", e)
-            raise e
+            raise
 
     def insert_embeddings(
         self,
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs: Any,
-    ) -> Tuple[int, Optional[Exception]]:
+    ) -> tuple[int, Exception]:
         workers = 10
         # Avoid exceeding MAX_ALLOWED_PACKET (default=64MB)
         max_batch_size = 64 * 1024 * 1024 // 24 // self.dim
         batch_size = len(embeddings) // workers
-        if batch_size > max_batch_size:
-            batch_size = max_batch_size
+        batch_size = min(batch_size, max_batch_size)
         with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
             futures = []
             for i in range(0, len(embeddings), batch_size):
@@ -227,8 +226,8 @@ class TiDB(VectorDB):
         self.cursor.execute(
             f"""
             SELECT id FROM {self.table_name}
-            ORDER BY {self.search_fn}(embedding, "{str(query)}") LIMIT {k};
-            """
+            ORDER BY {self.search_fn}(embedding, "{query!s}") LIMIT {k};
+            """  # noqa: S608
         )
         result = self.cursor.fetchall()
         return [int(i[0]) for i in result]

vectordb_bench/backend/clients/vespa/cli.py

@@ -0,0 +1,47 @@
+from typing import Annotated, Unpack
+
+import click
+from pydantic import SecretStr
+
+from vectordb_bench.backend.clients import DB
+from vectordb_bench.cli.cli import (
+    CommonTypedDict,
+    HNSWFlavor1,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+
+
+class VespaTypedDict(CommonTypedDict, HNSWFlavor1):
+    uri: Annotated[
+        str,
+        click.option("--uri", "-u", type=str, help="uri connection string", default="http://127.0.0.1"),
+    ]
+    port: Annotated[
+        int,
+        click.option("--port", "-p", type=int, help="connection port", default=8080),
+    ]
+    quantization: Annotated[
+        str, click.option("--quantization", type=click.Choice(["none", "binary"], case_sensitive=False), default="none")
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(VespaTypedDict)
+def Vespa(**params: Unpack[VespaTypedDict]):
+    from .config import VespaConfig, VespaHNSWConfig
+
+    case_params = {
+        "quantization_type": params["quantization"],
+        "M": params["m"],
+        "efConstruction": params["ef_construction"],
+        "ef": params["ef_search"],
+    }
+
+    run(
+        db=DB.Vespa,
+        db_config=VespaConfig(url=SecretStr(params["uri"]), port=params["port"]),
+        db_case_config=VespaHNSWConfig(**{k: v for k, v in case_params.items() if v}),
+        **params,
+    )

vectordb_bench/backend/clients/vespa/config.py

@@ -0,0 +1,51 @@
+from typing import Literal, TypeAlias
+
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, MetricType
+
+VespaMetric: TypeAlias = Literal["euclidean", "angular", "dotproduct", "prenormalized-angular", "hamming", "geodegrees"]
+
+VespaQuantizationType: TypeAlias = Literal["none", "binary"]
+
+
+class VespaConfig(DBConfig):
+    url: SecretStr = "http://127.0.0.1"
+    port: int = 8080
+
+    def to_dict(self):
+        return {
+            "url": self.url.get_secret_value(),
+            "port": self.port,
+        }
+
+
+class VespaHNSWConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType = MetricType.COSINE
+    quantization_type: VespaQuantizationType = "none"
+    M: int = 16
+    efConstruction: int = 200
+    ef: int = 100
+
+    def index_param(self) -> dict:
+        return {
+            "distance_metric": self.parse_metric(self.metric_type),
+            "max_links_per_node": self.M,
+            "neighbors_to_explore_at_insert": self.efConstruction,
+        }
+
+    def search_param(self) -> dict:
+        return {}
+
+    def parse_metric(self, metric_type: MetricType) -> VespaMetric:
+        match metric_type:
+            case MetricType.COSINE:
+                return "angular"
+            case MetricType.L2:
+                return "euclidean"
+            case MetricType.DP | MetricType.IP:
+                return "dotproduct"
+            case MetricType.HAMMING:
+                return "hamming"
+            case _:
+                raise NotImplementedError

vectordb_bench/backend/clients/vespa/util.py

@@ -0,0 +1,15 @@
+"""Utility functions for supporting binary quantization
+
+From https://docs.vespa.ai/en/binarizing-vectors.html#appendix-conversion-to-int8
+"""
+
+import numpy as np
+
+
+def binarize_tensor(tensor: list[float]) -> list[int]:
+    """
+    Binarize a floating-point list by thresholding at zero
+    and packing the bits into bytes.
+    """
+    tensor = np.array(tensor)
+    return np.packbits(np.where(tensor > 0, 1, 0), axis=0).astype(np.int8).tolist()