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

vectordb_bench/backend/clients/__init__.py

@@ -38,9 +38,13 @@ class DB(Enum):
     Chroma = "Chroma"
     AWSOpenSearch = "OpenSearch"
     AliyunElasticsearch = "AliyunElasticsearch"
+    MariaDB = "MariaDB"
     Test = "test"
     AliyunOpenSearch = "AliyunOpenSearch"
     MongoDB = "MongoDB"
+    TiDB = "TiDB"
+    Clickhouse = "Clickhouse"
+    Vespa = "Vespa"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901
@@ -115,6 +119,11 @@ class DB(Enum):
 
             return AWSOpenSearch
 
+        if self == DB.Clickhouse:
+            from .clickhouse.clickhouse import Clickhouse
+
+            return Clickhouse
+
         if self == DB.AlloyDB:
             from .alloydb.alloydb import AlloyDB
 
@@ -135,11 +144,26 @@ 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
 
             return Test
 
+        if self == DB.Vespa:
+            from .vespa.vespa import Vespa
+
+            return Vespa
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -216,6 +240,11 @@ class DB(Enum):
 
             return AWSOpenSearchConfig
 
+        if self == DB.Clickhouse:
+            from .clickhouse.config import ClickhouseConfig
+
+            return ClickhouseConfig
+
         if self == DB.AlloyDB:
             from .alloydb.config import AlloyDBConfig
 
@@ -236,15 +265,30 @@ 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
 
             return TestConfig
 
+        if self == DB.Vespa:
+            from .vespa.config import VespaConfig
+
+            return VespaConfig
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
-    def case_config_cls(  # noqa: PLR0911
+    def case_config_cls(  # noqa: C901, PLR0911, PLR0912
         self,
         index_type: IndexType | None = None,
     ) -> type[DBCaseConfig]:
@@ -288,6 +332,11 @@ class DB(Enum):
 
             return AWSOpenSearchIndexConfig
 
+        if self == DB.Clickhouse:
+            from .clickhouse.config import ClickhouseHNSWConfig
+
+            return ClickhouseHNSWConfig
+
         if self == DB.PgVectorScale:
             from .pgvectorscale.config import _pgvectorscale_case_config
 
@@ -318,6 +367,21 @@ 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
+
+        if self == DB.Vespa:
+            from .vespa.config import VespaHNSWConfig
+
+            return VespaHNSWConfig
+
         # DB.Pinecone, DB.Chroma, DB.Redis
         return EmptyDBCaseConfig
 

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"
@@ -161,7 +162,7 @@ class VectorDB(ABC):
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs,
-    ) -> (int, Exception):
+    ) -> tuple[int, Exception]:
         """Insert the embeddings to the vector database. The default number of embeddings for
         each insert_embeddings is 5000.
 

vectordb_bench/backend/clients/chroma/chroma.py

@@ -65,7 +65,7 @@ class ChromaClient(VectorDB):
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs: Any,
-    ) -> (int, Exception):
+    ) -> tuple[int, Exception]:
         """Insert embeddings into the database.
 
         Args:
@@ -74,7 +74,7 @@ class ChromaClient(VectorDB):
             kwargs: other arguments
 
         Returns:
