add Databend benchmark

README.md

@@ -59,6 +59,7 @@ All the database client supported
 | vespa                    | `pip install 'vectordb-bench[vespa]'`         |
 | oceanbase                | `pip install 'vectordb-bench[oceanbase]'`     |
 | hologres                 | `pip install 'vectordb-bench[hologres]'`      |
+| databend                 | `pip install 'vectordb-bench[databend]'`      |
 
 ### Run
 

install/requirements_py3.11.txt

@@ -24,4 +24,5 @@ scikit-learn
 pymilvus
 clickhouse_connect
 pyvespa
-mysql-connector-python
+mysql-connector-python
+databend-driver

pyproject.toml

@@ -74,6 +74,7 @@ all = [
     "pyvespa",
     "lancedb",
     "mysql-connector-python",
+    "databend-driver",
 ]
 
 qdrant          = [ "qdrant-client" ]
@@ -98,6 +99,7 @@ clickhouse      = [ "clickhouse-connect" ]
 vespa           = [ "pyvespa" ]
 lancedb         = [ "lancedb" ]
 oceanbase       = [ "mysql-connector-python" ]
+databend        = [ "databend-driver" ]
 
 [project.urls]
 "repository" = "https://github.com/zilliztech/VectorDBBench"

vectordb_bench/backend/clients/__init__.py

@@ -51,6 +51,7 @@ class DB(Enum):
     OceanBase = "OceanBase"
     S3Vectors = "S3Vectors"
     Hologres = "Alibaba Cloud Hologres"
+    Databend = "Databend"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
@@ -200,6 +201,11 @@ def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
 
             return Hologres
 
+        if self == DB.Databend:
+            from .databend.databend import Databend
+
+            return Databend
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -351,6 +357,11 @@ def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901, PLR0915
 
             return HologresConfig
 
+        if self == DB.Databend:
+            from .databend.config import DatabendConfig
+
+            return DatabendConfig
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -475,7 +486,10 @@ def case_config_cls(  # noqa: C901, PLR0911, PLR0912
         if self == DB.Hologres:
             from .hologres.config import HologresIndexConfig
 
-            return HologresIndexConfig
+        if self == DB.Databend:
+            from .databend.config import DatabendIndexConfig
+
+            return DatabendIndexConfig
 
         # DB.Pinecone, DB.Chroma, DB.Redis
         return EmptyDBCaseConfig

vectordb_bench/backend/clients/databend/cli.py

@@ -0,0 +1,47 @@
+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 DatabendIndexConfig
+
+
+class DatabendTypedDict(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=8000, help="DB Port")]
+    user: Annotated[int, click.option("--user", type=str, default="root", help="DB user")]
+
+
+class DatabendHNSWTypedDict(CommonTypedDict, DatabendTypedDict, HNSWFlavor2): ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(DatabendHNSWTypedDict)
+def Databend(**parameters: Unpack[DatabendHNSWTypedDict]):
+    from .config import DatabendConfig
+
+    run(
+        db=DB.Databend,
+        db_config=DatabendConfig(
+            db_label=parameters["db_label"],
+            user=parameters["user"],
+            password=SecretStr(parameters["password"]) if parameters["password"] else None,
+            host=parameters["host"],
+            port=parameters["port"],
+        ),
+        db_case_config=DatabendIndexConfig(
+            metric_type=None,
+            m=parameters["m"],
+            ef_construct=parameters["ef_construction"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/databend/config.py

@@ -0,0 +1,72 @@
+from abc import abstractmethod
+from typing import TypedDict
+
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+
+
+class DatabendConfigDict(TypedDict):
+    user: str
+    password: str
+    host: str
+    port: int
+    database: str
+    secure: bool
+
+
+class DatabendConfig(DBConfig):
+    user: str = "root"
+    password: SecretStr
+    host: str = "localhost"
+    port: int = 8000
+    db_name: str = "default"
+    secure: bool = False
+
+    def to_dict(self) -> DatabendConfigDict:
+        pwd_str = self.password.get_secret_value()
+        return {
+            "host": self.host,
+            "port": self.port,
+            "database": self.db_name,
+            "user": self.user,
+            "password": pwd_str,
+            "secure": self.secure,
+        }
+
+
+class DatabendIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+    create_index_before_load: bool = True
+    create_index_after_load: bool = False
+    m: int | None
+    ef_construct: int | 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 "l2"
+        if self.metric_type == MetricType.COSINE:
+            return "cosine"
+        return "cosine"
+
+    @abstractmethod
+    def session_param(self):
+        pass
+
+    def index_param(self) -> dict:
+        return {
+            "m": self.m,
+            "metric_type": self.parse_metric_str(),
+            "ef_construct": self.ef_construct,
+        }
+
+    def search_param(self) -> dict:
+        return {}
+
+    def session_param(self) -> dict:
+        return {}