vectordb-bench 0.0.22__tar.gz → 0.0.23__tar.gz

{vectordb_bench-0.0.22 → vectordb_bench-0.0.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: vectordb-bench
-Version: 0.0.22
+Version: 0.0.23
 Summary: VectorDBBench is not just an offering of benchmark results for mainstream vector databases and cloud services, it's your go-to tool for the ultimate performance and cost-effectiveness comparison. Designed with ease-of-use in mind, VectorDBBench is devised to help users, even non-professionals, reproduce results or test new systems, making the hunt for the optimal choice amongst a plethora of cloud services and open-source vector databases a breeze.
 Author-email: XuanYang-cn <xuan.yang@zilliz.com>
 Project-URL: repository, https://github.com/zilliztech/VectorDBBench
@@ -48,6 +48,8 @@ Requires-Dist: opensearch-py; extra == "all"
 Requires-Dist: memorydb; extra == "all"
 Requires-Dist: alibabacloud_ha3engine_vector; extra == "all"
 Requires-Dist: alibabacloud_searchengine20211025; extra == "all"
+Requires-Dist: mariadb; extra == "all"
+Requires-Dist: PyMySQL; extra == "all"
 Provides-Extra: qdrant
 Requires-Dist: qdrant-client; extra == "qdrant"
 Provides-Extra: pinecone
@@ -75,6 +77,10 @@ Requires-Dist: alibabacloud_ha3engine_vector; extra == "aliyun-opensearch"
 Requires-Dist: alibabacloud_searchengine20211025; extra == "aliyun-opensearch"
 Provides-Extra: mongodb
 Requires-Dist: pymongo; extra == "mongodb"
+Provides-Extra: mariadb
+Requires-Dist: mariadb; extra == "mariadb"
+Provides-Extra: tidb
+Requires-Dist: PyMySQL; extra == "tidb"
 
 # VectorDBBench: A Benchmark Tool for VectorDB
 
@@ -133,6 +139,7 @@ All the database client supported
 | awsopensearch            | `pip install vectordb-bench[opensearch]` |
 | aliyun_opensearch        | `pip install vectordb-bench[aliyun_opensearch]` |
 | mongodb                  | `pip install vectordb-bench[mongodb]`       |
+| tidb                     | `pip install vectordb-bench[tidb]`          |
 
 ### Run
 
@@ -209,7 +216,11 @@ Options:
   --ef-construction INTEGER       hnsw ef-construction
   --ef-search INTEGER             hnsw ef-search
   --quantization-type [none|bit|halfvec]
-                                  quantization type for vectors
+                                  quantization type for vectors (in index)
+  --table-quantization-type [none|bit|halfvec]
+                                  quantization type for vectors (in table). If
+                                  equal to bit, the parameter
+                                  quantization_type will be set to bit too.
   --custom-case-name TEXT         Custom case name i.e. PerformanceCase1536D50K
   --custom-case-description TEXT  Custom name description
   --custom-case-load-timeout INTEGER

{vectordb_bench-0.0.22 → vectordb_bench-0.0.23}/README.md

@@ -55,6 +55,7 @@ All the database client supported
 | awsopensearch            | `pip install vectordb-bench[opensearch]` |
 | aliyun_opensearch        | `pip install vectordb-bench[aliyun_opensearch]` |
 | mongodb                  | `pip install vectordb-bench[mongodb]`       |
+| tidb                     | `pip install vectordb-bench[tidb]`          |
 
 ### Run
 
@@ -131,7 +132,11 @@ Options:
   --ef-construction INTEGER       hnsw ef-construction
   --ef-search INTEGER             hnsw ef-search
   --quantization-type [none|bit|halfvec]
-                                  quantization type for vectors
+                                  quantization type for vectors (in index)
+  --table-quantization-type [none|bit|halfvec]
+                                  quantization type for vectors (in table). If
+                                  equal to bit, the parameter
+                                  quantization_type will be set to bit too.
   --custom-case-name TEXT         Custom case name i.e. PerformanceCase1536D50K
   --custom-case-description TEXT  Custom name description
   --custom-case-load-timeout INTEGER

{vectordb_bench-0.0.22 → vectordb_bench-0.0.23}/pyproject.toml

@@ -68,6 +68,8 @@ all = [
     "memorydb",
     "alibabacloud_ha3engine_vector",
     "alibabacloud_searchengine20211025",
+    "mariadb",
+    "PyMySQL",
 ]
 
 qdrant          = [ "qdrant-client" ]
@@ -86,6 +88,8 @@ chromadb        = [ "chromadb" ]
 opensearch      = [ "opensearch-py" ]
 aliyun_opensearch = [ "alibabacloud_ha3engine_vector", "alibabacloud_searchengine20211025"]
 mongodb         = [ "pymongo" ]
+mariadb         = [ "mariadb" ]
+tidb            = [ "PyMySQL" ]
 
 [project.urls]
 "repository" = "https://github.com/zilliztech/VectorDBBench"

{vectordb_bench-0.0.22 → vectordb_bench-0.0.23}/vectordb_bench/backend/clients/__init__.py

@@ -38,9 +38,11 @@ class DB(Enum):
     Chroma = "Chroma"
     AWSOpenSearch = "OpenSearch"
     AliyunElasticsearch = "AliyunElasticsearch"
+    MariaDB = "MariaDB"
     Test = "test"
     AliyunOpenSearch = "AliyunOpenSearch"
     MongoDB = "MongoDB"
+    TiDB = "TiDB"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901
@@ -135,6 +137,16 @@ class DB(Enum):
 
             return MongoDB
 
+        if self == DB.MariaDB:
+            from .mariadb.mariadb import MariaDB
+
+            return MariaDB
+
+        if self == DB.TiDB:
+            from .tidb.tidb import TiDB
+
+            return TiDB
+
         if self == DB.Test:
             from .test.test import Test
 
@@ -236,6 +248,16 @@ class DB(Enum):
 
             return MongoDBConfig
 
+        if self == DB.MariaDB:
+            from .mariadb.config import MariaDBConfig
+
+            return MariaDBConfig
+
+        if self == DB.TiDB:
+            from .tidb.config import TiDBConfig
+
+            return TiDBConfig
+
         if self == DB.Test:
             from .test.config import TestConfig
 
@@ -318,6 +340,16 @@ class DB(Enum):
 
             return MongoDBIndexConfig
 
+        if self == DB.MariaDB:
+            from .mariadb.config import _mariadb_case_config
+
+            return _mariadb_case_config.get(index_type)
+
+        if self == DB.TiDB:
+            from .tidb.config import TiDBIndexConfig
+
+            return TiDBIndexConfig
+
         # DB.Pinecone, DB.Chroma, DB.Redis
         return EmptyDBCaseConfig
 

{vectordb_bench-0.0.22 → vectordb_bench-0.0.23}/vectordb_bench/backend/clients/api.py

@@ -25,6 +25,7 @@ class IndexType(str, Enum):
     ES_HNSW = "hnsw"
     ES_IVFFlat = "ivfflat"
     GPU_IVF_FLAT = "GPU_IVF_FLAT"
+    GPU_BRUTE_FORCE = "GPU_BRUTE_FORCE"
     GPU_IVF_PQ = "GPU_IVF_PQ"
     GPU_CAGRA = "GPU_CAGRA"
     SCANN = "scann"

vectordb_bench-0.0.23/vectordb_bench/backend/clients/mariadb/cli.py

@@ -0,0 +1,107 @@
+from typing import Annotated, Optional, Unpack
+
+import click
+import os
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    HNSWFlavor1,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from vectordb_bench.backend.clients import DB
+
+
+class MariaDBTypedDict(CommonTypedDict):
+    user_name: Annotated[
+        str, click.option("--username",
+                          type=str,
+                          help="Username",
+                          required=True,
+                          ),
+    ]
+    password: Annotated[
+        str, click.option("--password",
+                     type=str,
+                     help="Password",
+                     required=True,
+                     ),
+    ]
+
+    host: Annotated[
+        str, click.option("--host",
+                          type=str,
+                          help="Db host",
+                          default="127.0.0.1",
+                          ),
+    ]
+
+    port: Annotated[
+        int, click.option("--port",
+                          type=int,
+                          default=3306,
+                          help="Db Port",
+                          ),
+    ]
+
+    storage_engine: Annotated[
+        int, click.option("--storage-engine",
+                          type=click.Choice(["InnoDB", "MyISAM"]),
+                          help="DB storage engine",
+                          required=True,
+                          ),
+    ]
+
+class MariaDBHNSWTypedDict(MariaDBTypedDict):
+    ...
+    m: Annotated[
+        Optional[int], click.option("--m",
+                                    type=int,
+                                    help="M parameter in MHNSW vector indexing",
+                                    required=False,
+                                    ),
+    ]
+
+    ef_search: Annotated[
+        Optional[int], click.option("--ef-search",
+                                    type=int,
+                                    help="MariaDB system variable mhnsw_min_limit",
+                                    required=False,
+                                    ),
+    ]
+
+    max_cache_size: Annotated[
+        Optional[int], click.option("--max-cache-size",
+                                    type=int,
+                                    help="MariaDB system variable mhnsw_max_cache_size",
+                                    required=False,
+                                    ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MariaDBHNSWTypedDict)
+def MariaDBHNSW(
+    **parameters: Unpack[MariaDBHNSWTypedDict],
+):
+    from .config import MariaDBConfig, MariaDBHNSWConfig
+
+    run(
+        db=DB.MariaDB,
+        db_config=MariaDBConfig(
+            db_label=parameters["db_label"],
+            user_name=parameters["username"],
+            password=SecretStr(parameters["password"]),
+            host=parameters["host"],
+            port=parameters["port"],
+        ),
+        db_case_config=MariaDBHNSWConfig(
+            M=parameters["m"],
+            ef_search=parameters["ef_search"],
+            storage_engine=parameters["storage_engine"],
+            max_cache_size=parameters["max_cache_size"],
+        ),
+        **parameters,
+    )

vectordb_bench-0.0.23/vectordb_bench/backend/clients/mariadb/config.py

@@ -0,0 +1,71 @@
+from pydantic import SecretStr, BaseModel
+from typing import TypedDict
+from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
+
+class MariaDBConfigDict(TypedDict):
+    """These keys will be directly used as kwargs in mariadb connection string,
+        so the names must match exactly mariadb API"""
+
+    user: str
+    password: str
+    host: str
+    port: int
+
+
+class MariaDBConfig(DBConfig):
+    user_name: str = "root"
+    password: SecretStr
+    host: str = "127.0.0.1"
+    port: int = 3306
+
+    def to_dict(self) -> MariaDBConfigDict:
+        pwd_str = self.password.get_secret_value()
+        return {
+            "host": self.host,
+            "port": self.port,
+            "user": self.user_name,
+            "password": pwd_str,
+        }
+
+
+class MariaDBIndexConfig(BaseModel):
+    """Base config for MariaDB"""
+
+    metric_type: MetricType | None = None
+
+    def parse_metric(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "euclidean"
+        elif self.metric_type == MetricType.COSINE:
+            return "cosine"
+        else:
+            raise ValueError(f"Metric type {self.metric_type} is not supported!")
+
+class MariaDBHNSWConfig(MariaDBIndexConfig, DBCaseConfig):
+    M: int | None
+    ef_search: int | None
+    index: IndexType = IndexType.HNSW
+    storage_engine: str = "InnoDB"
+    max_cache_size: int | None
+
+    def index_param(self) -> dict:
+        return {
+            "storage_engine": self.storage_engine,
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "M": self.M,
+            "max_cache_size": self.max_cache_size,
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "ef_search": self.ef_search,
+        }
+
+
+_mariadb_case_config = {
+        IndexType.HNSW: MariaDBHNSWConfig,
+}
+
+

vectordb_bench-0.0.23/vectordb_bench/backend/clients/mariadb/mariadb.py

@@ -0,0 +1,214 @@
+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
+
+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.name = "MariaDB"
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.db_name = "vectordbbench"
+        self.table_name = collection_name
+        self.dim = dim
+
+        # construct basic units
+        self.conn, self.cursor = self._create_connection(**self.db_config)
+        
+        if drop_old:
+            self._drop_db()
+            self._create_db_table(dim)
+
+        self.cursor.close()
+        self.conn.close()
+        self.cursor = None
+        self.conn = None
+
+
+    @staticmethod
+    def _create_connection(**kwargs) -> Tuple[mariadb.Connection, mariadb.Cursor]:
+        conn = mariadb.connect(**kwargs)
+        cursor = conn.cursor()
+
+        assert conn is not None, "Connection is not initialized"
+        assert cursor is not None, "Cursor is not initialized"
+
+        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"
+        log.info(f"{self.name} client drop db : {self.db_name}")
+
+        # flush tables before dropping database to avoid some locking issue
+        self.cursor.execute("FLUSH TABLES")
+        self.cursor.execute(f"DROP DATABASE IF EXISTS {self.db_name}")
+        self.cursor.execute("COMMIT")
+        self.cursor.execute("FLUSH TABLES")
+
+    def _create_db_table(self, dim: int):
+        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 database : {self.db_name}")
+            self.cursor.execute(f"CREATE DATABASE {self.db_name}")
+
+            log.info(f"{self.name} client create table : {self.table_name}")
+            self.cursor.execute(f"USE {self.db_name}")
+
+            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}"
+            )
+            raise e from None
+
+
+    @contextmanager
+    def init(self) -> None:
+        """ create and destory connections to database.
+
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+        """
+        self.conn, self.cursor = self._create_connection(**self.db_config)
+
+        index_param = self.case_config.index_param()
+        search_param = self.case_config.search_param()
+
+        # maximize allowed package size
+        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"]}")
+            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"
+
+        try:
+            yield
+        finally:
+            self.cursor.close()
+            self.conn.close()
+            self.cursor = None
+            self.conn = None
+    
+
+    def ready_to_load(self) -> bool:
+        pass
+
+    def optimize(self) -> None:
+        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:
+            index_options = f"DISTANCE={index_param['metric_type']}"
+            if index_param["index_type"] == "HNSW" and index_param["M"] != None:
+                index_options += f" M={index_param['M']}"
+
+            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}"
+            )
+            raise e from None
+
+        pass
+
+    @staticmethod
+    def vector_to_hex(v):
+        return np.array(v, 'float32').tobytes()
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> Tuple[int, Optional[Exception]]:
+        """Insert embeddings into the database.
+        Should call self.init() first.
+        """
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        try:
+            metadata_arr = np.array(metadata)
+            embeddings_arr = np.array(embeddings)
+
+            batch_data = []
+            for i, row in enumerate(metadata_arr):
+                batch_data.append((int(row), self.vector_to_hex(embeddings_arr[i])));
+
+            self.cursor.executemany(self.insert_sql, batch_data)
+            self.cursor.execute("COMMIT")
+            self.cursor.execute("FLUSH TABLES")
+
+            return len(metadata), None
+        except Exception as e:
+            log.warning(
+                f"Failed to insert data into Vector table ({self.table_name}), error: {e}"
+            )
+            return 0, e
+    
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        timeout: int | None = None,
+        **kwargs: Any,
+    ) -> (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()
+
+        if filters:
+            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()]
+

