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

vectordb_bench/backend/clients/pgvector/pgvector.py

@@ -94,7 +94,7 @@ class PgVector(VectorDB):
         reranking = self.case_config.search_param()["reranking"]
         column_name = (
             sql.SQL("binary_quantize({0})").format(sql.Identifier("embedding"))
-            if index_param["quantization_type"] == "bit"
+            if index_param["quantization_type"] == "bit" and index_param["table_quantization_type"] != "bit"
             else sql.SQL("embedding")
         )
         search_vector = (
@@ -104,7 +104,8 @@ class PgVector(VectorDB):
         )
 
         # The following sections assume that the quantization_type value matches the quantization function name
-        if index_param["quantization_type"] is not None:
+        if index_param["quantization_type"] != index_param["table_quantization_type"]:
+            # Reranking makes sense only if table quantization is not "bit"
             if index_param["quantization_type"] == "bit" and reranking:
                 # Embeddings needs to be passed to binary_quantize function if quantization_type is bit
                 search_query = sql.Composed(
@@ -113,7 +114,7 @@ class PgVector(VectorDB):
                             """
                             SELECT i.id
                             FROM (
-                                SELECT id, embedding {reranking_metric_fun_op} %s::vector AS distance
+                                SELECT id, embedding {reranking_metric_fun_op} %s::{table_quantization_type} AS distance
                                 FROM public.{table_name} {where_clause}
                                 ORDER BY {column_name}::{quantization_type}({dim})
                             """,
@@ -123,6 +124,8 @@ class PgVector(VectorDB):
                             reranking_metric_fun_op=sql.SQL(
                                 self.case_config.search_param()["reranking_metric_fun_op"],
                             ),
+                            search_vector=search_vector,
+                            table_quantization_type=sql.SQL(index_param["table_quantization_type"]),
                             quantization_type=sql.SQL(index_param["quantization_type"]),
                             dim=sql.Literal(self.dim),
                             where_clause=sql.SQL("WHERE id >= %s") if filtered else sql.SQL(""),
@@ -130,7 +133,7 @@ class PgVector(VectorDB):
                         sql.SQL(self.case_config.search_param()["metric_fun_op"]),
                         sql.SQL(
                             """
-                                {search_vector}
+                                {search_vector}::{quantization_type}({dim})
                                 LIMIT {quantized_fetch_limit}
                             ) i
                             ORDER BY i.distance
@@ -138,6 +141,8 @@ class PgVector(VectorDB):
                             """,
                         ).format(
                             search_vector=search_vector,
+                            quantization_type=sql.SQL(index_param["quantization_type"]),
+                            dim=sql.Literal(self.dim),
                             quantized_fetch_limit=sql.Literal(
                                 self.case_config.search_param()["quantized_fetch_limit"],
                             ),
@@ -160,10 +165,12 @@ class PgVector(VectorDB):
                             where_clause=sql.SQL("WHERE id >= %s") if filtered else sql.SQL(""),
                         ),
                         sql.SQL(self.case_config.search_param()["metric_fun_op"]),
-                        sql.SQL(" {search_vector} LIMIT %s::int").format(
+                        sql.SQL(" {search_vector}::{quantization_type}({dim}) LIMIT %s::int").format(
                             search_vector=search_vector,
+                            quantization_type=sql.SQL(index_param["quantization_type"]),
+                            dim=sql.Literal(self.dim),
                         ),
-                    ],
+                    ]
                 )
         else:
             search_query = sql.Composed(
@@ -175,8 +182,12 @@ class PgVector(VectorDB):
                         where_clause=sql.SQL("WHERE id >= %s") if filtered else sql.SQL(""),
                     ),
                     sql.SQL(self.case_config.search_param()["metric_fun_op"]),
-                    sql.SQL(" %s::vector LIMIT %s::int"),
-                ],
+                    sql.SQL(" {search_vector}::{quantization_type}({dim}) LIMIT %s::int").format(
+                        search_vector=search_vector,
+                        quantization_type=sql.SQL(index_param["quantization_type"]),
+                        dim=sql.Literal(self.dim),
+                    ),
+                ]
             )
 
         return search_query
@@ -323,7 +334,7 @@ class PgVector(VectorDB):
                 )
         with_clause = sql.SQL("WITH ({});").format(sql.SQL(", ").join(options)) if any(options) else sql.Composed(())
 
