From 5b305cd035941c24f2b864437092fd6b19a98a79 Mon Sep 17 00:00:00 2001 From: Alex Bondarev Date: Fri, 25 Sep 2026 14:39:58 -0500 Subject: [PATCH 1/5] feat(milvus): add native BF16 HNSW support Add IndexType.HNSW_BF16, HNSWBF16Config, and a MilvusHNSWBF16 CLI command. The client creates BFLOAT16_VECTOR fields and converts float32 insert batches and query vectors to round-to-nearest-even BF16 byte payloads accepted by pymilvus. The conversion uses NumPy bit operations, avoiding an additional ml_dtypes runtime dependency. Milvus still receives index_type "HNSW"; BF16 is selected by the collection field's data type, so FP32, FTS, and GPU configurations keep their existing behavior. Tests cover the rounding and special-value encoding, BFLOAT16_VECTOR schema selection against the FP32 default, insert and batch-search conversion, the CLI command wiring, and HNSW_BF16 case-config registration. Co-authored-by: RJ Silk --- tests/test_milvus.py | 154 +++++++++++++++++- tests/test_milvus_zilliz_cli.py | 35 ++++ vectordb_bench/backend/clients/api.py | 1 + vectordb_bench/backend/clients/milvus/cli.py | 20 +++ .../backend/clients/milvus/config.py | 12 ++ .../backend/clients/milvus/milvus.py | 31 +++- vectordb_bench/cli/vectordbbench.py | 3 +- 7 files changed, 248 insertions(+), 8 deletions(-) diff --git a/tests/test_milvus.py b/tests/test_milvus.py index 55128a1b2..3d5f79dba 100644 --- a/tests/test_milvus.py +++ b/tests/test_milvus.py @@ -8,14 +8,25 @@ from types import SimpleNamespace from unittest.mock import MagicMock, call +import numpy as np import pytest from pydantic import SecretStr +from pymilvus import DataType from vectordb_bench.backend.cases import CaseType from vectordb_bench.backend.clients import DB -from vectordb_bench.backend.clients.api import IndexType -from vectordb_bench.backend.clients.milvus.config import MilvusConfig, MilvusFtsConfig -from vectordb_bench.backend.clients.milvus.milvus import MILVUS_FORCE_MERGE_TARGET_SIZE_MB, Milvus +from vectordb_bench.backend.clients.api import IndexType, MetricType +from vectordb_bench.backend.clients.milvus.config import ( + HNSWBF16Config, + HNSWConfig, + MilvusConfig, + MilvusFtsConfig, +) +from vectordb_bench.backend.clients.milvus.milvus import ( + MILVUS_FORCE_MERGE_TARGET_SIZE_MB, + Milvus, + _float32_to_bf16_bytes, +) from vectordb_bench.backend.payload import PayloadProfile from vectordb_bench.backend.runner.mp_runner import MultiProcessingSearchRunner from vectordb_bench.interface import BenchMarkRunner @@ -24,6 +35,143 @@ log = logging.getLogger(__name__) +def test_float32_to_bf16_bytes_uses_round_to_nearest_even() -> None: + float32_bits = np.array( + [ + [0x3F807FFF, 0x3F808000, 0x3F808001, 0x3F818000], + [0x7F800000, 0xFF800000, 0x7FC00001, 0xFFC00001], + ], + dtype=np.uint32, + ) + + encoded = _float32_to_bf16_bytes(float32_bits.view(np.float32)) + actual = np.frombuffer(b"".join(encoded), dtype=" None: + config = HNSWBF16Config(M=30, efConstruction=360, ef=100, metric_type=MetricType.COSINE) + + assert config.index == IndexType.HNSW_BF16 + # Milvus has no HNSW_BF16 index type; BF16 is selected by the field's data type. + assert config.index_param() == { + "metric_type": "COSINE", + "index_type": "HNSW", + "params": {"M": 30, "efConstruction": 360}, + } + assert config.search_param() == {"metric_type": "COSINE", "params": {"ef": 100}} + + +def test_hnsw_bf16_index_type_resolves_to_bf16_case_config() -> None: + assert DB.Milvus.case_config_cls(IndexType.HNSW_BF16) is HNSWBF16Config + + +def _vector_field_schema_call( + monkeypatch: pytest.MonkeyPatch, + db_case_config: HNSWConfig, + dim: int = 4, +) -> call: + client = MagicMock() + client.has_collection.return_value = False + client_cls = MagicMock(return_value=client) + schema = MagicMock() + client_cls.create_schema.return_value = schema + client_cls.prepare_index_params.return_value = MagicMock() + monkeypatch.setattr("vectordb_bench.backend.clients.milvus.milvus.MilvusClient", client_cls) + + Milvus(dim=dim, db_config={"uri": "http://example.invalid"}, db_case_config=db_case_config) + + return next(c for c in schema.add_field.call_args_list if c.args[0] == "vector") + + +def test_milvus_bf16_case_config_creates_bfloat16_vector_field(monkeypatch: pytest.MonkeyPatch) -> None: + field_call = _vector_field_schema_call( + monkeypatch, + HNSWBF16Config(M=8, efConstruction=64, ef=32, metric_type=MetricType.COSINE), + ) + + assert field_call.args[1] == DataType.BFLOAT16_VECTOR + assert field_call.kwargs["dim"] == 4 + + +def test_milvus_fp32_case_config_still_creates_float_vector_field(monkeypatch: pytest.MonkeyPatch) -> None: + field_call = _vector_field_schema_call( + monkeypatch, + HNSWConfig(M=8, efConstruction=64, ef=32, metric_type=MetricType.COSINE), + ) + + assert field_call.args[1] == DataType.FLOAT_VECTOR + + +def test_milvus_fts_case_config_does_not_enable_bf16(monkeypatch: pytest.MonkeyPatch) -> None: + client = MagicMock() + client.has_collection.return_value = False + client_cls = MagicMock(return_value=client) + client_cls.create_schema.return_value = MagicMock() + client_cls.prepare_index_params.return_value = MagicMock() + monkeypatch.setattr("vectordb_bench.backend.clients.milvus.milvus.MilvusClient", client_cls) + + # An FTS config has no `index` attribute at all, so the BF16 probe must tolerate its absence. + db = Milvus(dim=4, db_config={"uri": "http://example.invalid"}, db_case_config=MilvusFtsConfig()) + + assert db._use_bf16 is False + + +def test_milvus_bf16_insert_converts_runner_batch_to_bf16_bytes() -> None: + client = MagicMock() + client.insert.side_effect = lambda _collection, rows: {"insert_count": len(rows)} + + db = object.__new__(Milvus) + db.client = client + db.collection_name = "test_collection" + db._primary_field = "pk" + db._scalar_id_field = "id" + db._vector_field = "vector" + db.with_scalar_labels = False + db._use_bf16 = True + + embeddings = [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]] + count, err = db.insert_embeddings(embeddings=embeddings, metadata=[0, 1, 2]) + + assert count == 3 + assert err is None + rows = client.insert.call_args.args[1] + assert [row["vector"] for row in rows] == _float32_to_bf16_bytes(embeddings) + # Two bytes per dimension, where a float32 payload would use four. + assert {len(row["vector"]) for row in rows} == {4} + + +def test_milvus_bf16_batch_search_converts_every_query() -> None: + captured = {} + + def search(**kwargs): + captured.update(kwargs) + return [[{"pk": 0}], [{"pk": 1}], [{"pk": 2}]] + + db = object.__new__(Milvus) + db.client = SimpleNamespace(search=search) + db.collection_name = "test_collection" + db._vector_field = "vector" + db._primary_field = "pk" + db.case_config = SimpleNamespace(search_param=lambda: {"metric_type": "COSINE"}) + db.expr = "" + db._use_bf16 = True + + queries = [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]] + + assert db.search_embeddings(queries, k=1) == [[0], [1], [2]] + # Guards against encoding only the first query of a batch. + assert captured["data"] == _float32_to_bf16_bytes(queries) + + def test_milvus_vector_payload_requests_vector_field_and_returns_ids(): captured = {} diff --git a/tests/test_milvus_zilliz_cli.py b/tests/test_milvus_zilliz_cli.py index d01628f10..e3f307c5a 100644 --- a/tests/test_milvus_zilliz_cli.py +++ b/tests/test_milvus_zilliz_cli.py @@ -2,7 +2,9 @@ from click.testing import CliRunner from pytest import MonkeyPatch +from vectordb_bench.backend.clients.api import IndexType from vectordb_bench.backend.clients.milvus import cli as milvus_cli +from vectordb_bench.backend.clients.milvus.config import HNSWBF16Config from vectordb_bench.backend.clients.zilliz_cloud import cli as zilliz_cli from vectordb_bench.cli import cli as common_cli @@ -124,6 +126,39 @@ def fake_run(**kwargs): assert captured["db_case_config"].force_merge_target_size_mb == 256 +def test_milvus_hnsw_bf16_cli_builds_bf16_case_config(monkeypatch: MonkeyPatch) -> None: + captured = {} + + def fake_run(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(milvus_cli, "run", fake_run) + + result = CliRunner().invoke( + milvus_cli.MilvusHNSWBF16, + [ + "--uri", + "http://localhost:19530", + "--collection-name", + "bench_bf16", + "--m", + "30", + "--ef-construction", + "360", + "--ef-search", + "100", + "--dry-run", + ], + ) + + assert result.exit_code == 0, result.output + case_config = captured["db_case_config"] + assert isinstance(case_config, HNSWBF16Config) + assert case_config.index == IndexType.HNSW_BF16 + assert (case_config.M, case_config.efConstruction, case_config.ef) == (30, 360, 100) + assert captured["db_config"].collection_name == "bench_bf16" + + @pytest.mark.parametrize( ("level_args", "expected_level"), [(["--level", "2"], 2), ([], None)], diff --git a/vectordb_bench/backend/clients/api.py b/vectordb_bench/backend/clients/api.py index e1c74e9df..862929a5d 100644 --- a/vectordb_bench/backend/clients/api.py +++ b/vectordb_bench/backend/clients/api.py @@ -25,6 +25,7 @@ class IndexType(StrEnum): HNSW_SQ = "HNSW_SQ" HNSW_BQ = "HNSW_BQ" HNSW_PQ = "HNSW_PQ" + HNSW_BF16 = "HNSW_BF16" HNSW_PRQ = "HNSW_PRQ" DISKANN = "DISKANN" STREAMING_DISKANN = "DISKANN" diff --git a/vectordb_bench/backend/clients/milvus/cli.py b/vectordb_bench/backend/clients/milvus/cli.py index 4df345ae7..0ceaecb47 100644 --- a/vectordb_bench/backend/clients/milvus/cli.py +++ b/vectordb_bench/backend/clients/milvus/cli.py @@ -206,6 +206,26 @@ def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]): ) +@cli.command() +@click_parameter_decorators_from_typed_dict(MilvusHNSWTypedDict) +def MilvusHNSWBF16(**parameters: Unpack[MilvusHNSWTypedDict]): + from .config import HNSWBF16Config + + run( + db=DBTYPE, + db_config=_build_milvus_config(parameters), + db_case_config=_apply_milvus_case_defaults( + HNSWBF16Config( + M=parameters["m"], + efConstruction=parameters["ef_construction"], + ef=parameters["ef_search"], + ), + parameters, + ), + **parameters, + ) + + class MilvusRefineTypedDict(TypedDict): refine: Annotated[ bool, diff --git a/vectordb_bench/backend/clients/milvus/config.py b/vectordb_bench/backend/clients/milvus/config.py index 1bc3b569c..edf797259 100644 --- a/vectordb_bench/backend/clients/milvus/config.py +++ b/vectordb_bench/backend/clients/milvus/config.py @@ -102,6 +102,17 @@ def search_param(self) -> dict: } +class HNSWBF16Config(HNSWConfig): + index: IndexType = IndexType.HNSW_BF16 + + def index_param(self) -> dict: + return { + "metric_type": self.parse_metric(), + "index_type": "HNSW", + "params": {"M": self.M, "efConstruction": self.efConstruction}, + } + + class HNSWSQConfig(HNSWConfig, DBCaseConfig): index: IndexType = IndexType.HNSW_SQ sq_type: SQType = SQType.SQ8 @@ -618,6 +629,7 @@ def search_param(self) -> dict: IndexType.AUTOINDEX: AutoIndexConfig, IndexType.FTS: MilvusFtsConfig, IndexType.HNSW: HNSWConfig, + IndexType.HNSW_BF16: HNSWBF16Config, IndexType.HNSW_SQ: HNSWSQConfig, IndexType.HNSW_PQ: HNSWPQConfig, IndexType.HNSW_PRQ: HNSWPRQConfig, diff --git a/vectordb_bench/backend/clients/milvus/milvus.py b/vectordb_bench/backend/clients/milvus/milvus.py index 4296e6740..8d2ae6be4 100644 --- a/vectordb_bench/backend/clients/milvus/milvus.py +++ b/vectordb_bench/backend/clients/milvus/milvus.py @@ -6,12 +6,13 @@ from contextlib import contextmanager from typing import Any +import numpy as np from pymilvus import DataType, Function, FunctionType, MilvusClient, MilvusException from vectordb_bench.backend.filter import Filter, FilterOp from vectordb_bench.backend.payload import PayloadProfile -from ..api import VectorDB +from ..api import IndexType, VectorDB from .config import MilvusFtsConfig, MilvusIndexConfig log = logging.getLogger(__name__) @@ -21,6 +22,19 @@ MILVUS_FORCE_MERGE_RETRY_INTERVAL_SECONDS = 30 +def _float32_to_bf16_bytes(vectors: Iterable[list[float]]) -> list[bytes]: + """Convert float32 vectors to round-to-nearest-even BF16 byte strings.""" + # pymilvus accepts raw BF16 bytes; keep this tested bit encoding to avoid an ml_dtypes dependency. + bits = np.asarray(vectors, dtype="> np.uint32(16)) & np.uint32(1)) + bf16 = ((bits + rounding_bias) >> np.uint32(16)).astype(" np.uint32(0x7F800000) + bf16[nan_mask] = ((bits[nan_mask] >> np.uint32(16)) | np.uint32(0x40)).astype(" bool: @@ -89,6 +104,8 @@ def __init__( # noqa: PLR0912, PLR0915 if self.with_scalar_labels: self._scalar_payload_label_field = "scalar_label" + self._use_bf16 = getattr(self.case_config, "index", None) == IndexType.HNSW_BF16 + client = MilvusClient( uri=self.db_config.get("uri"), user=self.db_config.get("user"), @@ -134,7 +151,11 @@ def __init__( # noqa: PLR0912, PLR0915 else: schema.add_field(self._primary_field, DataType.INT64, is_primary=True) schema.add_field(self._scalar_id_field, DataType.INT64) - schema.add_field(self._vector_field, DataType.FLOAT_VECTOR, dim=dim) + schema.add_field( + self._vector_field, + DataType.BFLOAT16_VECTOR if self._use_bf16 else DataType.FLOAT_VECTOR, + dim=dim, + ) if self.multitenant_tenant_labels: schema.add_field( @@ -429,12 +450,13 @@ def insert_embeddings( assert self.client is not None assert len(embeddings) == len(metadata) + vectors = _float32_to_bf16_bytes(embeddings) if self._use_bf16 else embeddings rows = [] for i in range(len(embeddings)): row = { self._primary_field: metadata[i], self._scalar_id_field: metadata[i], - self._vector_field: embeddings[i], + self._vector_field: vectors[i], } if tenant_labels_data is not None: row[self._multitenant_partition_key_field] = tenant_labels_data[i] @@ -566,9 +588,10 @@ def search_embeddings( tenant_expr = f"{tenant_field} == '{tenant}'" expr = tenant_expr if not expr else f"({expr}) and ({tenant_expr})" + search_data = _float32_to_bf16_bytes(queries) if self._use_bf16 else queries search_kwargs = { "collection_name": self.collection_name, - "data": queries, + "data": search_data, "anns_field": self._vector_field, "search_params": self.case_config.search_param(), "limit": k, diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index 1bbc462ef..49af7afd0 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -27,7 +27,7 @@ from ..backend.clients.lindorm.cli import LindormHNSW, LindormIVFBQ, LindormIVFPQ from ..backend.clients.mariadb.cli import MariaDBHNSW from ..backend.clients.memorydb.cli import MemoryDB -from ..backend.clients.milvus.cli import MilvusAutoIndex, MilvusFTS +from ..backend.clients.milvus.cli import MilvusAutoIndex, MilvusFTS, MilvusHNSWBF16 from ..backend.clients.oceanbase.cli import OceanBaseHNSW, OceanBaseIVF from ..backend.clients.oss_opensearch.cli import OSSOpenSearch from ..backend.clients.pgdiskann.cli import PgDiskAnn @@ -70,6 +70,7 @@ cli.add_command(ZillizAutoIndex) cli.add_command(MilvusAutoIndex) cli.add_command(MilvusFTS) +cli.add_command(MilvusHNSWBF16) cli.add_command(AWSOpenSearch) cli.add_command(OSSOpenSearch) cli.add_command(PgVectorScaleDiskAnn) From 11ff32b64e2bfcb9bd837dd318a2a4f7ecef7670 Mon Sep 17 00:00:00 2001 From: Alex Bondarev Date: Fri, 16 Oct 2026 10:44:09 -0500 Subject: [PATCH 2/5] refactor(milvus): select BF16 with vector_data_type Replace the HNSW_BF16 index type and the milvushnswbf16 command with a vector_data_type field on HNSWConfig. Only milvushnsw has the --vector-data-type option. Reject a reused collection if its vector field type or dimension does not match the configuration. --- README.md | 9 ++ tests/test_milvus.py | 108 +++++++++++++++--- tests/test_milvus_zilliz_cli.py | 35 ++++-- vectordb_bench/backend/clients/api.py | 1 - vectordb_bench/backend/clients/milvus/cli.py | 36 +++--- .../backend/clients/milvus/config.py | 19 ++- .../backend/clients/milvus/milvus.py | 50 +++++--- vectordb_bench/cli/vectordbbench.py | 3 +- 8 files changed, 185 insertions(+), 76 deletions(-) diff --git a/README.md b/README.md index 0ffa98f0e..42eba23f1 100644 --- a/README.md +++ b/README.md @@ -138,6 +138,15 @@ vectordbbench milvusautoindex \ ``` +`milvushnsw` accepts `--vector-data-type` (YAML: `vector_data_type`) to store vectors as `FLOAT32` (default) or `BFLOAT16`. The index type stays `HNSW`. With `--skip-drop-old` or `--skip-load`, the run fails if the vector field of the collection has a different data type or dimension. Results record the value in `task_config.db_case_config.vector_data_type`. + +```shell +vectordbbench milvushnsw \ + --uri http://localhost:19530 \ + --vector-data-type BFLOAT16 \ + +``` + Use `--nq` (default: `1`) to set the number of query vectors per concurrent search request. Milvus and Zilliz Cloud support batching for single-tenant vector performance cases, including filtered searches. This example reuses a collection containing the selected diff --git a/tests/test_milvus.py b/tests/test_milvus.py index 3d5f79dba..34d3c4685 100644 --- a/tests/test_milvus.py +++ b/tests/test_milvus.py @@ -17,10 +17,10 @@ from vectordb_bench.backend.clients import DB from vectordb_bench.backend.clients.api import IndexType, MetricType from vectordb_bench.backend.clients.milvus.config import ( - HNSWBF16Config, HNSWConfig, MilvusConfig, MilvusFtsConfig, + VectorDataType, ) from vectordb_bench.backend.clients.milvus.milvus import ( MILVUS_FORCE_MERGE_TARGET_SIZE_MB, @@ -57,21 +57,15 @@ def test_float32_to_bf16_bytes_uses_round_to_nearest_even() -> None: np.testing.assert_array_equal(actual, expected) -def test_hnsw_bf16_config_requests_plain_hnsw_index_type() -> None: - config = HNSWBF16Config(M=30, efConstruction=360, ef=100, metric_type=MetricType.COSINE) +def test_hnsw_config_vector_data_type_does_not_change_index_params() -> None: + fp32 = HNSWConfig(M=30, efConstruction=360, ef=100, metric_type=MetricType.COSINE) + bf16 = HNSWConfig(M=30, efConstruction=360, ef=100, metric_type=MetricType.COSINE, vector_data_type="BFLOAT16") - assert config.index == IndexType.HNSW_BF16 - # Milvus has no HNSW_BF16 index type; BF16 is selected by the field's data type. - assert config.index_param() == { - "metric_type": "COSINE", - "index_type": "HNSW", - "params": {"M": 30, "efConstruction": 360}, - } - assert config.search_param() == {"metric_type": "COSINE", "params": {"ef": 100}} - - -def test_hnsw_bf16_index_type_resolves_to_bf16_case_config() -> None: - assert DB.Milvus.case_config_cls(IndexType.HNSW_BF16) is HNSWBF16Config + assert fp32.vector_data_type == VectorDataType.FLOAT32 + assert bf16.vector_data_type == VectorDataType.BFLOAT16 + assert bf16.index == IndexType.HNSW + assert bf16.index_param() == fp32.index_param() + assert bf16.search_param() == fp32.search_param() def _vector_field_schema_call( @@ -95,7 +89,13 @@ def _vector_field_schema_call( def test_milvus_bf16_case_config_creates_bfloat16_vector_field(monkeypatch: pytest.MonkeyPatch) -> None: field_call = _vector_field_schema_call( monkeypatch, - HNSWBF16Config(M=8, efConstruction=64, ef=32, metric_type=MetricType.COSINE), + HNSWConfig( + M=8, + efConstruction=64, + ef=32, + metric_type=MetricType.COSINE, + vector_data_type=VectorDataType.BFLOAT16, + ), ) assert field_call.args[1] == DataType.BFLOAT16_VECTOR @@ -119,7 +119,7 @@ def test_milvus_fts_case_config_does_not_enable_bf16(monkeypatch: pytest.MonkeyP client_cls.prepare_index_params.return_value = MagicMock() monkeypatch.setattr("vectordb_bench.backend.clients.milvus.milvus.MilvusClient", client_cls) - # An FTS config has no `index` attribute at all, so the BF16 probe must tolerate its absence. + # An FTS config has no `vector_data_type` attribute, so the BF16 probe must tolerate its absence. db = Milvus(dim=4, db_config={"uri": "http://example.invalid"}, db_case_config=MilvusFtsConfig()) assert db._use_bf16 is False @@ -235,10 +235,10 @@ def perf_counter(): assert [c.kwargs["data"] for c in db.client.search.call_args_list] == expected_batches -def _fake_milvus_client(monkeypatch, *, collection_exists=False, properties=None): +def _fake_milvus_client(monkeypatch, *, collection_exists=False, properties=None, fields=None): client = MagicMock() client.has_collection.return_value = collection_exists - client.describe_collection.return_value = {"properties": properties or {}} + client.describe_collection.return_value = {"properties": properties or {}, "fields": fields or []} client_cls = MagicMock(return_value=client) client_cls.create_schema.return_value = MagicMock() client_cls.prepare_index_params.return_value = MagicMock() @@ -276,6 +276,76 @@ def test_milvus_rejects_existing_collection_with_incompatible_properties(monkeyp _create_milvus_with_collection_properties(monkeypatch, collection_exists=True) +def _vector_field(data_type: DataType, dim: int = 4) -> dict: + return {"name": "vector", "type": data_type, "params": {"dim": dim}} + + +def _reuse_hnsw_collection(vector_data_type: VectorDataType, dim: int = 4) -> Milvus: + return Milvus( + dim=dim, + db_config={"uri": "http://example.invalid"}, + db_case_config=HNSWConfig( + M=8, + efConstruction=64, + ef=32, + metric_type=MetricType.COSINE, + vector_data_type=vector_data_type, + ), + ) + + +def test_milvus_reuses_existing_collection_with_matching_vector_field(monkeypatch): + client = _fake_milvus_client( + monkeypatch, + collection_exists=True, + fields=[_vector_field(DataType.BFLOAT16_VECTOR)], + ) + + _reuse_hnsw_collection(VectorDataType.BFLOAT16) + + client.create_collection.assert_not_called() + + +@pytest.mark.parametrize( + ("fields", "vector_data_type", "match"), + [ + ( + [_vector_field(DataType.FLOAT_VECTOR)], + VectorDataType.BFLOAT16, + r"expected BFLOAT16_VECTOR \(dim=4\), got FLOAT_VECTOR \(dim=4\)", + ), + ( + [_vector_field(DataType.BFLOAT16_VECTOR)], + VectorDataType.FLOAT32, + r"expected FLOAT_VECTOR \(dim=4\), got BFLOAT16_VECTOR \(dim=4\)", + ), + ( + [_vector_field(DataType.FLOAT_VECTOR, dim=8)], + VectorDataType.FLOAT32, + r"expected FLOAT_VECTOR \(dim=4\), got FLOAT_VECTOR \(dim=8\)", + ), + ([], VectorDataType.FLOAT32, "has no vector field 'vector'"), + ], +) +def test_milvus_rejects_existing_collection_with_incompatible_vector_field( + monkeypatch, fields, vector_data_type, match +): + client = _fake_milvus_client(monkeypatch, collection_exists=True, fields=fields) + + with pytest.raises(ValueError, match=match): + _reuse_hnsw_collection(vector_data_type) + + client.close.assert_called_once() + + +def test_milvus_fts_reuse_skips_vector_field_check(monkeypatch): + client = _fake_milvus_client(monkeypatch, collection_exists=True) + + Milvus(dim=4, db_config={"uri": "http://example.invalid"}, db_case_config=MilvusFtsConfig()) + + client.create_collection.assert_not_called() + + def test_milvus_force_merge_config_defaults_and_validation(): from vectordb_bench.backend.clients.milvus.config import AutoIndexConfig, MilvusFtsConfig diff --git a/tests/test_milvus_zilliz_cli.py b/tests/test_milvus_zilliz_cli.py index e3f307c5a..e88f0c40e 100644 --- a/tests/test_milvus_zilliz_cli.py +++ b/tests/test_milvus_zilliz_cli.py @@ -4,7 +4,7 @@ from vectordb_bench.backend.clients.api import IndexType from vectordb_bench.backend.clients.milvus import cli as milvus_cli -from vectordb_bench.backend.clients.milvus.config import HNSWBF16Config +from vectordb_bench.backend.clients.milvus.config import HNSWConfig, VectorDataType from vectordb_bench.backend.clients.zilliz_cloud import cli as zilliz_cli from vectordb_bench.cli import cli as common_cli @@ -126,7 +126,15 @@ def fake_run(**kwargs): assert captured["db_case_config"].force_merge_target_size_mb == 256 -def test_milvus_hnsw_bf16_cli_builds_bf16_case_config(monkeypatch: MonkeyPatch) -> None: +@pytest.mark.parametrize( + ("vector_data_type_args", "expected"), + [([], VectorDataType.FLOAT32), (["--vector-data-type", "bfloat16"], VectorDataType.BFLOAT16)], +) +def test_milvus_hnsw_cli_sets_vector_data_type( + monkeypatch: MonkeyPatch, + vector_data_type_args: list[str], + expected: VectorDataType, +) -> None: captured = {} def fake_run(**kwargs): @@ -135,28 +143,39 @@ def fake_run(**kwargs): monkeypatch.setattr(milvus_cli, "run", fake_run) result = CliRunner().invoke( - milvus_cli.MilvusHNSWBF16, + milvus_cli.MilvusHNSW, [ "--uri", "http://localhost:19530", - "--collection-name", - "bench_bf16", "--m", "30", "--ef-construction", "360", "--ef-search", "100", + *vector_data_type_args, "--dry-run", ], ) assert result.exit_code == 0, result.output case_config = captured["db_case_config"] - assert isinstance(case_config, HNSWBF16Config) - assert case_config.index == IndexType.HNSW_BF16 + assert isinstance(case_config, HNSWConfig) + assert case_config.index == IndexType.HNSW assert (case_config.M, case_config.efConstruction, case_config.ef) == (30, 360, 100) - assert captured["db_config"].collection_name == "bench_bf16" + assert case_config.vector_data_type == expected + + +@pytest.mark.parametrize("command", [milvus_cli.MilvusHNSWSQ, milvus_cli.MilvusHNSWPQ, milvus_cli.MilvusHNSWPRQ]) +def test_milvus_quantized_hnsw_clis_reject_vector_data_type(command) -> None: + result = CliRunner().invoke( + command, + ["--uri", "http://localhost:19530", "--vector-data-type", "BFLOAT16", "--dry-run"], + ) + + assert result.exit_code == 2 + assert "No such option" in result.output + assert "--vector-data-type" in result.output @pytest.mark.parametrize( diff --git a/vectordb_bench/backend/clients/api.py b/vectordb_bench/backend/clients/api.py index 862929a5d..e1c74e9df 100644 --- a/vectordb_bench/backend/clients/api.py +++ b/vectordb_bench/backend/clients/api.py @@ -25,7 +25,6 @@ class IndexType(StrEnum): HNSW_SQ = "HNSW_SQ" HNSW_BQ = "HNSW_BQ" HNSW_PQ = "HNSW_PQ" - HNSW_BF16 = "HNSW_BF16" HNSW_PRQ = "HNSW_PRQ" DISKANN = "DISKANN" STREAMING_DISKANN = "DISKANN" diff --git a/vectordb_bench/backend/clients/milvus/cli.py b/vectordb_bench/backend/clients/milvus/cli.py index 0ceaecb47..46db50167 100644 --- a/vectordb_bench/backend/clients/milvus/cli.py +++ b/vectordb_bench/backend/clients/milvus/cli.py @@ -186,39 +186,33 @@ def MilvusFlat(**parameters: Unpack[MilvusFlatTypedDict]): class MilvusHNSWTypedDict(CommonTypedDict, MilvusTypedDict, HNSWFlavor3): ... -@cli.command() -@click_parameter_decorators_from_typed_dict(MilvusHNSWTypedDict) -def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]): - from .config import HNSWConfig - - run( - db=DBTYPE, - db_config=_build_milvus_config(parameters), - db_case_config=_apply_milvus_case_defaults( - HNSWConfig( - M=parameters["m"], - efConstruction=parameters["ef_construction"], - ef=parameters["ef_search"], - ), - parameters, +class MilvusHNSWCommandTypedDict(MilvusHNSWTypedDict): + vector_data_type: Annotated[ + str, + click.option( + "--vector-data-type", + type=click.Choice(["FLOAT32", "BFLOAT16"], case_sensitive=False), + default="FLOAT32", + show_default=True, + help="Data type of the stored vectors. Supported values: FLOAT32,BFLOAT16", ), - **parameters, - ) + ] @cli.command() -@click_parameter_decorators_from_typed_dict(MilvusHNSWTypedDict) -def MilvusHNSWBF16(**parameters: Unpack[MilvusHNSWTypedDict]): - from .config import HNSWBF16Config +@click_parameter_decorators_from_typed_dict(MilvusHNSWCommandTypedDict) +def MilvusHNSW(**parameters: Unpack[MilvusHNSWCommandTypedDict]): + from .config import HNSWConfig run( db=DBTYPE, db_config=_build_milvus_config(parameters), db_case_config=_apply_milvus_case_defaults( - HNSWBF16Config( + HNSWConfig( M=parameters["m"], efConstruction=parameters["ef_construction"], ef=parameters["ef_search"], + vector_data_type=parameters["vector_data_type"], ), parameters, ), diff --git a/vectordb_bench/backend/clients/milvus/config.py b/vectordb_bench/backend/clients/milvus/config.py index edf797259..59f147635 100644 --- a/vectordb_bench/backend/clients/milvus/config.py +++ b/vectordb_bench/backend/clients/milvus/config.py @@ -1,3 +1,4 @@ +from enum import StrEnum from typing import Annotated, ClassVar from pydantic import BaseModel, Field, SecretStr, field_validator @@ -82,11 +83,17 @@ def search_param(self) -> dict: return params +class VectorDataType(StrEnum): + FLOAT32 = "FLOAT32" + BFLOAT16 = "BFLOAT16" + + class HNSWConfig(MilvusIndexConfig, DBCaseConfig): M: int efConstruction: int ef: int | None = None index: IndexType = IndexType.HNSW + vector_data_type: VectorDataType = VectorDataType.FLOAT32 def index_param(self) -> dict: return { @@ -102,17 +109,6 @@ def search_param(self) -> dict: } -class HNSWBF16Config(HNSWConfig): - index: IndexType = IndexType.HNSW_BF16 - - def index_param(self) -> dict: - return { - "metric_type": self.parse_metric(), - "index_type": "HNSW", - "params": {"M": self.M, "efConstruction": self.efConstruction}, - } - - class HNSWSQConfig(HNSWConfig, DBCaseConfig): index: IndexType = IndexType.HNSW_SQ sq_type: SQType = SQType.SQ8 @@ -629,7 +625,6 @@ def search_param(self) -> dict: IndexType.AUTOINDEX: AutoIndexConfig, IndexType.FTS: MilvusFtsConfig, IndexType.HNSW: HNSWConfig, - IndexType.HNSW_BF16: HNSWBF16Config, IndexType.HNSW_SQ: HNSWSQConfig, IndexType.HNSW_PQ: HNSWPQConfig, IndexType.HNSW_PRQ: HNSWPRQConfig, diff --git a/vectordb_bench/backend/clients/milvus/milvus.py b/vectordb_bench/backend/clients/milvus/milvus.py index 8d2ae6be4..bacc04b3e 100644 --- a/vectordb_bench/backend/clients/milvus/milvus.py +++ b/vectordb_bench/backend/clients/milvus/milvus.py @@ -12,8 +12,8 @@ from vectordb_bench.backend.filter import Filter, FilterOp from vectordb_bench.backend.payload import PayloadProfile -from ..api import IndexType, VectorDB -from .config import MilvusFtsConfig, MilvusIndexConfig +from ..api import VectorDB +from .config import MilvusFtsConfig, MilvusIndexConfig, VectorDataType log = logging.getLogger(__name__) @@ -51,7 +51,7 @@ def supports_full_text_search(cls) -> bool: def has_text_field(self) -> bool: return bool(getattr(self, "_is_fts", False) and getattr(self, "_text_field", None)) - def __init__( # noqa: PLR0912, PLR0915 + def __init__( # noqa: PLR0915 self, dim: int, db_config: dict, @@ -104,7 +104,7 @@ def __init__( # noqa: PLR0912, PLR0915 if self.with_scalar_labels: self._scalar_payload_label_field = "scalar_label" - self._use_bf16 = getattr(self.case_config, "index", None) == IndexType.HNSW_BF16 + self._use_bf16 = getattr(self.case_config, "vector_data_type", None) == VectorDataType.BFLOAT16 client = MilvusClient( uri=self.db_config.get("uri"), @@ -197,18 +197,42 @@ def __init__( # noqa: PLR0912, PLR0915 self.collection_name, replica_number=self.db_config.get("replica_number", 1), ) - elif collection_properties: - actual_properties = client.describe_collection(self.collection_name).get("properties") or {} - if any(actual_properties.get(key) != value for key, value in collection_properties.items()): - client.close() - msg = ( - f"{self.name} collection {self.collection_name} has incompatible collection properties: " - f"expected {collection_properties}, got {actual_properties}. Drop and recreate the collection." - ) - raise ValueError(msg) + else: + self._validate_existing_collection(client, dim, collection_properties) client.close() + def _validate_existing_collection(self, client: MilvusClient, dim: int, collection_properties: dict) -> None: + desc = client.describe_collection(self.collection_name) + actual_properties = desc.get("properties") or {} + if any(actual_properties.get(key) != value for key, value in collection_properties.items()): + client.close() + msg = ( + f"{self.name} collection {self.collection_name} has incompatible collection properties: " + f"expected {collection_properties}, got {actual_properties}. Drop and recreate the collection." + ) + raise ValueError(msg) + if self._is_fts: + return + + expected_type = DataType.BFLOAT16_VECTOR if self._use_bf16 else DataType.FLOAT_VECTOR + fields_by_name = {self._field_property(field, "name"): field for field in desc.get("fields", [])} + vector_field = fields_by_name.get(self._vector_field) + if vector_field is None: + problem = f"has no vector field '{self._vector_field}'" + else: + actual_type = DataType(self._field_property(vector_field, "type")) + actual_dim = self._field_property(vector_field, "dim") + if actual_type == expected_type and actual_dim == dim: + return + problem = ( + f"has incompatible vector field '{self._vector_field}': " + f"expected {expected_type.name} (dim={dim}), got {actual_type.name} (dim={actual_dim})" + ) + client.close() + msg = f"{self.name} collection {self.collection_name} {problem}. Drop and recreate the collection." + raise ValueError(msg) + def _build_index_params(self): index_params = MilvusClient.prepare_index_params() if self._is_fts: diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index 49af7afd0..1bbc462ef 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -27,7 +27,7 @@ from ..backend.clients.lindorm.cli import LindormHNSW, LindormIVFBQ, LindormIVFPQ from ..backend.clients.mariadb.cli import MariaDBHNSW from ..backend.clients.memorydb.cli import MemoryDB -from ..backend.clients.milvus.cli import MilvusAutoIndex, MilvusFTS, MilvusHNSWBF16 +from ..backend.clients.milvus.cli import MilvusAutoIndex, MilvusFTS from ..backend.clients.oceanbase.cli import OceanBaseHNSW, OceanBaseIVF from ..backend.clients.oss_opensearch.cli import OSSOpenSearch from ..backend.clients.pgdiskann.cli import PgDiskAnn @@ -70,7 +70,6 @@ cli.add_command(ZillizAutoIndex) cli.add_command(MilvusAutoIndex) cli.add_command(MilvusFTS) -cli.add_command(MilvusHNSWBF16) cli.add_command(AWSOpenSearch) cli.add_command(OSSOpenSearch) cli.add_command(PgVectorScaleDiskAnn) From d91c90b4ec03148dc76489e52eeb035017dbef61 Mon Sep 17 00:00:00 2001 From: Alex Bondarev Date: Mon, 19 Oct 2026 09:55:57 -0500 Subject: [PATCH 3/5] style(vespa): import datetime as dt ruff reports ICN001 for this import, and make lint fails on main. --- vectordb_bench/backend/clients/vespa/vespa.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vectordb_bench/backend/clients/vespa/vespa.py b/vectordb_bench/backend/clients/vespa/vespa.py index 12fa65d5e..99a4e46c4 100644 --- a/vectordb_bench/backend/clients/vespa/vespa.py +++ b/vectordb_bench/backend/clients/vespa/vespa.py @@ -1,4 +1,4 @@ -import datetime +import datetime as dt import json import logging import math @@ -490,7 +490,7 @@ def _create_application_package(self): ValidationID, ) - tomorrow = datetime.date.today() + datetime.timedelta(days=1) + tomorrow = dt.date.today() + dt.timedelta(days=1) if self._is_fts: fields = [ From 6f1796e692dac96bc18cae5040b4db0a890f1030 Mon Sep 17 00:00:00 2001 From: Alex Bondarev Date: Mon, 19 Oct 2026 09:55:57 -0500 Subject: [PATCH 4/5] fix(milvus): use BF16 width in payload estimates The vector payload estimate used 4 bytes per dimension for all runs. HNSWConfig now gives 2 bytes per dimension for BFLOAT16. The task runner passes this width to the estimate. FLOAT32 stays at 4 bytes. --- tests/test_milvus.py | 51 ++++++++++++++++++- vectordb_bench/backend/cases.py | 12 +++-- .../backend/clients/milvus/config.py | 4 ++ vectordb_bench/backend/payload.py | 6 +-- vectordb_bench/backend/task_runner.py | 12 +++-- 5 files changed, 73 insertions(+), 12 deletions(-) diff --git a/tests/test_milvus.py b/tests/test_milvus.py index 34d3c4685..c3b9643b8 100644 --- a/tests/test_milvus.py +++ b/tests/test_milvus.py @@ -27,10 +27,12 @@ Milvus, _float32_to_bf16_bytes, ) +from vectordb_bench.backend.data_source import DatasetSource from vectordb_bench.backend.payload import PayloadProfile from vectordb_bench.backend.runner.mp_runner import MultiProcessingSearchRunner +from vectordb_bench.backend.task_runner import CaseRunner, RunningStatus from vectordb_bench.interface import BenchMarkRunner -from vectordb_bench.models import CaseConfig, TaskConfig +from vectordb_bench.models import CaseConfig, TaskConfig, TaskStage log = logging.getLogger(__name__) @@ -172,6 +174,53 @@ def search(**kwargs): assert captured["data"] == _float32_to_bf16_bytes(queries) +@pytest.mark.parametrize( + ("vector_data_type", "expected_bytes"), + [(VectorDataType.FLOAT32, 30_920), (VectorDataType.BFLOAT16, 15_560)], +) +def test_milvus_vector_payload_estimate_uses_vector_data_type_width( + monkeypatch: pytest.MonkeyPatch, + vector_data_type: VectorDataType, + expected_bytes: int, +) -> None: + case_config = CaseConfig(case_id=CaseType.Performance768D1M, k=10, payload_profile=PayloadProfile.VECTOR) + runner = CaseRunner( + run_id="bf16-payload", + config=TaskConfig( + db=DB.Milvus, + db_config=DB.Milvus.config_cls(), + db_case_config=HNSWConfig(M=16, efConstruction=100, ef=64, vector_data_type=vector_data_type), + case_config=case_config, + stages=[TaskStage.SEARCH_CONCURRENT, TaskStage.SEARCH_SERIAL], + ), + ca=case_config.case, + status=RunningStatus.PENDING, + dataset_source=DatasetSource.S3, + ) + + class SerialRunner: + def run(self): + return (0.9, 0.9, 0.01, 0.01, 0.01, {10: 0.9}), 0.1 + + class ConcurrentRunner: + def run(self): + return 100.0, [1], [100.0], [0.01], [0.01], [0.01], [0.01] + + def stop(self): + return None + + def fake_init_search_runners(self): + self.serial_search_runner = SerialRunner() + self.search_runner = ConcurrentRunner() + + monkeypatch.setattr(CaseRunner, "_init_search_runners", fake_init_search_runners) + + metrics = runner._run_perf_case(drop_old=False) + + assert metrics.payload_profile == PayloadProfile.VECTOR.value + assert metrics.payload_estimated_bytes_per_query == expected_bytes + + def test_milvus_vector_payload_requests_vector_field_and_returns_ids(): captured = {} diff --git a/vectordb_bench/backend/cases.py b/vectordb_bench/backend/cases.py index 67c6a9c8c..4c92d8248 100644 --- a/vectordb_bench/backend/cases.py +++ b/vectordb_bench/backend/cases.py @@ -138,10 +138,14 @@ class Case(BaseModel): def filters(self) -> Filter: return non_filter - def estimated_payload_bytes_per_query(self, k: int | None) -> int: + def estimated_payload_bytes_per_query(self, k: int | None, vector_bytes_per_dim: int = 4) -> int: if k is None: k = config.K_DEFAULT - return self.payload_profile.estimated_bytes_per_query(k=k, dim=self.dataset.data.dim) + return self.payload_profile.estimated_bytes_per_query( + k=k, + dim=self.dataset.data.dim, + vector_bytes_per_dim=vector_bytes_per_dim, + ) @property def is_multitenant(self) -> bool: @@ -925,10 +929,10 @@ def filters(self) -> Filter: int_value = int(self.dataset.data.size * self.filter_rate) return NewIntFilter(filter_rate=self.filter_rate, int_field=FTS_FILTER_ID_FIELD, int_value=int_value) - def estimated_payload_bytes_per_query(self, k: int | None) -> int: + def estimated_payload_bytes_per_query(self, k: int | None, vector_bytes_per_dim: int = 4) -> int: if k is None: k = config.K_DEFAULT - return self.payload_profile.estimated_bytes_per_query(k=k, dim=0) + return self.payload_profile.estimated_bytes_per_query(k=k, dim=0, vector_bytes_per_dim=vector_bytes_per_dim) class FTSBm25Performance(FtsPerformanceCase): diff --git a/vectordb_bench/backend/clients/milvus/config.py b/vectordb_bench/backend/clients/milvus/config.py index 59f147635..09fcb61a9 100644 --- a/vectordb_bench/backend/clients/milvus/config.py +++ b/vectordb_bench/backend/clients/milvus/config.py @@ -95,6 +95,10 @@ class HNSWConfig(MilvusIndexConfig, DBCaseConfig): index: IndexType = IndexType.HNSW vector_data_type: VectorDataType = VectorDataType.FLOAT32 + @property + def vector_bytes_per_dim(self) -> int: + return 2 if self.vector_data_type == VectorDataType.BFLOAT16 else 4 + def index_param(self) -> dict: return { "metric_type": self.parse_metric(), diff --git a/vectordb_bench/backend/payload.py b/vectordb_bench/backend/payload.py index 04cba3b85..d8111b1f4 100644 --- a/vectordb_bench/backend/payload.py +++ b/vectordb_bench/backend/payload.py @@ -7,15 +7,15 @@ class PayloadProfile(StrEnum): SCALAR_LABEL = "scalar_label" TEXT = "text" - def estimated_bytes_per_query(self, *, k: int, dim: int) -> int: + def estimated_bytes_per_query(self, *, k: int, dim: int, vector_bytes_per_dim: int = 4) -> int: # Approximate payload size used for cloud leaderboard cost expansion. - # ID + distance is about 20 bytes per hit; vector is float32. + # ID + distance is about 20 bytes per hit; vectors default to float32 width. id_distance_bytes = 20 scalar_label_bytes = 16 if self == PayloadProfile.IDS_ONLY: return k * id_distance_bytes if self == PayloadProfile.VECTOR: - return k * (id_distance_bytes + dim * 4) + return k * (id_distance_bytes + dim * vector_bytes_per_dim) if self == PayloadProfile.SCALAR_LABEL: return k * (id_distance_bytes + scalar_label_bytes) if self == PayloadProfile.TEXT: diff --git a/vectordb_bench/backend/task_runner.py b/vectordb_bench/backend/task_runner.py index 6605b4857..280bacb46 100644 --- a/vectordb_bench/backend/task_runner.py +++ b/vectordb_bench/backend/task_runner.py @@ -367,6 +367,12 @@ def _run_capacity_case(self) -> Metric: log.info(f"Capacity case loading dataset reaches VectorDB's limit: max capacity = {count}") return Metric(max_load_count=count) + def _estimated_payload_bytes_per_query(self) -> int: + return self.ca.estimated_payload_bytes_per_query( + self.config.case_config.k, + vector_bytes_per_dim=getattr(self.config.db_case_config, "vector_bytes_per_dim", 4), + ) + def _run_perf_case(self, drop_old: bool = True) -> Metric: """run performance cases @@ -447,9 +453,7 @@ def _run_perf_case(self, drop_old: bool = True) -> Metric: } if hasattr(self.ca, "payload_profile"): m.payload_profile = self.ca.payload_profile.value - m.payload_estimated_bytes_per_query = self.ca.estimated_payload_bytes_per_query( - self.config.case_config.k - ) + m.payload_estimated_bytes_per_query = self._estimated_payload_bytes_per_query() except Exception as e: log.warning(f"Failed to run performance case, reason = {e}") traceback.print_exc() @@ -560,7 +564,7 @@ def _run_cloud_cold_latency_case(self, drop_old: bool = True) -> Metric: "cold_latency": self.cold_warm_search_runner.run(), } m.payload_profile = self.ca.payload_profile.value - m.payload_estimated_bytes_per_query = self.ca.estimated_payload_bytes_per_query(self.config.case_config.k) + m.payload_estimated_bytes_per_query = self._estimated_payload_bytes_per_query() except Exception as e: log.warning(f"Failed to run cloud cold latency case, reason = {e}") traceback.print_exc() From 644751027558e490dcedc79f6c7d1a62faea353f Mon Sep 17 00:00:00 2001 From: Alex Bondarev Date: Mon, 19 Oct 2026 11:06:20 -0500 Subject: [PATCH 5/5] test(milvus): cover BF16 payload width paths Test the cold-latency estimate, the default width for other configs, the FTS case and the NaN branch of the BF16 encoder. VectorDataType now gives the width of each data type, with a shared FLOAT32 default in payload.py. The README now names --skip-drop-old --skip-load as the reuse flags. --- README.md | 2 +- tests/test_milvus.py | 84 ++++++++++++------- vectordb_bench/backend/cases.py | 10 ++- .../backend/clients/milvus/config.py | 8 +- .../backend/clients/milvus/milvus.py | 2 +- vectordb_bench/backend/payload.py | 4 +- vectordb_bench/backend/task_runner.py | 3 +- 7 files changed, 74 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 42eba23f1..e92afbeee 100644 --- a/README.md +++ b/README.md @@ -138,7 +138,7 @@ vectordbbench milvusautoindex \ ``` -`milvushnsw` accepts `--vector-data-type` (YAML: `vector_data_type`) to store vectors as `FLOAT32` (default) or `BFLOAT16`. The index type stays `HNSW`. With `--skip-drop-old` or `--skip-load`, the run fails if the vector field of the collection has a different data type or dimension. Results record the value in `task_config.db_case_config.vector_data_type`. +`milvushnsw` accepts `--vector-data-type` (YAML: `vector_data_type`) to store vectors as `FLOAT32` (default) or `BFLOAT16`. The index type stays `HNSW`. With `--skip-drop-old --skip-load`, the run fails if the vector field of the collection has a different data type or dimension. Results record the value in `task_config.db_case_config.vector_data_type`. ```shell vectordbbench milvushnsw \ diff --git a/tests/test_milvus.py b/tests/test_milvus.py index c3b9643b8..95a6bc92b 100644 --- a/tests/test_milvus.py +++ b/tests/test_milvus.py @@ -15,8 +15,9 @@ from vectordb_bench.backend.cases import CaseType from vectordb_bench.backend.clients import DB -from vectordb_bench.backend.clients.api import IndexType, MetricType +from vectordb_bench.backend.clients.api import DBCaseConfig, IndexType, MetricType from vectordb_bench.backend.clients.milvus.config import ( + AutoIndexConfig, HNSWConfig, MilvusConfig, MilvusFtsConfig, @@ -32,7 +33,7 @@ from vectordb_bench.backend.runner.mp_runner import MultiProcessingSearchRunner from vectordb_bench.backend.task_runner import CaseRunner, RunningStatus from vectordb_bench.interface import BenchMarkRunner -from vectordb_bench.models import CaseConfig, TaskConfig, TaskStage +from vectordb_bench.models import CaseConfig, TaskConfig log = logging.getLogger(__name__) @@ -42,6 +43,7 @@ def test_float32_to_bf16_bytes_uses_round_to_nearest_even() -> None: [ [0x3F807FFF, 0x3F808000, 0x3F808001, 0x3F818000], [0x7F800000, 0xFF800000, 0x7FC00001, 0xFFC00001], + [0x7F800001, 0x7FFFFFFF, 0xFFFFFFFF, 0xFF800001], ], dtype=np.uint32, ) @@ -53,6 +55,7 @@ def test_float32_to_bf16_bytes_uses_round_to_nearest_even() -> None: [ [0x3F80, 0x3F80, 0x3F81, 0x3F82], [0x7F80, 0xFF80, 0x7FC0, 0xFFC0], + [0x7FC0, 0x7FFF, 0xFFFF, 0xFFC0], ], dtype=np.uint16, ) @@ -174,51 +177,70 @@ def search(**kwargs): assert captured["data"] == _float32_to_bf16_bytes(queries) -@pytest.mark.parametrize( - ("vector_data_type", "expected_bytes"), - [(VectorDataType.FLOAT32, 30_920), (VectorDataType.BFLOAT16, 15_560)], -) -def test_milvus_vector_payload_estimate_uses_vector_data_type_width( - monkeypatch: pytest.MonkeyPatch, - vector_data_type: VectorDataType, - expected_bytes: int, -) -> None: - case_config = CaseConfig(case_id=CaseType.Performance768D1M, k=10, payload_profile=PayloadProfile.VECTOR) - runner = CaseRunner( - run_id="bf16-payload", +def _case_runner(db_case_config: DBCaseConfig, case_config: CaseConfig) -> CaseRunner: + return CaseRunner( + run_id="payload-estimate", config=TaskConfig( db=DB.Milvus, db_config=DB.Milvus.config_cls(), - db_case_config=HNSWConfig(M=16, efConstruction=100, ef=64, vector_data_type=vector_data_type), + db_case_config=db_case_config, case_config=case_config, - stages=[TaskStage.SEARCH_CONCURRENT, TaskStage.SEARCH_SERIAL], + stages=[], ), ca=case_config.case, status=RunningStatus.PENDING, dataset_source=DatasetSource.S3, ) - class SerialRunner: - def run(self): - return (0.9, 0.9, 0.01, 0.01, 0.01, {10: 0.9}), 0.1 - class ConcurrentRunner: - def run(self): - return 100.0, [1], [100.0], [0.01], [0.01], [0.01], [0.01] +@pytest.mark.parametrize( + ("db_case_config", "expected_bytes"), + [ + (HNSWConfig(M=16, efConstruction=100, ef=64), 30_920), + (HNSWConfig(M=16, efConstruction=100, ef=64, vector_data_type=VectorDataType.BFLOAT16), 15_560), + (AutoIndexConfig(), 30_920), + ], +) +def test_milvus_vector_payload_estimate_uses_vector_data_type_width( + db_case_config: DBCaseConfig, + expected_bytes: int, +) -> None: + case_config = CaseConfig(case_id=CaseType.Performance768D1M, k=10, payload_profile=PayloadProfile.VECTOR) + + metrics = _case_runner(db_case_config, case_config)._run_perf_case(drop_old=False) + + assert metrics.payload_profile == PayloadProfile.VECTOR.value + assert metrics.payload_estimated_bytes_per_query == expected_bytes - def stop(self): - return None - def fake_init_search_runners(self): - self.serial_search_runner = SerialRunner() - self.search_runner = ConcurrentRunner() +@pytest.mark.parametrize( + ("vector_data_type", "expected_bytes"), + [(VectorDataType.FLOAT32, 309_200), (VectorDataType.BFLOAT16, 155_600)], +) +def test_milvus_cold_latency_payload_estimate_uses_vector_data_type_width( + monkeypatch: pytest.MonkeyPatch, + vector_data_type: VectorDataType, + expected_bytes: int, +) -> None: + runner = _case_runner( + HNSWConfig(M=16, efConstruction=100, ef=64, vector_data_type=vector_data_type), + CaseConfig(case_id=CaseType.CloudColdLatencyCase, custom_case={"query_count": 1, "payload_profile": "vector"}), + ) + runner.db = MagicMock() + monkeypatch.setattr( + "vectordb_bench.backend.task_runner.ColdWarmSearchRunner", + lambda **_kwargs: SimpleNamespace(run=lambda: {}), + ) - monkeypatch.setattr(CaseRunner, "_init_search_runners", fake_init_search_runners) + metric = runner._run_cloud_cold_latency_case(drop_old=False) - metrics = runner._run_perf_case(drop_old=False) + assert metric.payload_estimated_bytes_per_query == expected_bytes - assert metrics.payload_profile == PayloadProfile.VECTOR.value - assert metrics.payload_estimated_bytes_per_query == expected_bytes + +def test_milvus_fts_payload_estimate_accepts_vector_width() -> None: + runner = _case_runner(MilvusFtsConfig(), CaseConfig(case_id=CaseType.FTSBm25Performance, k=10)) + + assert runner._estimated_payload_bytes_per_query() == 200 def test_milvus_vector_payload_requests_vector_field_and_returns_ids(): diff --git a/vectordb_bench/backend/cases.py b/vectordb_bench/backend/cases.py index 4c92d8248..8cc27fe45 100644 --- a/vectordb_bench/backend/cases.py +++ b/vectordb_bench/backend/cases.py @@ -5,7 +5,7 @@ from vectordb_bench import config from vectordb_bench.backend.clients.api import MetricType from vectordb_bench.backend.filter import Filter, FilterOp, IntFilter, LabelFilter, NewIntFilter, NonFilter, non_filter -from vectordb_bench.backend.payload import PayloadProfile +from vectordb_bench.backend.payload import FLOAT32_BYTES_PER_DIM, PayloadProfile from vectordb_bench.base import BaseModel from vectordb_bench.frontend.components.custom.getCustomConfig import CustomDatasetConfig @@ -138,7 +138,9 @@ class Case(BaseModel): def filters(self) -> Filter: return non_filter - def estimated_payload_bytes_per_query(self, k: int | None, vector_bytes_per_dim: int = 4) -> int: + def estimated_payload_bytes_per_query( + self, k: int | None, vector_bytes_per_dim: int = FLOAT32_BYTES_PER_DIM + ) -> int: if k is None: k = config.K_DEFAULT return self.payload_profile.estimated_bytes_per_query( @@ -929,7 +931,9 @@ def filters(self) -> Filter: int_value = int(self.dataset.data.size * self.filter_rate) return NewIntFilter(filter_rate=self.filter_rate, int_field=FTS_FILTER_ID_FIELD, int_value=int_value) - def estimated_payload_bytes_per_query(self, k: int | None, vector_bytes_per_dim: int = 4) -> int: + def estimated_payload_bytes_per_query( + self, k: int | None, vector_bytes_per_dim: int = FLOAT32_BYTES_PER_DIM + ) -> int: if k is None: k = config.K_DEFAULT return self.payload_profile.estimated_bytes_per_query(k=k, dim=0, vector_bytes_per_dim=vector_bytes_per_dim) diff --git a/vectordb_bench/backend/clients/milvus/config.py b/vectordb_bench/backend/clients/milvus/config.py index 09fcb61a9..7482dac00 100644 --- a/vectordb_bench/backend/clients/milvus/config.py +++ b/vectordb_bench/backend/clients/milvus/config.py @@ -3,6 +3,8 @@ from pydantic import BaseModel, Field, SecretStr, field_validator +from vectordb_bench.backend.payload import FLOAT32_BYTES_PER_DIM + from ..api import DBCaseConfig, DBConfig, IndexType, MetricType, SQType @@ -87,6 +89,10 @@ class VectorDataType(StrEnum): FLOAT32 = "FLOAT32" BFLOAT16 = "BFLOAT16" + @property + def bytes_per_dim(self) -> int: + return {VectorDataType.FLOAT32: FLOAT32_BYTES_PER_DIM, VectorDataType.BFLOAT16: 2}[self] + class HNSWConfig(MilvusIndexConfig, DBCaseConfig): M: int @@ -97,7 +103,7 @@ class HNSWConfig(MilvusIndexConfig, DBCaseConfig): @property def vector_bytes_per_dim(self) -> int: - return 2 if self.vector_data_type == VectorDataType.BFLOAT16 else 4 + return self.vector_data_type.bytes_per_dim def index_param(self) -> dict: return { diff --git a/vectordb_bench/backend/clients/milvus/milvus.py b/vectordb_bench/backend/clients/milvus/milvus.py index bacc04b3e..8de35f18b 100644 --- a/vectordb_bench/backend/clients/milvus/milvus.py +++ b/vectordb_bench/backend/clients/milvus/milvus.py @@ -29,7 +29,7 @@ def _float32_to_bf16_bytes(vectors: Iterable[list[float]]) -> list[bytes]: rounding_bias = np.uint32(0x7FFF) + ((bits >> np.uint32(16)) & np.uint32(1)) bf16 = ((bits + rounding_bias) >> np.uint32(16)).astype(" np.uint32(0x7F800000) bf16[nan_mask] = ((bits[nan_mask] >> np.uint32(16)) | np.uint32(0x40)).astype(" int: + def estimated_bytes_per_query(self, *, k: int, dim: int, vector_bytes_per_dim: int = FLOAT32_BYTES_PER_DIM) -> int: # Approximate payload size used for cloud leaderboard cost expansion. # ID + distance is about 20 bytes per hit; vectors default to float32 width. id_distance_bytes = 20 diff --git a/vectordb_bench/backend/task_runner.py b/vectordb_bench/backend/task_runner.py index 280bacb46..457906b91 100644 --- a/vectordb_bench/backend/task_runner.py +++ b/vectordb_bench/backend/task_runner.py @@ -17,6 +17,7 @@ from .clients import DB, MetricType, api from .data_source import DatasetSource from .filter import FilterOp +from .payload import FLOAT32_BYTES_PER_DIM from .runner import ( ColdWarmSearchRunner, ConcurrentInsertRunner, @@ -370,7 +371,7 @@ def _run_capacity_case(self) -> Metric: def _estimated_payload_bytes_per_query(self) -> int: return self.ca.estimated_payload_bytes_per_query( self.config.case_config.k, - vector_bytes_per_dim=getattr(self.config.db_case_config, "vector_bytes_per_dim", 4), + vector_bytes_per_dim=getattr(self.config.db_case_config, "vector_bytes_per_dim", FLOAT32_BYTES_PER_DIM), ) def _run_perf_case(self, drop_old: bool = True) -> Metric: