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

vectordb_bench/backend/clients/__init__.py

@@ -43,6 +43,8 @@ class DB(Enum):
     AliyunOpenSearch = "AliyunOpenSearch"
     MongoDB = "MongoDB"
     TiDB = "TiDB"
+    Clickhouse = "Clickhouse"
+    Vespa = "Vespa"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901
@@ -117,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
 
@@ -152,6 +159,11 @@ class DB(Enum):
 
             return Test
 
+        if self == DB.Vespa:
+            from .vespa.vespa import Vespa
+
+            return Vespa
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -228,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
 
@@ -263,10 +280,15 @@ class DB(Enum):
 
             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]:
@@ -310,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
 
@@ -350,6 +377,11 @@ class DB(Enum):
 
             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

@@ -162,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

@@ -1,83 +1,98 @@
-from typing import Annotated, Optional, Unpack
+from typing import Annotated, Unpack
 
 import click
-import os
 from pydantic import SecretStr
 
+from vectordb_bench.backend.clients import DB
+
 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,
-                          ),
+        str,
+        click.option(
+            "--username",
+            type=str,
+            help="Username",
+            required=True,
+        ),
     ]
     password: Annotated[
-        str, click.option("--password",
-                     type=str,
-                     help="Password",
-                     required=True,
-                     ),
+        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",
-                          ),
+        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",
-                          ),
+        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,
-                          ),
+        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,
-                                    ),
+        int | None,
+        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,
-                                    ),
+        int | None,
+        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,
-                                    ),
+        int | None,
+        click.option(
+            "--max-cache-size",
+            type=int,
+            help="MariaDB system variable mhnsw_max_cache_size",
+            required=False,
+        ),
     ]
 
 

vectordb_bench/backend/clients/mariadb/config.py

@@ -1,10 +1,13 @@
-from pydantic import SecretStr, BaseModel
 from typing import TypedDict
-from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
+
+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"""
+    so the names must match exactly mariadb API"""
 
     user: str
     password: str
@@ -36,10 +39,11 @@ class MariaDBIndexConfig(BaseModel):
     def parse_metric(self) -> str:
         if self.metric_type == MetricType.L2:
             return "euclidean"
-        elif self.metric_type == MetricType.COSINE:
+        if self.metric_type == MetricType.COSINE:
             return "cosine"
-        else:
-            raise ValueError(f"Metric type {self.metric_type} is not supported!")
+        msg = f"Metric type {self.metric_type} is not supported!"
+        raise ValueError(msg)
+
 
 class MariaDBHNSWConfig(MariaDBIndexConfig, DBCaseConfig):
     M: int | None
@@ -65,7 +69,5 @@ class MariaDBHNSWConfig(MariaDBIndexConfig, DBCaseConfig):
 
 
 _mariadb_case_config = {
-        IndexType.HNSW: MariaDBHNSWConfig,
+    IndexType.HNSW: MariaDBHNSWConfig,
 }
-
-