feat: add zvec support (#692)

* feat: support zvec
* update default m

Author: egolearner

README.md

@@ -63,6 +63,7 @@ All the database client supported
 | tencent_es               | `pip install vectordb-bench[tencent_es]`    |
 | alisql                   | `pip install 'vectordb-bench[alisql]'`      |
 | doris                    | `pip install vectordb-bench[doris]`         |
+| zvec                     | `pip install vectordb-bench[zvec]`          |
 
 ### Run
 
@@ -413,6 +414,27 @@ Options:
   --help                          Show this message and exit.
   ```
 
+### Run Zvec from command line
+
+```bash
+vectordbbench zvec --path Performance768D10M --db-label 16c64g-v0.1 \
+    --case-type Performance768D10M --num-concurrency 12,14,16,18,20 \
+    --quantize-type int8 --ef-search 118 --is-using-refiner
+```
+To list the options for zvec, execute vectordbbench zvec --help
+```
+  --path TEXT                     collection path  [required]
+  --m INTEGER                     HNSW index parameter m.
+  --ef-construction INTEGER       HNSW index parameter ef_construction
+  --ef-search INTEGER             HNSW index parameter ef for search
+  --quantize-type TEXT            HNSW index quantize type, fp16/int8
+                                  supported
+  --is-using-refiner              is using refiner, suitable for quantized
+                                  index, recall `ef-search` results then
+                                  refine with unquantized vector to `topk`
+                                  results
+```
+
 ### Run Doris from command line
 
 Doris supports ann index with type hnsw from version 4.0.x

pyproject.toml

@@ -78,6 +78,7 @@ all = [
     "lancedb",
     "mysql-connector-python",
     "turbopuffer[fast]",
+    'zvec',
 ]
 
 qdrant          = [ "qdrant-client" ]
@@ -106,6 +107,7 @@ oceanbase       = [ "mysql-connector-python" ]
 alisql          = [ "mysql-connector-python" ]
 doris           = [ "doris-vector-search" ]
 turbopuffer     = [ "turbopuffer" ]
+zvec            = [ "zvec" ]
 
 [project.urls]
 Repository  = "https://github.com/zilliztech/VectorDBBench"

vectordb_bench/backend/clients/__init__.py

@@ -56,6 +56,7 @@ class DB(Enum):
     AliSQL = "AlibabaCloudRDSMySQL"
     Doris = "Doris"
     TurboPuffer = "TurboPuffer"
+    Zvec = "Zvec"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
@@ -228,6 +229,11 @@ def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
 
             return AliSQL
 
+        if self == DB.Zvec:
+            from .zvec.zvec import Zvec
+
+            return Zvec
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -402,6 +408,11 @@ def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901, PLR0915
 
             return AliSQLConfig
 
+        if self == DB.Zvec:
+            from .zvec.config import ZvecConfig
+
+            return ZvecConfig
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -533,6 +544,11 @@ def case_config_cls(  # noqa: C901, PLR0911, PLR0912, PLR0915
 
             return HologresIndexConfig
 
+        if self == DB.Zvec:
+            from .zvec.config import ZvecHNSWIndexConfig
+
+            return ZvecHNSWIndexConfig
+
         if self == DB.TencentElasticsearch:
             from .tencent_elasticsearch.config import TencentElasticsearchIndexConfig
 

vectordb_bench/backend/clients/zvec/cli.py

@@ -0,0 +1,70 @@
+from typing import Annotated, Unpack
+
+import click
+
+from ....cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from .. import DB
+
+
+class ZvecTypedDict(CommonTypedDict):
+    path: Annotated[
+        str,
+        click.option("--path", type=str, help="collection path", required=True),
+    ]
+
+
+class ZvecHNSWTypedDict(CommonTypedDict, ZvecTypedDict):
+    m: Annotated[
+        int,
+        click.option("--m", type=int, default=50, help="HNSW index parameter m."),
+    ]
+    ef_construct: Annotated[
+        int,
+        click.option("--ef-construction", type=int, default=500, help="HNSW index parameter ef_construction"),
+    ]
+    ef_search: Annotated[
+        int,
+        click.option("--ef-search", type=int, default=300, help="HNSW index parameter ef for search"),
+    ]
+    quantize_type: Annotated[
+        int,
+        click.option("--quantize-type", type=str, default="", help="HNSW index quantize type, fp16/int8 supported"),
+    ]
+    is_using_refiner: Annotated[
+        bool,
+        click.option(
+            "--is-using-refiner",
+            is_flag=True,
+            default=False,
+            help="is using refiner, suitable for quantized index, "
+            "recall `ef-search` results then refine with unquantized vector to `topk` results",
+        ),
+    ]
+
+
+# default to hnsw
+@cli.command()
+@click_parameter_decorators_from_typed_dict(ZvecHNSWTypedDict)
+def Zvec(**parameters: Unpack[ZvecHNSWTypedDict]):
+    from .config import ZvecConfig, ZvecHNSWIndexConfig
+
+    run(
+        db=DB.Zvec,
+        db_config=ZvecConfig(
+            db_label=parameters["db_label"],
+            path=parameters["path"],
+        ),
+        db_case_config=ZvecHNSWIndexConfig(
+            M=parameters["m"],
+            ef_construction=parameters["ef_construction"],
+            ef_search=parameters["ef_search"],
+            quantize_type=parameters["quantize_type"],
+            is_using_refiner=parameters["is_using_refiner"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/zvec/config.py

@@ -0,0 +1,36 @@
+from pydantic import BaseModel
+
+from ..api import DBCaseConfig, DBConfig, MetricType
+
+
+class ZvecConfig(DBConfig):
+    """Zvec connection configuration."""
+
+    db_label: str
+    path: str
+
+    def to_dict(self) -> dict:
+        return {
+            "path": self.path,
+        }
+
+
+class ZvecIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+
+    def index_param(self) -> dict:
+        return {}
+
+    def search_param(self) -> dict:
+        return {}
+
+
+class ZvecHNSWIndexConfig(ZvecIndexConfig):
+    M: int | None = 50
+    ef_construction: int | None = 500
+
+    ef_search: int | None = 300
+
+    quantize_type: str = ""
+
+    is_using_refiner: bool = False