diff --git a/README.md b/README.md index 0ffa98f0e..77989527f 100644 --- a/README.md +++ b/README.md @@ -572,6 +572,8 @@ Options: ### Run Zvec from command line +HNSW remains the default, so existing commands are unchanged: + ```bash vectordbbench zvec --path Performance768D10M --db-label 16c64g-v0.1 \ --case-type Performance768D10M --num-concurrency 12,14,16,18,20 \ @@ -591,6 +593,67 @@ To list the options for zvec, execute vectordbbench zvec --help results ``` + #### Zvec DiskANN + + DiskANN requires a zvec SDK exposing `DiskAnnIndexParam`, `DiskAnnQueryParam`, and `Collection.close()`. + This adapter uses zvec's native C++ DiskANN, not the Microsoft Rust implementation. + From this updated checkout on Linux, in an environment with benchmark dependencies and the DiskANN-enabled SDK + already installed, install the adapter and run its checks: + + ```bash + python -m pip install --no-deps -e . + python -m pip install pytest + vectordbbench zvec --help + python -m pytest tests/test_zvec.py -q + ZVEC_NATIVE_TESTS=1 python -m pytest tests/test_zvec.py -k native_diskann -v + ``` + + The native test builds 1,200 32-dimensional vectors in a temporary directory, checks insertion and index completion, + closes/reopens the collection, searches it, and exercises rebuilding. It is a functionality test, not a recall or + performance benchmark. It is skipped unless explicitly enabled on Linux. Run it successfully before a large benchmark. + + | Option | Default | Meaning | + | --- | --- | --- | + | `--index-type` | hnsw | Select hnsw or diskann | + | `--metric-type` | IP when omitted | IP, COSINE or L2; use the same metric when building and reopening | + | `--max-degree` | 64 | DiskANN graph degree, 1-100 | + | `--build-list-size` | 100 | Construction candidate list, 10-100 | + | `--search-list-size` | 300 | Query candidate list, at least 1; the engine uses at least top-k | + | `--pq-chunk-num` | 0 | PQ bytes per vector; 0 is automatic, otherwise at most min(dimension, 1024) | + + DiskANN stores FP32 vectors and uses internal PQ for graph traversal. HNSW options (`--m`, `--ef-construction`, + `--ef-search`, `--quantize-type`, `--is-using-refiner`) are rejected in DiskANN mode, and DiskANN options are rejected + in HNSW mode. PQ chunks are not equivalent to HNSW INT8 scalar quantization. + + After the checks pass, build a separate DiskANN collection. The load stage includes `optimize()`, which builds the + target index; flushing alone does not establish that a DiskANN index was built. `--drop-old` replaces any existing + DiskANN collection at this path. The adapter refuses to use or delete an existing non-DiskANN collection. + + ```bash + export DATASET_LOCAL_DIR="$PWD/dataset" + vectordbbench zvec --index-type diskann --path Performance768D10M_DiskANN \ + --case-type Performance768D10M --metric-type IP \ + --max-degree 64 --build-list-size 100 --pq-chunk-num 96 \ + --drop-old --load --skip-search-serial --skip-search-concurrent + ``` + + Then search the same index without reloading or rebuilding it: + + ```bash + vectordbbench zvec --index-type diskann --path Performance768D10M_DiskANN \ + --case-type Performance768D10M --metric-type IP \ + --max-degree 64 --build-list-size 100 --pq-chunk-num 96 \ + --search-list-size 300 --num-concurrency 12,14,16,18,20 \ + --skip-drop-old --skip-load --search-serial --search-concurrent + ``` + + Search-only runs validate the stored dimensions, metric, build parameters, positive document count and complete + indexing. Keep build parameters unchanged when reopening; sweep `--search-list-size` to compare QPS at matched recall. + The examples explicitly use IP, the adapter's historical default. A dataset labelled COSINE does not change that + default or prove that its vectors are normalized. For a COSINE comparison, build both algorithms with + `--metric-type COSINE` in separate paths; do not relabel an existing IP index as COSINE. Multiple search processes can + duplicate in-memory PQ codes and metadata, so measure total process memory as well as QPS. + ### Run Doris from command line Doris supports ann index with type hnsw from version 4.0.x diff --git a/tests/test_zvec.py b/tests/test_zvec.py new file mode 100644 index 000000000..18f132f2a --- /dev/null +++ b/tests/test_zvec.py @@ -0,0 +1,384 @@ +import importlib.util +import os +from pathlib import Path +import pickle +import random +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import pytest +from click.testing import CliRunner +from pydantic import ValidationError + +from vectordb_bench.backend.clients import DB +from vectordb_bench.backend.clients.api import IndexType, MetricType +from vectordb_bench.backend.clients.zvec.config import ZvecDiskANNIndexConfig, ZvecHNSWIndexConfig +from vectordb_bench.backend.filter import non_filter + + +def test_diskann_config_and_registry_roundtrip(): + config = ZvecDiskANNIndexConfig(metric_type=MetricType.COSINE, pq_chunk_num=96) + assert config.index_param() == {"max_degree": 64, "list_size": 100, "pq_chunk_num": 96} + assert config.search_param() == {"index": IndexType.DISKANN, "list_size": 300} + assert DB.Zvec.case_config_cls(IndexType.DISKANN) is ZvecDiskANNIndexConfig + saved = config.model_dump(mode="json") + assert DB.Zvec.case_config_cls(saved["index"])(**saved) == config + assert pickle.loads(pickle.dumps(config)) == config + assert DB.Zvec.case_config_cls() is ZvecHNSWIndexConfig + assert DB.Zvec.case_config_cls(IndexType.HNSW) is ZvecHNSWIndexConfig + + +@pytest.mark.parametrize( + "values", + [ + {"max_degree": 0}, + {"max_degree": 101}, + {"build_list_size": 9}, + {"build_list_size": 101}, + {"search_list_size": 0}, + {"pq_chunk_num": -1}, + {"pq_chunk_num": 1025}, + {"index": "HNSW"}, + ], +) +def test_diskann_config_rejects_unsupported_values(values): + with pytest.raises(ValidationError): + ZvecDiskANNIndexConfig(**values) + + +@pytest.fixture +def adapter(monkeypatch): + sdk = ModuleType("zvec") + sdk.IndexType = SimpleNamespace(DISKANN="DISKANN", HNSW="HNSW") + sdk.MetricType = SimpleNamespace(IP="IP", COSINE="COSINE", L2="L2") + sdk.QuantizeType = SimpleNamespace(UNDEFINED="UNDEFINED", INT8="INT8", FP16="FP16", INT4="INT4") + sdk.DataType = SimpleNamespace(INT64="INT64", STRING="STRING", VECTOR_FP32="VECTOR_FP32") + sdk.LogLevel = SimpleNamespace(WARN="WARN") + sdk.init = Mock() + sdk.CollectionOption = lambda **kwargs: SimpleNamespace(**kwargs) + sdk.OptimizeOption = lambda: SimpleNamespace() + sdk.InvertIndexParam = lambda **kwargs: SimpleNamespace(**kwargs) + sdk.FieldSchema = lambda *args, **kwargs: SimpleNamespace() + sdk.VectorSchema = lambda name, data_type, **kwargs: SimpleNamespace(name=name, data_type=data_type, **kwargs) + sdk.CollectionSchema = lambda **kwargs: SimpleNamespace(**kwargs) + sdk.Doc = lambda **kwargs: SimpleNamespace(**kwargs) + sdk.VectorQuery = lambda **kwargs: SimpleNamespace(**kwargs) + sdk.DiskAnnIndexParam = lambda **kwargs: SimpleNamespace(type="DISKANN", **kwargs) + sdk.DiskAnnQueryParam = lambda **kwargs: SimpleNamespace(type="DISKANN", **kwargs) + sdk.HnswIndexParam = lambda **kwargs: SimpleNamespace(type="HNSW", **kwargs) + sdk.HnswQueryParam = lambda **kwargs: SimpleNamespace(type="HNSW", **kwargs) + sdk.open = Mock() + sdk.create_and_open = Mock() + monkeypatch.setitem(sys.modules, "zvec", sdk) + source = Path(__file__).parents[1] / "vectordb_bench/backend/clients/zvec/zvec.py" + spec = importlib.util.spec_from_file_location("vectordb_bench.backend.clients.zvec._unit_adapter", source) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.Zvec, sdk + + +def make_collection(adapter, *, config=None, count=100, completeness=1.0, dimension=768): + client, _sdk = adapter + config = config or ZvecDiskANNIndexConfig(pq_chunk_num=96) + field = SimpleNamespace(dimension=dimension, index_param=client._parse_index_param(config)) + collection = Mock() + collection.schema.vector.return_value = field + collection.stats = SimpleNamespace(doc_count=count, index_completeness={"dense": completeness}) + return collection + + +def make_client(adapter, config=None): + client, _sdk = adapter + instance = object.__new__(client) + instance.case_config = config or ZvecDiskANNIndexConfig(pq_chunk_num=96) + instance.table_name = "test" + instance.dim = 768 + instance.path = "unused" + instance.with_scalar_labels = False + instance._scalar_label_field = "label" + instance.query_param = client._parse_query_param(instance.case_config) + instance.option = SimpleNamespace(read_only=False, enable_mmap=True) + instance.expr = "" + return instance + + +def test_sdk_parameter_mapping_does_not_conflate_build_and_query(adapter): + client, sdk = adapter + config = ZvecDiskANNIndexConfig(metric_type=MetricType.COSINE, pq_chunk_num=96) + index = client._parse_index_param(config) + query = client._parse_query_param(config) + assert index.type == sdk.IndexType.DISKANN + assert index.metric_type == sdk.MetricType.COSINE + assert index.quantize_type == sdk.QuantizeType.UNDEFINED + assert (index.max_degree, index.list_size, index.pq_chunk_num) == (64, 100, 96) + assert query.list_size == 300 + hnsw = ZvecHNSWIndexConfig(M=50, ef_construction=500, ef_search=118, quantize_type="int8", is_using_refiner=True) + index = client._parse_index_param(hnsw) + query = client._parse_query_param(hnsw) + assert (index.type, index.m, index.ef_construction, index.quantize_type) == ("HNSW", 50, 500, "INT8") + assert query.ef == 118 and query.is_using_refiner + assert index.metric_type == sdk.MetricType.IP + + +@pytest.mark.parametrize( + "symbol, method", + [ + ("DiskAnnIndexParam", "_parse_index_param"), + ("DiskAnnQueryParam", "_parse_query_param"), + ], +) +def test_diskann_reports_missing_sdk_support(adapter, monkeypatch, symbol, method): + client, sdk = adapter + monkeypatch.delattr(sdk, symbol) + with pytest.raises(ValueError, match=symbol): + getattr(client, method)(ZvecDiskANNIndexConfig()) + assert client._parse_query_param(ZvecHNSWIndexConfig()).type == "HNSW" + + +def test_diskann_creation_rejects_hnsw_directory_before_delete(adapter, tmp_path): + client, sdk = adapter + collection = make_collection(adapter) + collection.schema.vector.return_value.index_param.type = "HNSW" + sdk.open.return_value = collection + with pytest.raises(ValueError, match="non-DiskANN"): + client(768, {"path": str(tmp_path)}, ZvecDiskANNIndexConfig(), drop_old=True) + collection.destroy.assert_not_called() + collection.close.assert_called_once() + sdk.create_and_open.assert_not_called() + + +def test_diskann_rebuild_allows_changed_build_parameters(adapter, tmp_path): + client, sdk = adapter + collection = make_collection(adapter) + + def mark_destroyed(): + collection.close.side_effect = ValueError("collection is already destroyed.") + + collection.destroy.side_effect = mark_destroyed + sdk.open.return_value = collection + config = ZvecDiskANNIndexConfig(max_degree=32, pq_chunk_num=48) + client(768, {"path": str(tmp_path)}, config, drop_old=True) + assert not sdk.open.call_args.kwargs["option"].read_only + collection.destroy.assert_called_once() + collection.close.assert_not_called() + sdk.create_and_open.assert_called_once() + sdk.create_and_open.return_value.close.assert_called_once() + + +def test_diskann_failed_destroy_closes_collection_without_recreating(adapter, tmp_path): + client, sdk = adapter + collection = make_collection(adapter) + collection.destroy.side_effect = ValueError("destroy failed") + sdk.open.return_value = collection + with pytest.raises(ValueError, match="destroy failed"): + client(768, {"path": str(tmp_path)}, ZvecDiskANNIndexConfig(), drop_old=True) + collection.destroy.assert_called_once() + collection.close.assert_called_once() + sdk.create_and_open.assert_not_called() + + +def test_diskann_validates_pq_dimensions_before_open(adapter, tmp_path): + client, sdk = adapter + with pytest.raises(ValueError, match="pq_chunk_num"): + client(8, {"path": str(tmp_path)}, ZvecDiskANNIndexConfig(pq_chunk_num=96), drop_old=True) + sdk.open.assert_not_called() + sdk.create_and_open.assert_not_called() + + +def test_diskann_new_directory_and_search_only_reopen(adapter, tmp_path): + client, sdk = adapter + path = tmp_path / "diskann" + config = ZvecDiskANNIndexConfig(pq_chunk_num=96) + with pytest.raises(ValueError, match="does not exist"): + client(768, {"path": str(path)}, config, drop_old=False) + client(768, {"path": str(path)}, config, drop_old=True) + sdk.create_and_open.assert_called_once() + sdk.create_and_open.return_value.close.assert_called_once() + path.mkdir() + collection = make_collection(adapter, config=config) + sdk.open.return_value = collection + client(768, {"path": str(path)}, config, drop_old=False) + assert sdk.open.call_args.kwargs["option"].read_only + collection.destroy.assert_not_called() + collection.close.assert_called_once() + + +@pytest.mark.parametrize("change", ["metric", "pq", "degree", "dimension", "incomplete", "empty"]) +def test_diskann_search_only_checks_persisted_schema_and_readiness(adapter, tmp_path, change): + client, sdk = adapter + collection = make_collection(adapter) + field = collection.schema.vector.return_value + if change == "metric": + field.index_param.metric_type = "COSINE" + elif change == "pq": + field.index_param.pq_chunk_num = 48 + elif change == "degree": + field.index_param.max_degree = 32 + elif change == "dimension": + field.dimension = 128 + elif change == "incomplete": + collection.stats.index_completeness["dense"] = 0.5 + else: + collection.stats.doc_count = 0 + sdk.open.return_value = collection + with pytest.raises((ValueError, RuntimeError)): + client(768, {"path": str(tmp_path)}, ZvecDiskANNIndexConfig(pq_chunk_num=96), drop_old=False) + collection.close.assert_called_once() + collection.destroy.assert_not_called() + + +def test_diskann_session_closes_on_error(adapter): + instance = make_client(adapter) + collection = make_collection(adapter) + adapter[1].open.return_value = collection + with pytest.raises(RuntimeError, match="test body"): + with instance.init(): + raise RuntimeError("test body") + collection.close.assert_called_once() + assert instance.collection is None + + +def test_diskann_insert_checks_partial_and_malformed_statuses(adapter): + instance = make_client(adapter) + instance.collection = Mock() + good = SimpleNamespace(ok=lambda: True) + bad = SimpleNamespace(ok=lambda: False) + instance.collection.insert.return_value = [good, bad] + count, error = instance.insert_embeddings([[0.0], [1.0]], [10, 11]) + assert count == 1 and error.non_retryable and error.inserted_count == 1 + instance.collection.insert.return_value = [good] + count, error = instance.insert_embeddings([[0.0], [1.0]], [10, 11]) + assert count == 0 and error.non_retryable + instance.collection.insert.return_value = [good, good] + assert instance.insert_embeddings([[0.0], [1.0]], [10, 11]) == (2, None) + count, error = instance.insert_embeddings([[0.0]], [10, 11]) + assert count == 0 and error.non_retryable + + +def test_diskann_optimize_requires_exact_count_and_index_completion(adapter): + instance = make_client(adapter) + instance.collection = make_collection(adapter, count=99, completeness=0.0) + with pytest.raises(RuntimeError, match="expected 100"): + instance.optimize(data_size=100) + instance.collection.optimize.assert_not_called() + instance.collection.stats.doc_count = 100 + with pytest.raises(RuntimeError, match="not fully indexed"): + instance.optimize(data_size=100) + instance.collection.optimize.side_effect = lambda **_kwargs: instance.collection.stats.index_completeness.update( + dense=1.0 + ) + instance.optimize(data_size=100) + assert instance.collection.stats.index_completeness["dense"] == 1.0 + + +def test_search_passes_diskann_query_parameters_without_refiner(adapter): + instance = make_client(adapter) + instance.collection = Mock() + instance.collection.query.return_value = [SimpleNamespace(id="7"), SimpleNamespace(id="8")] + assert instance.search_embedding([0.1] * 768, k=2) == [7, 8] + arguments = instance.collection.query.call_args.kwargs + assert arguments["vectors"].param.type == "DISKANN" + assert arguments["vectors"].param.list_size == 300 + assert arguments["topk"] == 2 and arguments["output_fields"] == [] + + +def test_cli_defaults_and_diskann_selection(monkeypatch): + from vectordb_bench.backend.clients.zvec import cli as zvec_cli + + captured = {} + monkeypatch.setattr(zvec_cli, "run", lambda **values: captured.update(values)) + runner = CliRunner() + result = runner.invoke(zvec_cli.Zvec, ["--path", "unused", "--dry-run"]) + assert result.exit_code == 0, result.output + assert isinstance(captured["db_case_config"], ZvecHNSWIndexConfig) + assert captured["db_case_config"] == ZvecHNSWIndexConfig() + result = runner.invoke( + zvec_cli.Zvec, + [ + "--path", + "unused", + "--index-type", + "diskann", + "--metric-type", + "cosine", + "--max-degree", + "64", + "--build-list-size", + "100", + "--pq-chunk-num", + "96", + "--search-list-size", + "400", + "--case-type", + "Performance768D10M", + "--skip-drop-old", + "--skip-load", + "--dry-run", + ], + ) + assert result.exit_code == 0, result.output + config = captured["db_case_config"] + assert isinstance(config, ZvecDiskANNIndexConfig) + assert config.metric_type == MetricType.COSINE + assert config.search_list_size == 400 and config.build_list_size == 100 + assert captured["drop_old"] is False and captured["load"] is False + + +@pytest.mark.parametrize( + "arguments", + [ + ["--index-type", "diskann", "--ef-search", "118"], + ["--index-type", "diskann", "--m", "50"], + ["--index-type", "diskann", "--quantize-type", "int8"], + ["--index-type", "diskann", "--is-using-refiner"], + ["--index-type", "hnsw", "--pq-chunk-num", "96"], + ["--index-type", "diskann", "--max-degree", "101"], + ["--index-type", "diskann", "--search-list-size", "0"], + ], +) +def test_cli_rejects_mixed_or_invalid_parameters(monkeypatch, arguments): + from vectordb_bench.backend.clients.zvec import cli as zvec_cli + + run = Mock() + monkeypatch.setattr(zvec_cli, "run", run) + result = CliRunner().invoke(zvec_cli.Zvec, ["--path", "unused", "--dry-run", *arguments]) + assert result.exit_code == 2, result.output + run.assert_not_called() + + +@pytest.mark.skipif( + sys.platform != "linux" or os.getenv("ZVEC_NATIVE_TESTS") != "1", + reason="Opt-in Linux DiskANN smoke test: ZVEC_NATIVE_TESTS=1", +) +def test_native_diskann_build_reopen_and_search(tmp_path): + import zvec as sdk + + from vectordb_bench.backend.clients.zvec.zvec import Zvec + + if not hasattr(sdk, "DiskAnnIndexParam"): + pytest.fail("Installed zvec must support DiskAnnIndexParam") + if not callable(getattr(sdk.Collection, "close", None)): + pytest.fail("Installed zvec must support Collection.close()") + random_values = random.Random(71) + vectors = [[random_values.random() for _ in range(32)] for _ in range(1200)] + config = ZvecDiskANNIndexConfig(metric_type=MetricType.L2, max_degree=32, pq_chunk_num=8) + connection = {"path": str(tmp_path / "diskann")} + client = Zvec(32, connection, config, drop_old=True) + with client.init(): + for start in range(0, len(vectors), 100): + batch = vectors[start : start + 100] + count, error = client.insert_embeddings(batch, list(range(start, start + len(batch)))) + assert error is None and count == len(batch) + client.optimize(data_size=len(vectors)) + assert client.collection.schema.vector("dense").index_param.type == sdk.IndexType.DISKANN + reopened = Zvec(32, connection, config, drop_old=False) + reopened.prepare_filter(non_filter) + with reopened.init(): + assert reopened.collection.stats.doc_count == len(vectors) + result = reopened.search_embedding(vectors[17], k=10) + assert len(result) == 10 and len(set(result)) == 10 + assert 17 in result + rebuilt = Zvec(32, connection, config, drop_old=True) + with rebuilt.init(): + assert rebuilt.collection.stats.doc_count == 0 diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index dbd83bb66..127e523c9 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/vectordb_bench/backend/clients/__init__.py @@ -663,8 +663,10 @@ def case_config_cls( # noqa: C901, PLR0911, PLR0912, PLR0915 return HologresIndexConfig if self == DB.Zvec: - from .zvec.config import ZvecHNSWIndexConfig + from .zvec.config import ZvecDiskANNIndexConfig, ZvecHNSWIndexConfig + if index_type == IndexType.DISKANN: + return ZvecDiskANNIndexConfig return ZvecHNSWIndexConfig if self == DB.TencentElasticsearch: diff --git a/vectordb_bench/backend/clients/zvec/cli.py b/vectordb_bench/backend/clients/zvec/cli.py index 0a9e09761..b58efb27a 100644 --- a/vectordb_bench/backend/clients/zvec/cli.py +++ b/vectordb_bench/backend/clients/zvec/cli.py @@ -1,6 +1,7 @@ from typing import Annotated, Unpack import click +from click.core import ParameterSource from ....cli.cli import ( CommonTypedDict, @@ -8,7 +9,8 @@ click_parameter_decorators_from_typed_dict, run, ) -from .. import DB +from .. import DB, MetricType +from .config import ZvecConfig, ZvecDiskANNIndexConfig, ZvecHNSWIndexConfig class ZvecTypedDict(CommonTypedDict): @@ -23,7 +25,7 @@ class ZvecHNSWTypedDict(CommonTypedDict, ZvecTypedDict): int, click.option("--m", type=int, default=50, help="HNSW index parameter m."), ] - ef_construct: Annotated[ + ef_construction: Annotated[ int, click.option("--ef-construction", type=int, default=500, help="HNSW index parameter ef_construction"), ] @@ -32,7 +34,7 @@ class ZvecHNSWTypedDict(CommonTypedDict, ZvecTypedDict): click.option("--ef-search", type=int, default=300, help="HNSW index parameter ef for search"), ] quantize_type: Annotated[ - int, + str, click.option("--quantize-type", type=str, default="", help="HNSW index quantize type, fp16/int8 supported"), ] is_using_refiner: Annotated[ @@ -47,24 +49,100 @@ class ZvecHNSWTypedDict(CommonTypedDict, ZvecTypedDict): ] -# default to hnsw +class ZvecAlgorithmTypedDict(ZvecHNSWTypedDict): + index_type: Annotated[ + str, + click.option( + "--index-type", + type=click.Choice(["hnsw", "diskann"], case_sensitive=False), + default="hnsw", + show_default=True, + help="Zvec vector index algorithm.", + ), + ] + metric_type: Annotated[ + str | None, + click.option( + "--metric-type", + type=click.Choice(["IP", "COSINE", "L2"], case_sensitive=False), + default=None, + help="Distance metric; omitted keeps the existing IP default.", + ), + ] + max_degree: Annotated[ + int, + click.option("--max-degree", type=click.IntRange(1, 100), default=64, help="DiskANN graph out-degree."), + ] + build_list_size: Annotated[ + int, + click.option( + "--build-list-size", + type=click.IntRange(10, 100), + default=100, + help="DiskANN candidate list size during construction.", + ), + ] + search_list_size: Annotated[ + int, + click.option( + "--search-list-size", + type=click.IntRange(min=1), + default=300, + help="DiskANN query candidate list size; the engine uses at least top-k.", + ), + ] + pq_chunk_num: Annotated[ + int, + click.option( + "--pq-chunk-num", + type=click.IntRange(0, 1024), + default=0, + help="DiskANN PQ bytes per vector; 0 selects automatically, nonzero must not exceed dimension.", + ), + ] + + @cli.command() -@click_parameter_decorators_from_typed_dict(ZvecHNSWTypedDict) -def Zvec(**parameters: Unpack[ZvecHNSWTypedDict]): - from .config import ZvecConfig, ZvecHNSWIndexConfig +@click_parameter_decorators_from_typed_dict(ZvecAlgorithmTypedDict) +def Zvec(**parameters: Unpack[ZvecAlgorithmTypedDict]): + index_type = parameters["index_type"].lower() + context = click.get_current_context() + unsupported = ( + ("m", "ef_construction", "ef_search", "quantize_type", "is_using_refiner") + if index_type == "diskann" + else ("max_degree", "build_list_size", "search_list_size", "pq_chunk_num") + ) + for parameter in unsupported: + if context.get_parameter_source(parameter) not in (None, ParameterSource.DEFAULT): + message = f"--{parameter.replace('_', '-')} does not apply to {index_type}" + raise click.UsageError(message) - run( - db=DB.Zvec, - db_config=ZvecConfig( - db_label=parameters["db_label"], - path=parameters["path"], - ), - db_case_config=ZvecHNSWIndexConfig( + metric_type = parameters["metric_type"] + metric_type = MetricType(metric_type.upper()) if metric_type else None + if index_type == "diskann": + case_config = ZvecDiskANNIndexConfig( + metric_type=metric_type, + max_degree=parameters["max_degree"], + build_list_size=parameters["build_list_size"], + search_list_size=parameters["search_list_size"], + pq_chunk_num=parameters["pq_chunk_num"], + ) + else: + case_config = ZvecHNSWIndexConfig( + metric_type=metric_type, 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"], + ) + + run( + db=DB.Zvec, + db_config=ZvecConfig( + db_label=parameters["db_label"], + path=parameters["path"], ), + db_case_config=case_config, **parameters, ) diff --git a/vectordb_bench/backend/clients/zvec/config.py b/vectordb_bench/backend/clients/zvec/config.py index 4c649d08a..0d4f453df 100644 --- a/vectordb_bench/backend/clients/zvec/config.py +++ b/vectordb_bench/backend/clients/zvec/config.py @@ -1,6 +1,8 @@ -from pydantic import BaseModel +from typing import Literal -from ..api import DBCaseConfig, DBConfig, MetricType +from pydantic import BaseModel, Field + +from ..api import DBCaseConfig, DBConfig, IndexType, MetricType class ZvecConfig(DBConfig): @@ -34,3 +36,21 @@ class ZvecHNSWIndexConfig(ZvecIndexConfig): quantize_type: str = "" is_using_refiner: bool = False + + +class ZvecDiskANNIndexConfig(ZvecIndexConfig): + index: Literal[IndexType.DISKANN] = IndexType.DISKANN + max_degree: int = Field(default=64, ge=1, le=100) + build_list_size: int = Field(default=100, ge=10, le=100) + search_list_size: int = Field(default=300, ge=1) + pq_chunk_num: int = Field(default=0, ge=0, le=1024) + + def index_param(self) -> dict: + return { + "max_degree": self.max_degree, + "list_size": self.build_list_size, + "pq_chunk_num": self.pq_chunk_num, + } + + def search_param(self) -> dict: + return {"index": self.index, "list_size": self.search_list_size} diff --git a/vectordb_bench/backend/clients/zvec/zvec.py b/vectordb_bench/backend/clients/zvec/zvec.py index 63cd87c9b..a96daaab4 100644 --- a/vectordb_bench/backend/clients/zvec/zvec.py +++ b/vectordb_bench/backend/clients/zvec/zvec.py @@ -1,5 +1,8 @@ +from __future__ import annotations + import logging from contextlib import contextmanager +from pathlib import Path import zvec from zvec import ( @@ -18,8 +21,8 @@ from vectordb_bench.backend.filter import Filter, FilterOp -from ..api import MetricType, VectorDB -from .config import ZvecConfig, ZvecHNSWIndexConfig, ZvecIndexConfig +from ..api import MetricType, NonRetryableInsertError, PartialInsertError, VectorDB +from .config import ZvecConfig, ZvecDiskANNIndexConfig, ZvecHNSWIndexConfig, ZvecIndexConfig log = logging.getLogger(__name__) @@ -49,6 +52,7 @@ def __init__( self.table_name = collection_name self.dim = dim self.path = db_config["path"] + self.expr = "" # 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" @@ -88,6 +92,13 @@ def __init__( self.query_param = Zvec._parse_query_param(db_case_config) + if isinstance(db_case_config, ZvecDiskANNIndexConfig): + if dim <= 0 or db_case_config.pq_chunk_num > dim: + message = "DiskANN requires a positive dimension and pq_chunk_num <= dimension" + raise ValueError(message) + self._prepare_diskann_collection(drop_old) + return + if drop_old: try: collection = zvec.open(self.path) @@ -99,11 +110,73 @@ def __init__( else: collection = zvec.open(self.path) + def _prepare_diskann_collection(self, drop_old: bool) -> None: + if Path(self.path).exists(): + option = CollectionOption(read_only=not drop_old, enable_mmap=True) + collection = zvec.open(self.path, option=option) + destroyed = False + try: + self._validate_diskann_collection(collection, require_ready=not drop_old, rebuilding=drop_old) + if drop_old: + collection.destroy() + destroyed = True + finally: + if not destroyed: + collection.close() + elif not drop_old: + message = f"DiskANN collection does not exist at {self.path}; build it before search-only benchmarking" + raise ValueError(message) + if drop_old: + collection = zvec.create_and_open(path=self.path, schema=self.schema, option=self.option) + collection.close() + log.info( + "DiskANN selected: path=%s metric=%s build=%s search=%s stored_vectors=FP32", + self.path, + self._parse_metric(self.case_config.metric_type), + self.case_config.index_param(), + self.search_config, + ) + + def _validate_diskann_collection( + self, + collection: zvec.Collection, + *, + require_ready: bool, + rebuilding: bool = False, + ) -> None: + field = collection.schema.vector("dense") + if field is None or field.index_param.type != zvec.IndexType.DISKANN: + message = "Refusing to use or delete a non-DiskANN collection; choose a separate DiskANN path" + raise ValueError(message) + if rebuilding: + return + expected = self._parse_index_param(self.case_config) + actual = field.index_param + if field.dimension != self.dim or any( + getattr(actual, name) != getattr(expected, name) + for name in ("metric_type", "quantize_type", "max_degree", "list_size", "pq_chunk_num") + ): + message = "Stored DiskANN schema differs from requested dimensions, metric or build parameters" + raise ValueError(message) + if require_ready: + stats = collection.stats + if stats.doc_count <= 0 or stats.index_completeness.get("dense") != 1.0: + message = "DiskANN collection is empty or not fully indexed; complete load and optimize first" + raise RuntimeError(message) + @contextmanager def init(self): self.collection = zvec.open(self.path, self.option) - yield - self.collection = None + try: + if isinstance(self.case_config, ZvecDiskANNIndexConfig): + self._validate_diskann_collection(self.collection, require_ready=self.option.read_only) + yield + finally: + try: + if isinstance(self.case_config, ZvecDiskANNIndexConfig): + self.collection.close() + finally: + self.collection = None def insert_embeddings( self, @@ -111,7 +184,12 @@ def insert_embeddings( metadata: list[int], labels_data: list[str] | None = None, **kwargs, - ) -> tuple[int, Exception]: + ) -> tuple[int, Exception | None]: + if isinstance(self.case_config, ZvecDiskANNIndexConfig) and ( + len(embeddings) != len(metadata) + or (self.with_scalar_labels and (labels_data is None or len(labels_data) != len(metadata))) + ): + return 0, NonRetryableInsertError("Embedding, ID and label batch lengths must match") docs = [] for i, id_ in enumerate(metadata): embedding = embeddings[i] @@ -128,7 +206,14 @@ def insert_embeddings( ) ) try: - self.collection.insert(docs) + statuses = self.collection.insert(docs) + if isinstance(self.case_config, ZvecDiskANNIndexConfig): + if len(statuses) != len(docs): + return 0, NonRetryableInsertError("zvec insert returned an incomplete status list") + inserted = sum(status.ok() for status in statuses) + if inserted != len(docs): + message = f"zvec DiskANN batch insert failed for {len(docs) - inserted}/{len(docs)} documents" + return inserted, PartialInsertError(message, inserted_count=inserted) return len(metadata), None except Exception as e: log.warning(f"Failed to insert data into Zvec table ({self.table_name}), error: {e}") @@ -153,7 +238,24 @@ def search_embedding( return [int(result.id) for result in results] def optimize(self, data_size: int | None = None): + if isinstance(self.case_config, ZvecDiskANNIndexConfig): + self._validate_diskann_collection(self.collection, require_ready=False) + if data_size is not None and self.collection.stats.doc_count != data_size: + message = f"DiskANN document count {self.collection.stats.doc_count} differs from expected {data_size}" + raise RuntimeError(message) + log.info("Building DiskANN index via optimize: documents=%s", self.collection.stats.doc_count) self.collection.optimize(option=OptimizeOption()) + if isinstance(self.case_config, ZvecDiskANNIndexConfig): + self._validate_diskann_collection(self.collection, require_ready=True) + stats = self.collection.stats + if data_size is not None and stats.doc_count != data_size: + message = f"DiskANN document count after optimize is {stats.doc_count}, expected {data_size}" + raise RuntimeError(message) + log.info( + "DiskANN optimize completed: documents=%s index_completeness=%s", + stats.doc_count, + stats.index_completeness["dense"], + ) def prepare_filter(self, filters: Filter): self.option = CollectionOption(read_only=True, enable_mmap=True) @@ -181,7 +283,18 @@ def _parse_metric(cls, metric_type: MetricType) -> zvec.MetricType: return d[metric_type] @classmethod - def _parse_index_param(cls, index_config: ZvecIndexConfig) -> zvec.HnswIndexParam: + def _parse_index_param(cls, index_config: ZvecIndexConfig) -> zvec.HnswIndexParam | zvec.DiskAnnIndexParam: + if isinstance(index_config, ZvecDiskANNIndexConfig): + if not hasattr(zvec, "DiskAnnIndexParam"): + message = "Installed zvec SDK does not expose DiskAnnIndexParam; install a DiskANN-enabled zvec build" + raise ValueError(message) + return zvec.DiskAnnIndexParam( + metric_type=Zvec._parse_metric(index_config.metric_type), + max_degree=index_config.max_degree, + list_size=index_config.build_list_size, + pq_chunk_num=index_config.pq_chunk_num, + quantize_type=QuantizeType.UNDEFINED, + ) if isinstance(index_config, ZvecHNSWIndexConfig): return zvec.HnswIndexParam( metric_type=Zvec._parse_metric(index_config.metric_type), @@ -193,7 +306,12 @@ def _parse_index_param(cls, index_config: ZvecIndexConfig) -> zvec.HnswIndexPara raise ValueError(message) @classmethod - def _parse_query_param(cls, index_config: ZvecIndexConfig) -> zvec.HnswQueryParam: + def _parse_query_param(cls, index_config: ZvecIndexConfig) -> zvec.HnswQueryParam | zvec.DiskAnnQueryParam: + if isinstance(index_config, ZvecDiskANNIndexConfig): + if not hasattr(zvec, "DiskAnnQueryParam"): + message = "Installed zvec SDK does not expose DiskAnnQueryParam; install a DiskANN-enabled zvec build" + raise ValueError(message) + return zvec.DiskAnnQueryParam(list_size=index_config.search_list_size) if isinstance(index_config, ZvecHNSWIndexConfig): return zvec.HnswQueryParam( ef=index_config.ef_search,