feat: add support for PolarDB

- add PolarDB vector search client with FAISS_HNSW_FLAT, FAISS_HNSW_PQ, and FAISS_HNSW_SQ index types
- add CLI integration with hnswflat, hnswpq, and hnswsq benchmark commands
- add frontend (Streamlit) UI support with index type selection, HNSW/PQ/SQ parameter configuration

Author: nanlongyu

README.md

@@ -62,6 +62,7 @@ All the database client supported
 | hologres                 | `pip install vectordb-bench[hologres]`      |
 | tencent_es               | `pip install vectordb-bench[tencent_es]`    |
 | alisql                   | `pip install 'vectordb-bench[alisql]'`      |
+| polardb                  | `pip install vectordb-bench[polardb]`       |
 | doris                    | `pip install vectordb-bench[doris]`         |
 | zvec                     | `pip install vectordb-bench[zvec]`          |
 | endee                    | `pip install vectordb-bench[endee]`         |
@@ -527,6 +528,47 @@ To list the options for Lindorm, execute `vectordbbench lindormhnsw --help`, The
   --ef-search INTEGER             hnsw ef-search  [required]
 ```
 
+### Run PolarDB from command line
+
+PolarDB supports index types: faiss_hnsw_flat, faiss_hnsw_pq, and faiss_hnsw_sq.
+
+**Example: Run faiss_hnsw_flat benchmark**
+
+```shell
+vectordbbench polardbhnswflat \
+  --case-type Performance768D1M \
+  --username <db_user> \
+  --password '<db_password>' \
+  --host <db_host> \
+  --port 3306 \
+  --m 16 \
+  --ef-construction 256 \
+  --ef-search 256 \
+  --insert-workers 64 \
+  --num-concurrency '10,20,40,60,80' \
+  --concurrency-duration 60 \
+  --task-label <task_label> \
+  --db-label <db_label> \
+  --skip-search-serial \
+  --post-load-index
+```
+
+To list the options for PolarDB, execute `vectordbbench polardbhnswflat --help`. The following are some PolarDB-specific command-line options.
+
+```text
+  --username TEXT                  Username  [required]
+  --password TEXT                  Password
+  --host TEXT                      Db host  [default: 127.0.0.1]
+  --port INTEGER                   Db Port  [default: 3306]
+  --database TEXT                  Database name  [default: vectordbbench]
+  --m INTEGER                      M parameter (max_degree) in HNSW
+  --ef-construction INTEGER        ef_construction parameter in HNSW
+  --ef-search INTEGER              polar_vector_index_hnsw_ef_search session variable
+  --insert-workers INTEGER         Number of concurrent threads for data insertion
+  --post-load-index / --inline-index
+                                   Create index after load or inline at table creation
+```
+
 #### Using a configuration file.
 
 The vectordbbench command can optionally read some or all the options from a yaml formatted configuration file.

install/requirements_py3.11.txt

@@ -26,5 +26,6 @@ pymilvus
 clickhouse_connect
 pyvespa
 mysql-connector-python
+PyMySQL
 packaging
-hdrhistogram>=0.10.1
+hdrhistogram>=0.10.1

pyproject.toml

@@ -106,6 +106,7 @@ vespa           = [ "pyvespa" ]
 lancedb         = [ "lancedb" ]
 oceanbase       = [ "mysql-connector-python" ]
 alisql          = [ "mysql-connector-python" ]
+polardb         = [ "PyMySQL" ]
 doris           = [ "doris-vector-search" ]
 turbopuffer     = [ "turbopuffer" ]
 zvec            = [ "zvec" ]

vectordb_bench/backend/clients/__init__.py

@@ -59,6 +59,7 @@ class DB(Enum):
     Zvec = "Zvec"
     Endee = "Endee"
     Lindorm = "Lindorm"
+    PolarDB = "PolarDB"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
@@ -246,6 +247,11 @@ def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
 
             return LindormVector
 
+        if self == DB.PolarDB:
+            from .polardb.polardb import PolarDB
+
+            return PolarDB
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -435,6 +441,11 @@ def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901, PLR0915
 
             return LindormConfig
 
+        if self == DB.PolarDB:
+            from .polardb.config import PolarDBConfig
+
+            return PolarDBConfig
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -581,6 +592,11 @@ def case_config_cls(  # noqa: C901, PLR0911, PLR0912, PLR0915
 
             return AliSQLIndexConfig
 
+        if self == DB.PolarDB:
+            from .polardb.config import _polardb_case_config
+
+            return _polardb_case_config.get(index_type)
+
         if self == DB.Doris:
             from .doris.config import DorisCaseConfig
 

vectordb_bench/backend/clients/polardb/__init__.py

@@ -0,0 +1,248 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Annotated, Unpack
+
+if TYPE_CHECKING:
+    from .config import PolarDBConfig
+
+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 PolarDBTypedDict(CommonTypedDict):
+    user_name: Annotated[
+        str,
+        click.option(
+            "--username",
+            type=str,
+            help="Username",
+            required=True,
+        ),
+    ]
+    password: Annotated[
+        str,
+        click.option(
+            "--password",
+            type=str,
+            help="Password",
+            default="",
+        ),
+    ]
+
+    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",
+        ),
+    ]
+
+    database: Annotated[
+        str,
+        click.option(
+            "--database",
+            type=str,
+            help="Database name",
+            default="vectordbbench",
+        ),
+    ]
+
+    unix_socket: Annotated[
+        str,
+        click.option(
+            "--unix-socket",
+            type=str,
+            help="Unix socket path (overrides host/port if set)",
+            default="",
+        ),
+    ]
+
+
+class PolarDBHNSWTypedDict(PolarDBTypedDict):
+    m: Annotated[
+        int,
+        click.option(
+            "--m",
+            type=int,
+            help="M parameter (max_degree) in HNSW",
+            default=16,
+        ),
+    ]
+
+    ef_construction: Annotated[
+        int,
+        click.option(
+            "--ef-construction",
+            type=int,
+            help="ef_construction parameter in HNSW",
+            default=200,
+        ),
+    ]
+
+    ef_search: Annotated[
+        int,
+        click.option(
+            "--ef-search",
+            type=int,
+            help="polar_vector_index_hnsw_ef_search session variable",
+            default=64,
+        ),
+    ]
+
+    insert_workers: Annotated[
+        int,
+        click.option(
+            "--insert-workers",
+            type=int,
+            help="Number of concurrent threads for data insertion",
+            default=10,
+        ),
+    ]
+
+    post_load_index: Annotated[
+        bool,
+        click.option(
+            "--post-load-index/--inline-index",
+            type=bool,
+            help="Create vector index via ALTER TABLE after data load; "
+            "otherwise create index inline during table creation",
+            default=False,
+        ),
+    ]
+
+
+class PolarDBHNSWPQTypedDict(PolarDBHNSWTypedDict):
+    pq_m: Annotated[
+        int,
+        click.option(
+            "--pq-m",
+            type=int,
+            help="PQ subquantizer count (must divide dimension)",
+            default=1,
+        ),
+    ]
+
+    pq_nbits: Annotated[
+        int,
+        click.option(
+            "--pq-nbits",
+            type=int,
+            help="PQ bits per subquantizer (max 24)",
+            default=8,
+        ),
+    ]
+
+
+class PolarDBHNSWSQTypedDict(PolarDBHNSWTypedDict):
+    sq_type: Annotated[
+        str,
+        click.option(
+            "--sq-type",
+            type=str,
+            help="SQ quantizer type (8bit, 4bit, fp16, bf16, 6bit, etc.)",
+            default="8bit",
+        ),
+    ]
+
+
+def _build_db_config(parameters: dict) -> PolarDBConfig:
+    from .config import PolarDBConfig
+
+    pwd = parameters["password"]
+    sock = parameters["unix_socket"]
+    return PolarDBConfig(
+        db_label=parameters["db_label"],
+        user_name=parameters["username"],
+        password=SecretStr(pwd) if pwd else None,
+        host=parameters["host"],
+        port=parameters["port"],
+        database=parameters["database"],
+        unix_socket=sock if sock else None,
+    )
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PolarDBHNSWTypedDict)
+def PolarDBHNSWFlat(
+    **parameters: Unpack[PolarDBHNSWTypedDict],
+):
+    from .config import PolarDBHNSWFlatConfig
+
+    run(
+        db=DB.PolarDB,
+        db_config=_build_db_config(parameters),
+        db_case_config=PolarDBHNSWFlatConfig(
+            M=parameters["m"],
+            ef_construction=parameters["ef_construction"],
+            ef_search=parameters["ef_search"],
+            insert_workers=parameters["insert_workers"],
+            post_load_index=parameters["post_load_index"],
+        ),
+        **parameters,
+    )
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PolarDBHNSWPQTypedDict)
+def PolarDBHNSWPQ(
+    **parameters: Unpack[PolarDBHNSWPQTypedDict],
+):
+    from .config import PolarDBHNSWPQConfig
+
+    run(
+        db=DB.PolarDB,
+        db_config=_build_db_config(parameters),
+        db_case_config=PolarDBHNSWPQConfig(
+            M=parameters["m"],
+            ef_construction=parameters["ef_construction"],
+            ef_search=parameters["ef_search"],
+            insert_workers=parameters["insert_workers"],
+            post_load_index=parameters["post_load_index"],
+            pq_m=parameters["pq_m"],
+            pq_nbits=parameters["pq_nbits"],
+        ),
+        **parameters,
+    )
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PolarDBHNSWSQTypedDict)
+def PolarDBHNSWSQ(
+    **parameters: Unpack[PolarDBHNSWSQTypedDict],
+):
+    from .config import PolarDBHNSWSQConfig
+
+    run(
+        db=DB.PolarDB,
+        db_config=_build_db_config(parameters),
+        db_case_config=PolarDBHNSWSQConfig(
+            M=parameters["m"],
+            ef_construction=parameters["ef_construction"],
+            ef_search=parameters["ef_search"],
+            insert_workers=parameters["insert_workers"],
+            post_load_index=parameters["post_load_index"],
+            sq_type=parameters["sq_type"],
+        ),
+        **parameters,
+    )