diff --git a/README.md b/README.md index 0ffa98f0e..e92afbeee 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 --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 55128a1b2..95a6bc92b 100644 --- a/tests/test_milvus.py +++ b/tests/test_milvus.py @@ -8,22 +8,241 @@ 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 DBCaseConfig, IndexType, MetricType +from vectordb_bench.backend.clients.milvus.config import ( + AutoIndexConfig, + HNSWConfig, + MilvusConfig, + MilvusFtsConfig, + VectorDataType, +) +from vectordb_bench.backend.clients.milvus.milvus import ( + MILVUS_FORCE_MERGE_TARGET_SIZE_MB, + 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 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], + [0x7F800001, 0x7FFFFFFF, 0xFFFFFFFF, 0xFF800001], + ], + dtype=np.uint32, + ) + + encoded = _float32_to_bf16_bytes(float32_bits.view(np.float32)) + actual = np.frombuffer(b"".join(encoded), dtype=" 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 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( + 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, + HNSWConfig( + M=8, + efConstruction=64, + ef=32, + metric_type=MetricType.COSINE, + vector_data_type=VectorDataType.BFLOAT16, + ), + ) + + 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 `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 + + +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 _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=db_case_config, + case_config=case_config, + stages=[], + ), + ca=case_config.case, + status=RunningStatus.PENDING, + dataset_source=DatasetSource.S3, + ) + + +@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 + + +@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: {}), + ) + + metric = runner._run_cloud_cold_latency_case(drop_old=False) + + assert metric.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(): captured = {} @@ -87,10 +306,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() @@ -128,6 +347,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 d01628f10..e88f0c40e 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 HNSWConfig, VectorDataType from vectordb_bench.backend.clients.zilliz_cloud import cli as zilliz_cli from vectordb_bench.cli import cli as common_cli @@ -124,6 +126,58 @@ def fake_run(**kwargs): assert captured["db_case_config"].force_merge_target_size_mb == 256 +@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): + captured.update(kwargs) + + monkeypatch.setattr(milvus_cli, "run", fake_run) + + result = CliRunner().invoke( + milvus_cli.MilvusHNSW, + [ + "--uri", + "http://localhost:19530", + "--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, HNSWConfig) + assert case_config.index == IndexType.HNSW + assert (case_config.M, case_config.efConstruction, case_config.ef) == (30, 360, 100) + 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( ("level_args", "expected_level"), [(["--level", "2"], 2), ([], None)], diff --git a/vectordb_bench/backend/cases.py b/vectordb_bench/backend/cases.py index 67c6a9c8c..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,10 +138,16 @@ 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 = FLOAT32_BYTES_PER_DIM + ) -> 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 +931,12 @@ 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 = 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) + 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/cli.py b/vectordb_bench/backend/clients/milvus/cli.py index 4df345ae7..46db50167 100644 --- a/vectordb_bench/backend/clients/milvus/cli.py +++ b/vectordb_bench/backend/clients/milvus/cli.py @@ -186,9 +186,22 @@ def MilvusFlat(**parameters: Unpack[MilvusFlatTypedDict]): class MilvusHNSWTypedDict(CommonTypedDict, MilvusTypedDict, HNSWFlavor3): ... +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", + ), + ] + + @cli.command() -@click_parameter_decorators_from_typed_dict(MilvusHNSWTypedDict) -def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]): +@click_parameter_decorators_from_typed_dict(MilvusHNSWCommandTypedDict) +def MilvusHNSW(**parameters: Unpack[MilvusHNSWCommandTypedDict]): from .config import HNSWConfig run( @@ -199,6 +212,7 @@ def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]): 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 1bc3b569c..7482dac00 100644 --- a/vectordb_bench/backend/clients/milvus/config.py +++ b/vectordb_bench/backend/clients/milvus/config.py @@ -1,7 +1,10 @@ +from enum import StrEnum from typing import Annotated, ClassVar 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 @@ -82,11 +85,25 @@ def search_param(self) -> dict: return params +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 efConstruction: int ef: int | None = None index: IndexType = IndexType.HNSW + vector_data_type: VectorDataType = VectorDataType.FLOAT32 + + @property + def vector_bytes_per_dim(self) -> int: + 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 4296e6740..8de35f18b 100644 --- a/vectordb_bench/backend/clients/milvus/milvus.py +++ b/vectordb_bench/backend/clients/milvus/milvus.py @@ -6,13 +6,14 @@ 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 .config import MilvusFtsConfig, MilvusIndexConfig +from .config import MilvusFtsConfig, MilvusIndexConfig, VectorDataType 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: @@ -36,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, @@ -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, "vector_data_type", None) == VectorDataType.BFLOAT16 + 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( @@ -176,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: @@ -429,12 +474,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 +612,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/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 = [ diff --git a/vectordb_bench/backend/payload.py b/vectordb_bench/backend/payload.py index 04cba3b85..727483470 100644 --- a/vectordb_bench/backend/payload.py +++ b/vectordb_bench/backend/payload.py @@ -1,5 +1,7 @@ from enum import StrEnum +FLOAT32_BYTES_PER_DIM = 4 + class PayloadProfile(StrEnum): IDS_ONLY = "ids_only" @@ -7,15 +9,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 = FLOAT32_BYTES_PER_DIM) -> 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..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, @@ -367,6 +368,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", FLOAT32_BYTES_PER_DIM), + ) + def _run_perf_case(self, drop_old: bool = True) -> Metric: """run performance cases @@ -447,9 +454,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 +565,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()