-        if index_param["quantization_type"] is not None:
+        if index_param["quantization_type"] != index_param["table_quantization_type"]:
             index_create_sql = sql.SQL(
                 """
                 CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name}
@@ -365,14 +376,23 @@ class PgVector(VectorDB):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
 
+        index_param = self.case_config.index_param()
+
         try:
             log.info(f"{self.name} client create table : {self.table_name}")
 
             # create table
             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 {table_quantization_type}({dim}));
+                    """
+                ).format(
+                    table_name=sql.Identifier(self.table_name),
+                    table_quantization_type=sql.SQL(index_param["table_quantization_type"]),
+                    dim=dim,
+                )
             )
             self.cursor.execute(
                 sql.SQL(
@@ -393,18 +413,41 @@ class PgVector(VectorDB):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
 
+        index_param = self.case_config.index_param()
+
         try:
             metadata_arr = np.array(metadata)
             embeddings_arr = np.array(embeddings)
 
-            with self.cursor.copy(
-                sql.SQL("COPY public.{table_name} FROM STDIN (FORMAT BINARY)").format(
-                    table_name=sql.Identifier(self.table_name),
-                ),
-            ) as copy:
-                copy.set_types(["bigint", "vector"])
-                for i, row in enumerate(metadata_arr):
-                    copy.write_row((row, embeddings_arr[i]))
+            if index_param["table_quantization_type"] == "bit":
+                with self.cursor.copy(
+                    sql.SQL("COPY public.{table_name} FROM STDIN (FORMAT TEXT)").format(
+                        table_name=sql.Identifier(self.table_name)
+                    )
+                ) as copy:
+                    # Same logic as pgvector binary_quantize
+                    for i, row in enumerate(metadata_arr):
+                        embeddings_bit = ""
+                        for embedding in embeddings_arr[i]:
+                            if embedding > 0:
+                                embeddings_bit += "1"
+                            else:
+                                embeddings_bit += "0"
+                        copy.write_row((str(row), embeddings_bit))
+            else:
+                with self.cursor.copy(
+                    sql.SQL("COPY public.{table_name} FROM STDIN (FORMAT BINARY)").format(
+                        table_name=sql.Identifier(self.table_name)
+                    )
+                ) as copy:
+                    if index_param["table_quantization_type"] == "halfvec":
+                        copy.set_types(["bigint", "halfvec"])
+                        for i, row in enumerate(metadata_arr):
+                            copy.write_row((row, np.float16(embeddings_arr[i])))
+                    else:
+                        copy.set_types(["bigint", "vector"])
+                        for i, row in enumerate(metadata_arr):
+                            copy.write_row((row, embeddings_arr[i]))
             self.conn.commit()
 
             if kwargs.get("last_batch"):

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,49 @@
+from pydantic import SecretStr, BaseModel, validator
+from ..api import DBConfig, DBCaseConfig, 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
+
+    @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 {
+            "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"
+        elif self.metric_type == MetricType.COSINE:
+            return "vec_cosine_distance"
+        else:
+            raise ValueError(f"Unsupported metric type: {self.metric_type}")
+
+    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,234 @@
+import concurrent.futures
+import io
+import logging
+import time
+from contextlib import contextmanager
+from typing import Any, Optional, Tuple
+
+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 e
+
+    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 e
+
+    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}"
+                    """
+                )
+                result = cursor.fetchone()
+                return result[0]
+        except Exception as e:
+            log.warning("Failed to check TiFlash replica progress: %s", e)
+            raise e
+
+    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}")
+                result = cursor.fetchone()
+                return result[0]
+        except Exception as e:
+            log.warning("Failed to wait TiFlash to catch up: %s", e)
+            raise e
+
+    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 e
+
+    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}"
+                    """
+                )
+                result = cursor.fetchone()
+                return result[0]
+        except Exception as e:
+            log.warning("Failed to read TiFlash index pending rows: %s", e)
+            raise e
+
+    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 ")
+                for i in range(offset, offset + size):
+                    if i > offset:
+                        buf.write(",")
+                    buf.write(f'({metadata[i]}, "{str(embeddings[i])}")')
+                cursor.execute(buf.getvalue())
+                conn.commit()
+        except Exception as e:
+            log.warning("Failed to insert data into table: %s", e)
+            raise e
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> Tuple[int, Optional[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
+        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, "{str(query)}") LIMIT {k};
+            """
+        )
+        result = self.cursor.fetchall()
+        return [int(i[0]) for i in result]

vectordb_bench/cli/vectordbbench.py

@@ -1,5 +1,6 @@
 from ..backend.clients.alloydb.cli import AlloyDBScaNN
 from ..backend.clients.aws_opensearch.cli import AWSOpenSearch
+from ..backend.clients.mariadb.cli import MariaDBHNSW
 from ..backend.clients.memorydb.cli import MemoryDB
 from ..backend.clients.milvus.cli import MilvusAutoIndex
 from ..backend.clients.pgdiskann.cli import PgDiskAnn
@@ -10,6 +11,7 @@ from ..backend.clients.redis.cli import Redis
 from ..backend.clients.test.cli import Test
 from ..backend.clients.weaviate_cloud.cli import Weaviate
 from ..backend.clients.zilliz_cloud.cli import ZillizAutoIndex
+from ..backend.clients.tidb.cli import TiDB
 from .cli import cli
 
 cli.add_command(PgVectorHNSW)
@@ -25,6 +27,8 @@ cli.add_command(AWSOpenSearch)
 cli.add_command(PgVectorScaleDiskAnn)
 cli.add_command(PgDiskAnn)
 cli.add_command(AlloyDBScaNN)
+cli.add_command(MariaDBHNSW)
+cli.add_command(TiDB)
 
 
 if __name__ == "__main__":