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

vectordb_bench/backend/clients/tidb/cli.py

@@ -0,0 +1,98 @@
+from typing import Annotated, Unpack
+
+import click
+from pydantic import SecretStr
+
+from vectordb_bench.backend.clients import DB
+
+from ....cli.cli import CommonTypedDict, cli, click_parameter_decorators_from_typed_dict, run
+
+
+class TiDBTypedDict(CommonTypedDict):
+    user_name: Annotated[
+        str,
+        click.option(
+            "--username",
+            type=str,
+            help="Username",
+            default="root",
+            show_default=True,
+            required=True,
+        ),
+    ]
+    password: Annotated[
+        str,
+        click.option(
+            "--password",
+            type=str,
+            default="",
+            show_default=True,
+            help="Password",
+        ),
+    ]
+    host: Annotated[
+        str,
+        click.option(
+            "--host",
+            type=str,
+            default="127.0.0.1",
+            show_default=True,
+            required=True,
+            help="Db host",
+        ),
+    ]
+    port: Annotated[
+        int,
+        click.option(
+            "--port",
+            type=int,
+            default=4000,
+            show_default=True,
+            required=True,
+            help="Db Port",
+        ),
+    ]
+    db_name: Annotated[
+        str,
+        click.option(
+            "--db-name",
+            type=str,
+            default="test",
+            show_default=True,
+            required=True,
+            help="Db name",
+        ),
+    ]
+    ssl: Annotated[
+        bool,
+        click.option(
+            "--ssl/--no-ssl",
+            default=False,
+            show_default=True,
+            is_flag=True,
+            help="Enable or disable SSL, for TiDB Serverless SSL must be enabled",
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(TiDBTypedDict)
+def TiDB(
+    **parameters: Unpack[TiDBTypedDict],
+):
+    from .config import TiDBConfig, TiDBIndexConfig
+
+    run(
+        db=DB.TiDB,
+        db_config=TiDBConfig(
+            db_label=parameters["db_label"],
+            user_name=parameters["username"],
+            password=SecretStr(parameters["password"]),
+            host=parameters["host"],
+            port=parameters["port"],
+            db_name=parameters["db_name"],
+            ssl=parameters["ssl"],
+        ),
+        db_case_config=TiDBIndexConfig(),
+        **parameters,
+    )

vectordb_bench/backend/clients/tidb/config.py

@@ -0,0 +1,46 @@
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, MetricType
+
+
+class TiDBConfig(DBConfig):
+    user_name: str = "root"
+    password: SecretStr
+    host: str = "127.0.0.1"
+    port: int = 4000
+    db_name: str = "test"
+    ssl: bool = False
+
+    def to_dict(self) -> dict:
+        pwd_str = self.password.get_secret_value()
+        return {
+            "host": self.host,
+            "port": self.port,
+            "user": self.user_name,
+            "password": pwd_str,
+            "database": self.db_name,
+            "ssl_verify_cert": self.ssl,
+            "ssl_verify_identity": self.ssl,
+        }
+
+
+class TiDBIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+
+    def get_metric_fn(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "vec_l2_distance"
+        if self.metric_type == MetricType.COSINE:
+            return "vec_cosine_distance"
+        msg = f"Unsupported metric type: {self.metric_type}"
+        raise ValueError(msg)
+
+    def index_param(self) -> dict:
+        return {
+            "metric_fn": self.get_metric_fn(),
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_fn": self.get_metric_fn(),
+        }

vectordb_bench/backend/clients/tidb/tidb.py

@@ -0,0 +1,233 @@
+import concurrent.futures
+import io
+import logging
+import time
+from contextlib import contextmanager
+from typing import Any
+
+import pymysql
+
+from ..api import VectorDB
+from .config import TiDBIndexConfig
+
+log = logging.getLogger(__name__)
+
+
+class TiDB(VectorDB):
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: TiDBIndexConfig,
+        collection_name: str = "vector_bench_test",
+        drop_old: bool = False,
+        **kwargs,
+    ):
+        self.name = "TiDB"
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.table_name = collection_name
+        self.dim = dim
+        self.conn = None  # To be inited by init()
+        self.cursor = None  # To be inited by init()
+
+        self.search_fn = db_case_config.search_param()["metric_fn"]
+
+        if drop_old:
+            self._drop_table()
+            self._create_table()
+
+    @contextmanager
+    def init(self):
+        with self._get_connection() as (conn, cursor):
+            self.conn = conn
+            self.cursor = cursor
+            try:
+                yield
+            finally:
+                self.conn = None
+                self.cursor = None
+
+    @contextmanager
+    def _get_connection(self):
+        with pymysql.connect(**self.db_config) as conn:
+            conn.autocommit = False
+            with conn.cursor() as cursor:
+                yield conn, cursor
+
+    def _drop_table(self):
+        try:
+            with self._get_connection() as (conn, cursor):
+                cursor.execute(f"DROP TABLE IF EXISTS {self.table_name}")
+                conn.commit()
+        except Exception as e:
+            log.warning("Failed to drop table: %s error: %s", self.table_name, e)
+            raise
+
+    def _create_table(self):
+        try:
+            index_param = self.case_config.index_param()
+            with self._get_connection() as (conn, cursor):
+                cursor.execute(
+                    f"""
+                    CREATE TABLE {self.table_name} (
+                        id BIGINT PRIMARY KEY,
+                        embedding VECTOR({self.dim}) NOT NULL,
+                        VECTOR INDEX (({index_param["metric_fn"]}(embedding)))
+                    );
+                    """
+                )
+                conn.commit()
+        except Exception as e:
+            log.warning("Failed to create table: %s error: %s", self.table_name, e)
+            raise
+
+    def ready_to_load(self) -> bool:
+        pass
+
+    def optimize(self, data_size: int | None = None) -> None:
+        while True:
+            progress = self._optimize_check_tiflash_replica_progress()
+            if progress != 1:
+                log.info("Data replication not ready, progress: %d", progress)
+                time.sleep(2)
+            else:
+                break
+
+        log.info("Waiting TiFlash to catch up...")
+        self._optimize_wait_tiflash_catch_up()
+
+        log.info("Start compacting TiFlash replica...")
+        self._optimize_compact_tiflash()
+
+        log.info("Waiting index build to finish...")
+        log_reduce_seq = 0
+        while True:
+            pending_rows = self._optimize_get_tiflash_index_pending_rows()
+            if pending_rows > 0:
+                if log_reduce_seq % 15 == 0:
+                    log.info("Index not fully built, pending rows: %d", pending_rows)
+                log_reduce_seq += 1
+                time.sleep(2)
+            else:
+                break
+
+        log.info("Index build finished successfully.")
+
+    def _optimize_check_tiflash_replica_progress(self):
+        try:
+            database = self.db_config["database"]
+            with self._get_connection() as (_, cursor):
+                cursor.execute(
+                    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
+
+    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}")  # noqa: S608
+                result = cursor.fetchone()
+                return result[0]
+        except Exception as e:
+            log.warning("Failed to wait TiFlash to catch up: %s", e)
+            raise
+
+    def _optimize_compact_tiflash(self):
+        try:
+            with self._get_connection() as (conn, cursor):
+                cursor.execute(f"ALTER TABLE {self.table_name} COMPACT")
+                conn.commit()
+        except Exception as e:
+            log.warning("Failed to compact table: %s", e)
+            raise
+
+    def _optimize_get_tiflash_index_pending_rows(self):
+        try:
+            database = self.db_config["database"]
+            with self._get_connection() as (_, cursor):
+                cursor.execute(
+                    f"""
+                    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
+
+    def _insert_embeddings_serial(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        offset: int,
+        size: int,
+    ) -> Exception:
+        try:
+            with self._get_connection() as (conn, cursor):
+                buf = io.StringIO()
+                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]}, "{embeddings[i]!s}")')
+                cursor.execute(buf.getvalue())
+                conn.commit()
+        except Exception as e:
+            log.warning("Failed to insert data into table: %s", e)
+            raise
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> 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
+        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):
+                offset = i
+                size = min(batch_size, len(embeddings) - i)
+                future = executor.submit(self._insert_embeddings_serial, embeddings, metadata, offset, size)
+                futures.append(future)
+            done, pending = concurrent.futures.wait(futures, return_when=concurrent.futures.FIRST_EXCEPTION)
+            executor.shutdown(wait=False)
+            for future in done:
+                future.result()
+            for future in pending:
+                future.cancel()
+        return len(metadata), None
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        timeout: int | None = None,
+        **kwargs: Any,
+    ) -> list[int]:
+        self.cursor.execute(
+            f"""
+            SELECT id FROM {self.table_name}
+            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()