diff --git a/README.md b/README.md index bd7568da1..8a28bf469 100644 --- a/README.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index ef4792bee..7a3e8bf71 100644 --- a/pyproject.toml +++ b/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" diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index 3bd2d3189..f31f93d29 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/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 diff --git a/vectordb_bench/backend/clients/zvec/cli.py b/vectordb_bench/backend/clients/zvec/cli.py new file mode 100644 index 000000000..74fe786af --- /dev/null +++ b/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=100, 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, + ) diff --git a/vectordb_bench/backend/clients/zvec/config.py b/vectordb_bench/backend/clients/zvec/config.py new file mode 100644 index 000000000..dd2fb1a01 --- /dev/null +++ b/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 = 100 + ef_construction: int | None = 500 + + ef_search: int | None = 300 + + quantize_type: str = "" + + is_using_refiner: bool = False diff --git a/vectordb_bench/backend/clients/zvec/zvec.py b/vectordb_bench/backend/clients/zvec/zvec.py new file mode 100644 index 000000000..63cd87c9b --- /dev/null +++ b/vectordb_bench/backend/clients/zvec/zvec.py @@ -0,0 +1,214 @@ +import logging +from contextlib import contextmanager + +import zvec +from zvec import ( + CollectionOption, + CollectionSchema, + DataType, + Doc, + FieldSchema, + InvertIndexParam, + LogLevel, + OptimizeOption, + QuantizeType, + VectorQuery, + VectorSchema, +) + +from vectordb_bench.backend.filter import Filter, FilterOp + +from ..api import MetricType, VectorDB +from .config import ZvecConfig, ZvecHNSWIndexConfig, ZvecIndexConfig + +log = logging.getLogger(__name__) + +zvec.init(log_level=LogLevel.WARN) + + +class Zvec(VectorDB): + supported_filter_types: list[FilterOp] = [ + FilterOp.NonFilter, + FilterOp.NumGE, + FilterOp.StrEqual, + ] + + def __init__( + self, + dim: int, + db_config: ZvecConfig, + db_case_config: ZvecIndexConfig, + collection_name: str = "vector_bench_test", + drop_old: bool = False, + with_scalar_labels: bool = False, + **kwargs, + ): + self.name = "Zvec" + self.db_config = db_config + self.case_config = db_case_config + self.table_name = collection_name + self.dim = dim + self.path = db_config["path"] + # avoid the search_param being called every time during the search process + self.search_config = db_case_config.search_param() + self._scalar_id_field = "id" + self._scalar_label_field = "label" + self.with_scalar_labels = with_scalar_labels + + log.info(f"Search config: {self.search_config}") + + fields = [ + FieldSchema( + "id", DataType.INT64, nullable=False, index_param=InvertIndexParam(enable_range_optimization=True) + ), + ] + if with_scalar_labels: + fields.append( + FieldSchema( + self._scalar_label_field, + DataType.STRING, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=False), + ) + ) + self.schema = CollectionSchema( + name=self.table_name, + fields=fields, + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=dim, + index_param=Zvec._parse_index_param(db_case_config), + ), + ], + ) + + self.option = CollectionOption(read_only=False, enable_mmap=True) + + self.query_param = Zvec._parse_query_param(db_case_config) + + if drop_old: + try: + collection = zvec.open(self.path) + collection.destroy() + except Exception as e: + log.warning(f"Failed to drop table {self.table_name}: {e}") + + collection = zvec.create_and_open(path=self.path, schema=self.schema, option=self.option) + else: + collection = zvec.open(self.path) + + @contextmanager + def init(self): + self.collection = zvec.open(self.path, self.option) + yield + self.collection = None + + def insert_embeddings( + self, + embeddings: list[list[float]], + metadata: list[int], + labels_data: list[str] | None = None, + **kwargs, + ) -> tuple[int, Exception]: + docs = [] + for i, id_ in enumerate(metadata): + embedding = embeddings[i] + fields = ( + {"id": id_} if not self.with_scalar_labels else {"id": id_, self._scalar_label_field: labels_data[i]} + ) + docs.append( + Doc( + id=f"{id_}", + fields=fields, + vectors={ + "dense": embedding, + }, + ) + ) + try: + self.collection.insert(docs) + return len(metadata), None + except Exception as e: + log.warning(f"Failed to insert data into Zvec table ({self.table_name}), error: {e}") + return 0, e + + def search_embedding( + self, + query: list[float], + k: int = 100, + filters: dict | None = None, + ) -> list[int]: + if filters: + results = [] + else: + results = self.collection.query( + output_fields=[], + topk=k, + filter=self.expr, + vectors=VectorQuery(field_name="dense", vector=query, param=self.query_param), + ) + + return [int(result.id) for result in results] + + def optimize(self, data_size: int | None = None): + self.collection.optimize(option=OptimizeOption()) + + def prepare_filter(self, filters: Filter): + self.option = CollectionOption(read_only=True, enable_mmap=True) + log.debug("set readonly: %s", self.option.read_only) + + if filters.type == FilterOp.NonFilter: + self.expr = "" + elif filters.type == FilterOp.NumGE: + self.expr = f"{self._scalar_id_field} >= {filters.int_value}" + elif filters.type == FilterOp.StrEqual: + self.expr = f"{self._scalar_label_field} = '{filters.label_value}'" + else: + msg = f"Not support Filter for zvec - {filters}" + raise ValueError(msg) + + @classmethod + def _parse_metric(cls, metric_type: MetricType) -> zvec.MetricType: + if not metric_type: + return zvec.MetricType.IP + d = { + MetricType.COSINE: zvec.MetricType.COSINE, + MetricType.L2: zvec.MetricType.L2, + MetricType.IP: zvec.MetricType.IP, + } + return d[metric_type] + + @classmethod + def _parse_index_param(cls, index_config: ZvecIndexConfig) -> zvec.HnswIndexParam: + if isinstance(index_config, ZvecHNSWIndexConfig): + return zvec.HnswIndexParam( + metric_type=Zvec._parse_metric(index_config.metric_type), + m=index_config.M, + ef_construction=index_config.ef_construction, + quantize_type=Zvec._parse_quantize_type(index_config.quantize_type), + ) + message = f"Not support index type - {index_config}" + raise ValueError(message) + + @classmethod + def _parse_query_param(cls, index_config: ZvecIndexConfig) -> zvec.HnswQueryParam: + if isinstance(index_config, ZvecHNSWIndexConfig): + return zvec.HnswQueryParam( + ef=index_config.ef_search, + is_using_refiner=index_config.is_using_refiner, + ) + message = f"Not support index type - {index_config}" + raise ValueError(message) + + @classmethod + def _parse_quantize_type(cls, quantize_type: str) -> QuantizeType: + if not quantize_type: + return QuantizeType.UNDEFINED + d = { + "FP16": QuantizeType.FP16, + "INT8": QuantizeType.INT8, + "INT4": QuantizeType.INT4, + } + return d[quantize_type.upper()] diff --git a/vectordb_bench/backend/runner/mp_runner.py b/vectordb_bench/backend/runner/mp_runner.py index 9133e407a..b7823af37 100644 --- a/vectordb_bench/backend/runner/mp_runner.py +++ b/vectordb_bench/backend/runner/mp_runner.py @@ -1,4 +1,5 @@ import concurrent +import contextlib import logging import multiprocessing as mp import random @@ -66,6 +67,11 @@ def search( with cond: cond.wait() + # NOTE: Zvec allows multiple read-only opens, or one read-write open. + # Use prepare_filter to switch to read-only mode. + with contextlib.suppress(Exception): + self.db.prepare_filter(self.filters) + with self.db.init(): self.db.prepare_filter(self.filters) num, idx = len(test_data), random.randint(0, len(test_data) - 1) diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index a53f0d314..9d071716f 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -32,6 +32,7 @@ from ..backend.clients.vespa.cli import Vespa from ..backend.clients.weaviate_cloud.cli import Weaviate from ..backend.clients.zilliz_cloud.cli import ZillizAutoIndex +from ..backend.clients.zvec.cli import Zvec from .batch_cli import BatchCli from .cli import cli @@ -70,6 +71,7 @@ cli.add_command(AliSQLHNSW) cli.add_command(Doris) cli.add_command(TurboPuffer) +cli.add_command(Zvec) if __name__ == "__main__": diff --git a/vectordb_bench/frontend/config/styles.py b/vectordb_bench/frontend/config/styles.py index bce4561fd..e5f4eb1d9 100644 --- a/vectordb_bench/frontend/config/styles.py +++ b/vectordb_bench/frontend/config/styles.py @@ -71,6 +71,7 @@ def getPatternShape(i): DB.Doris: "https://doris.apache.org/images/logo.svg", DB.TurboPuffer: "https://turbopuffer.com/logo2.png", DB.CockroachDB: "https://raw.githubusercontent.com/cockroachdb/cockroach/master/docs/media/cockroach_db.png", + DB.Zvec: "https://zvec.org/img/zvec-logo-light.svg", } # RedisCloud color: #0D6EFD @@ -89,6 +90,7 @@ def getPatternShape(i): DB.TiDB.value: "#0D6EFD", DB.Vespa.value: "#61d790", DB.Doris.value: "#52CAA3", + DB.Zvec.value: "#219FFF", } COLORS_10 = [