{vectordb_bench-0.0.22 → vectordb_bench-0.0.23}/vectordb_bench/backend/clients/milvus/cli.py

@@ -194,6 +194,56 @@ 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[
+        str,
+        click.option("--metric-type", type=str, required=True, help="Metric type for brute force search"),
+    ]
+    limit: Annotated[
+        int,
+        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]):
+    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 MilvusGPUIVFPQTypedDict(
     CommonTypedDict,

{vectordb_bench-0.0.22 → vectordb_bench-0.0.23}/vectordb_bench/backend/clients/milvus/config.py

@@ -40,6 +40,7 @@ class MilvusIndexConfig(BaseModel):
             IndexType.GPU_CAGRA,
             IndexType.GPU_IVF_FLAT,
             IndexType.GPU_IVF_PQ,
+            IndexType.GPU_BRUTE_FORCE,
         ]
 
     def parse_metric(self) -> str:
@@ -184,6 +185,37 @@ class GPUIVFFlatConfig(MilvusIndexConfig, DBCaseConfig):
         }
 
 
+class GPUBruteForceConfig(MilvusIndexConfig, DBCaseConfig):
+    limit: int = 10  # Default top-k for search
+    metric_type: str  # Metric type (e.g., 'L2', 'IP', etc.)
+    index: IndexType = IndexType.GPU_BRUTE_FORCE  # Index type set to GPU_BRUTE_FORCE
+
+    def index_param(self) -> dict:
+        """
+        Returns the parameters for creating the GPU_BRUTE_FORCE index.
+        No additional parameters required for index building.
+        """
+        return {
+            "metric_type": self.parse_metric(),  # Metric type for distance calculation (L2, IP, etc.)
+            "index_type": self.index.value,  # GPU_BRUTE_FORCE index type
+            "params": {},  # No additional parameters for GPU_BRUTE_FORCE
+        }
+
+    def search_param(self) -> dict:
+        """
+        Returns the parameters for performing a search on the GPU_BRUTE_FORCE index.
+        Only metric_type and top-k (limit) are needed for search.
+        """
+        return {
+            "metric_type": self.parse_metric(),  # Metric type for search
+            "params": {
+                "nprobe": 1,  # For GPU_BRUTE_FORCE, set nprobe to 1 (brute force search)
+                "limit": self.limit,  # Top-k for search
+            },
+        }
+
+
+
 class GPUIVFPQConfig(MilvusIndexConfig, DBCaseConfig):
     nlist: int = 1024
     m: int = 0
@@ -261,4 +293,5 @@ _milvus_case_config = {
     IndexType.GPU_IVF_FLAT: GPUIVFFlatConfig,
     IndexType.GPU_IVF_PQ: GPUIVFPQConfig,
     IndexType.GPU_CAGRA: GPUCAGRAConfig,
+    IndexType.GPU_BRUTE_FORCE: GPUBruteForceConfig,
 }