-            (int, Exception): number of embeddings inserted and exception if any
+            tuple[int, Exception]: number of embeddings inserted and exception if any
         """
         ids = [str(i) for i in metadata]
         metadata = [{"id": int(i)} for i in metadata]

vectordb_bench/backend/clients/clickhouse/cli.py

@@ -0,0 +1,66 @@
+from typing import Annotated, TypedDict, Unpack
+
+import click
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    HNSWFlavor2,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from .. import DB
+from .config import ClickhouseHNSWConfig
+
+
+class ClickhouseTypedDict(TypedDict):
+    password: Annotated[str, click.option("--password", type=str, help="DB password")]
+    host: Annotated[str, click.option("--host", type=str, help="DB host", required=True)]
+    port: Annotated[int, click.option("--port", type=int, default=8123, help="DB Port")]
+    user: Annotated[int, click.option("--user", type=str, default="clickhouse", help="DB user")]
+    ssl: Annotated[
+        bool,
+        click.option(
+            "--ssl/--no-ssl",
+            is_flag=True,
+            show_default=True,
+            default=True,
+            help="Enable or disable SSL for Clickhouse",
+        ),
+    ]
+    ssl_ca_certs: Annotated[
+        str,
+        click.option(
+            "--ssl-ca-certs",
+            show_default=True,
+            help="Path to certificate authority file to use for SSL",
+        ),
+    ]
+
+
+class ClickhouseHNSWTypedDict(CommonTypedDict, ClickhouseTypedDict, HNSWFlavor2): ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(ClickhouseHNSWTypedDict)
+def Clickhouse(**parameters: Unpack[ClickhouseHNSWTypedDict]):
+    from .config import ClickhouseConfig
+
+    run(
+        db=DB.Clickhouse,
+        db_config=ClickhouseConfig(
+            db_label=parameters["db_label"],
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            host=parameters["host"],
+            port=parameters["port"],
+            ssl=parameters["ssl"],
+            ssl_ca_certs=parameters["ssl_ca_certs"],
+        ),
+        db_case_config=ClickhouseHNSWConfig(
+            M=parameters["m"],
+            efConstruction=parameters["ef_construction"],
+            ef=parameters["ef_runtime"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/clickhouse/clickhouse.py

@@ -0,0 +1,156 @@
+"""Wrapper around the Clickhouse vector database over VectorDB"""
+
+import logging
+from contextlib import contextmanager
+from typing import Any
+
+import clickhouse_connect
+
+from ..api import DBCaseConfig, VectorDB
+
+log = logging.getLogger(__name__)
+
+
+class Clickhouse(VectorDB):
+    """Use SQLAlchemy instructions"""
+
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: DBCaseConfig,
+        collection_name: str = "CHVectorCollection",
+        drop_old: bool = False,
+        **kwargs,
+    ):
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.table_name = collection_name
+        self.dim = dim
+
+        self._index_name = "clickhouse_index"
+        self._primary_field = "id"
+        self._vector_field = "embedding"
+
+        # construct basic units
+        self.conn = clickhouse_connect.get_client(
+            host=self.db_config["host"],
+            port=self.db_config["port"],
+            username=self.db_config["user"],
+            password=self.db_config["password"],
+            database=self.db_config["dbname"],
+        )
+
+        if drop_old:
+            log.info(f"Clickhouse client drop table : {self.table_name}")
+            self._drop_table()
+            self._create_table(dim)
+
+        self.conn.close()
+        self.conn = None
+
+    @contextmanager
+    def init(self):
+        """
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+            >>>     self.search_embedding()
+        """
+
+        self.conn = clickhouse_connect.get_client(
+            host=self.db_config["host"],
+            port=self.db_config["port"],
+            username=self.db_config["user"],
+            password=self.db_config["password"],
+            database=self.db_config["dbname"],
+        )
+
+        try:
+            yield
+        finally:
+            self.conn.close()
+            self.conn = None
+
+    def _drop_table(self):
+        assert self.conn is not None, "Connection is not initialized"
+
+        self.conn.command(f'DROP TABLE IF EXISTS {self.db_config["dbname"]}.{self.table_name}')
+
+    def _create_table(self, dim: int):
+        assert self.conn is not None, "Connection is not initialized"
+
+        try:
+            # create table
+            self.conn.command(
+                f'CREATE TABLE IF NOT EXISTS {self.db_config["dbname"]}.{self.table_name} \
+                    (id UInt32, embedding Array(Float64)) ENGINE = MergeTree() ORDER BY id;'
+            )
+
+        except Exception as e:
+            log.warning(f"Failed to create Clickhouse table: {self.table_name} error: {e}")
+            raise e from None
+
+    def ready_to_load(self):
+        pass
+
+    def optimize(self, data_size: int | None = None):
+        pass
+
+    def ready_to_search(self):
+        pass
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> tuple[int, Exception]:
+        assert self.conn is not None, "Connection is not initialized"
+
+        try:
+            # do not iterate for bulk insert
+            items = [metadata, embeddings]
+
+            self.conn.insert(
+                table=self.table_name,
+                data=items,
+                column_names=["id", "embedding"],
+                column_type_names=["UInt32", "Array(Float64)"],
+                column_oriented=True,
+            )
+            return len(metadata), None
+        except Exception as e:
+            log.warning(f"Failed to insert data into Clickhouse 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,
+    ) -> list[int]:
+        assert self.conn is not None, "Connection is not initialized"
+
+        index_param = self.case_config.index_param()  # noqa: F841
+        search_param = self.case_config.search_param()
+
+        if filters:
+            gt = filters.get("id")
+            filter_sql = (
+                f'SELECT id, {search_param["metric_type"]}(embedding,{query}) AS score '  # noqa: S608
+                f'FROM {self.db_config["dbname"]}.{self.table_name}  '
+                f"WHERE id > {gt} "
+                f"ORDER BY score LIMIT {k};"
+            )
+            result = self.conn.query(filter_sql).result_rows
+            return [int(row[0]) for row in result]
+        else:  # noqa: RET505
+            select_sql = (
+                f'SELECT id, {search_param["metric_type"]}(embedding,{query}) AS score '  # noqa: S608
+                f'FROM {self.db_config["dbname"]}.{self.table_name}  '
+                f"ORDER BY score LIMIT {k};"
+            )
+            result = self.conn.query(select_sql).result_rows
+            return [int(row[0]) for row in result]

vectordb_bench/backend/clients/clickhouse/config.py

@@ -0,0 +1,60 @@
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+
+
+class ClickhouseConfig(DBConfig):
+    user_name: str = "clickhouse"
+    password: SecretStr
+    host: str = "localhost"
+    port: int = 8123
+    db_name: str = "default"
+
+    def to_dict(self) -> dict:
+        pwd_str = self.password.get_secret_value()
+        return {
+            "host": self.host,
+            "port": self.port,
+            "dbname": self.db_name,
+            "user": self.user_name,
+            "password": pwd_str,
+        }
+
+
+class ClickhouseIndexConfig(BaseModel):
+
+    metric_type: MetricType | None = None
+
+    def parse_metric(self) -> str:
+        if not self.metric_type:
+            return ""
+        return self.metric_type.value
+
+    def parse_metric_str(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "L2Distance"
+        if self.metric_type == MetricType.COSINE:
+            return "cosineDistance"
+        msg = f"Not Support for {self.metric_type}"
+        raise RuntimeError(msg)
+        return None
+
+
+class ClickhouseHNSWConfig(ClickhouseIndexConfig, DBCaseConfig):
+    M: int | None
+    efConstruction: int | None
+    ef: int | None = None
+    index: IndexType = IndexType.HNSW
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric_str(),
+            "index_type": self.index.value,
+            "params": {"M": self.M, "efConstruction": self.efConstruction},
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "met˝ric_type": self.parse_metric_str(),
+            "params": {"ef": self.ef},
+        }

vectordb_bench/backend/clients/elastic_cloud/elastic_cloud.py

@@ -81,7 +81,7 @@ class ElasticCloud(VectorDB):
         embeddings: Iterable[list[float]],
         metadata: list[int],
         **kwargs,
-    ) -> (int, Exception):
+    ) -> tuple[int, Exception]:
         """Insert the embeddings to the elasticsearch."""
         assert self.client is not None, "should self.init() first"
 

vectordb_bench/backend/clients/mariadb/cli.py

@@ -0,0 +1,122 @@
+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 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[
+        int | None,
+        click.option(
+            "--m",
+            type=int,
+            help="M parameter in MHNSW vector indexing",
+            required=False,
+        ),
+    ]
+
+    ef_search: Annotated[
+        int | None,
+        click.option(
+            "--ef-search",
+            type=int,
+            help="MariaDB system variable mhnsw_min_limit",
+            required=False,
+        ),
+    ]
+
+    max_cache_size: Annotated[
+        int | None,
+        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/backend/clients/mariadb/config.py

@@ -0,0 +1,73 @@
+from typing import TypedDict
+
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+
+
+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"
+        if self.metric_type == MetricType.COSINE:
+            return "cosine"
+        msg = f"Metric type {self.metric_type} is not supported!"
+        raise ValueError(msg)
+
+
+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,
+}