diff --git a/README.md b/README.md index b08458c79..b6391fdee 100644 --- a/README.md +++ b/README.md @@ -943,6 +943,47 @@ We've developed lots of comprehensive benchmark cases to test vector databases' - **Large Dataset:** Similar to the XLarge Dataset case, but uses a slightly smaller dataset (10M-1024dim, 10M-768dim, 5M-1536dim). - **Medium Dataset:** A case using a medium dataset (1M-1024dim, 1M-768dim, 500K-1536dim). - **Small Dataset:** For development (100K-768dim, 50K-1536dim). + +##### LAION-100M Large-TopK + +`Performance768D100M` selects its query and ground-truth files from `--k`: + +| Requested K | Query file | Ground-truth file | Queries | GT width | +|---:|---|---|---:|---:| +| `1..1,000` | `test.parquet` | `neighbors.parquet` | 1,000 | 1,000 | +| `1,001..100,000` | `test_nq200.parquet` | `neighbors_top100k_nq200.parquet` | 200 | 100,000 | +| `100,001..1,000,000` | `test_nq200.parquet` | `neighbors_top1m_nq200.parquet` | 200 | 1,000,000 | + +K must be positive, and LAION-100M rejects values above 1,000,000. Integer-filter runs use the smallest published GT width that covers K: + +| Integer filter rate | Published GT widths | Maximum K | +|---:|---:|---:| +| 50%, 60%, 70%, 80%, 90%, 95%, 98%, 99% | 100K, 1M | 1M | +| 99.5% | 100K, 500K | 500K | +| 99.8% | 100K, 200K | 200K | +| 99.9% | 100K | 100K | + +K up to 1,000 keeps the original 1,000-query filtered artifacts; larger K uses `test_nq200.parquet`. Other integer filter rates and label-filter runs above K=1,000 are rejected before database initialization. VDBBench validates query IDs, row counts, and GT width before issuing a search. + +```bash +vectordbbench milvusautoindex --uri http://localhost:19530 --case-type NewIntFilterPerformanceCase --dataset-with-size-type "Large LAION (768dim, 100M)" --filter-rate 0.99 --k 1000000 +``` + +Wide GT remains in Parquet/Arrow form and is opened inside the serial-search subprocess one query row at a time. Results include primary `recall@K`, `recall_at` for the available cutoffs among 100, 1K, 10K, 100K, and 1M, plus serial and concurrent p50/p95/p99 latency. Concurrent throughput continues to use the configured fixed-duration phase. + +For Milvus and Zilliz Cloud performance runs with K above 16,384, VDBBench automatically creates new collections with `query_mode=large_topk` before creating the vector index. Reused collections are validated and rejected when that property is missing or incompatible. The target database, selected mode, and requested K are written to the run log. Self-hosted Milvus must be 2.6.14 or later, the release that introduced the `query_mode=large_topk` collection property; earlier servers accept the property without honoring it. Other backends must already permit the requested K; VDBBench forwards K unchanged and does not alter their collection properties. + +##### Performance Response Payloads + +Every vector search performance case can be configured for either an IDs-only response or a response that also includes each result vector. IDs only remains the default. Run the scenarios separately from the CLI: + +```bash +vectordbbench milvusautoindex --uri http://localhost:19530 --case-type Performance768D100M --k 1000000 --payload-profile ids_only +vectordbbench milvusautoindex --uri http://localhost:19530 --case-type Performance768D100M --k 1000000 --payload-profile vector +``` + +The frontend can select one or both scenarios for Milvus and Zilliz Cloud. Each scenario produces independent P99 latency, QPS, and recall metrics. `qps` remains the highest observed QPS among the configured concurrency levels; VDBBench does not discover a backend concurrency limit. + #### Filtering Search Performance Case - **Int-Filter Cases:** Evaluates search performance with int-based filter expression (e.g. "id >= 2,000"). - **Label-Filter Cases:** Evaluates search performance with label-based filter expressions (e.g., "color == 'red'"). The test includes randomly generated labels to simulate real-world filtering scenarios. diff --git a/docs/release/2026-08-large-topk.md b/docs/release/2026-08-large-topk.md new file mode 100644 index 000000000..d938fb5ae --- /dev/null +++ b/docs/release/2026-08-large-topk.md @@ -0,0 +1,42 @@ +# LAION-100M Large-TopK + +VDBBench now supports K values through 1,000,000 on the existing `Performance768D100M` case. + +## Dataset Selection + +The case selects hosted LAION artifacts from K: + +- K up to 1,000 uses `test.parquet` and `neighbors.parquet`. +- K from 1,001 through 100,000 uses the 200-query `test_nq200.parquet` and `neighbors_top100k_nq200.parquet` files. +- K from 100,001 through 1,000,000 uses `test_nq200.parquet` and `neighbors_top1m_nq200.parquet`. + +Integer-filter runs use their original 1,000-query artifacts through K=1,000. Above that, VDBBench uses the smallest published 200-query GT tier that covers K: + +| Integer filter rate | Published GT widths | Maximum K | +|---:|---:|---:| +| 50%, 60%, 70%, 80%, 90%, 95%, 98%, 99% | 100K, 1M | 1M | +| 99.5% | 100K, 500K | 500K | +| 99.8% | 100K, 200K | 200K | +| 99.9% | 100K | 100K | + +The loader verifies query ID alignment, row count, and ground-truth width. Unsupported integer rates, requests above the selected filter's maximum K, label-filter runs above K=1,000, and unfiltered LAION K above 1,000,000 fail before database initialization. LAION-backed workloads that do not measure recall, such as cold latency, keep the standard 1,000-query artifacts while forwarding their configured K to the backend. + +## Memory And Metrics + +Ground truth is represented by a local Parquet path and compact metadata in the parent process. The serial-search subprocess opens that path and reads one Arrow/NumPy neighbor row at a time, avoiding conversion of the full wide GT into Python integer lists. + +Recall and NDCG now use O(K) hash lookups. Recall counts each matching returned ID at most once, so duplicate IDs do not inflate the score. A large-TopK serial run reports: + +- primary recall at the requested K; +- `recall_at` for each supported cutoff no greater than K; +- serial p50, p95, and p99 latency. + +Serial and concurrent latency fields are stored in seconds, matching the existing p95/p99 fields; the frontend converts them to milliseconds for display. `recall_at` values are ratios from 0 to 1. Existing result files load with zero/empty defaults for the new fields. + +## Milvus Collection Mode + +For Milvus and Zilliz Cloud performance runs with K above 16,384, VDBBench sets `query_mode=large_topk` when it creates the collection, before creating the vector index. The run log records the target database, the requested K, and the selected query mode. When reusing a collection, VDBBench validates the property and fails before loading or searching if the collection is incompatible. + +Self-hosted Milvus requires 2.6.14 or later, the release that introduced the `query_mode=large_topk` collection property. Earlier servers accept the property without honoring it, so the requested K still fails against the default 16,384 limit, and the reuse check rejects the collection on the next run. + +Other backends are unchanged. Their target collection must already support the requested result count before the benchmark starts. diff --git a/pyproject.toml b/pyproject.toml index 223291721..acd30275b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ "psutil", "polars", "plotly", + "matplotlib", "environs", "pydantic>=2.0,<3", "scikit-learn", diff --git a/tests/test_case_runner_reuse.py b/tests/test_case_runner_reuse.py index ce06e5ba0..44315c635 100644 --- a/tests/test_case_runner_reuse.py +++ b/tests/test_case_runner_reuse.py @@ -6,6 +6,7 @@ from vectordb_bench.backend.clients.doris.config import DorisCaseConfig, DorisConfig from vectordb_bench.backend.clients.pinecone.config import PineconeConfig from vectordb_bench.backend.clients.turbopuffer.config import TurboPufferConfig, TurboPufferIndexConfig +from vectordb_bench.backend.clients.zilliz_cloud.config import AutoIndexConfig, ZillizCloudConfig from vectordb_bench.backend.data_source import DatasetSource from vectordb_bench.backend.dataset import DatasetWithSizeType from vectordb_bench.backend.task_runner import CaseRunner, RunningStatus, TaskRunner @@ -23,6 +24,7 @@ def make_runner( db: DB = DB.TurboPuffer, db_config=None, db_case_config=None, + k: int = 100, stages: list[TaskStage] | None = None, insert_batch_size: int = DEFAULT_INSERT_BATCH_SIZE, ) -> CaseRunner: @@ -33,6 +35,8 @@ def make_runner( db_config = PineconeConfig(api_key="key", index_name="idx") elif db == DB.Doris: db_config = DorisConfig(password=SecretStr("")) + elif db == DB.ZillizCloud: + db_config = ZillizCloudConfig(uri=SecretStr("http://example.invalid")) else: db_config = DB.Test.config_cls() if db_case_config is None: @@ -40,6 +44,8 @@ def make_runner( db_case_config = TurboPufferIndexConfig(metric_type=MetricType.COSINE) elif db == DB.Doris: db_case_config = DorisCaseConfig(metric_type=MetricType.COSINE) + elif db == DB.ZillizCloud: + db_case_config = AutoIndexConfig(metric_type=MetricType.COSINE) else: db_case_config = EmptyDBCaseConfig() @@ -47,7 +53,7 @@ def make_runner( db=db, db_config=db_config, db_case_config=db_case_config, - case_config=CaseConfig(case_id=case_id, custom_case=custom_case or {}), + case_config=CaseConfig(case_id=case_id, custom_case=custom_case or {}, k=k), stages=stages or [TaskStage.DROP_OLD, TaskStage.LOAD, TaskStage.SEARCH_SERIAL], insert_batch_size=insert_batch_size, ) @@ -123,6 +129,13 @@ def test_reuse_key_distinguishes_insert_batch_size(): ) +def test_reuse_key_distinguishes_zilliz_default_and_large_topk_modes(): + assert_not_reusable( + make_runner(db=DB.ZillizCloud, case_id=CaseType.Performance768D100M, k=16_384), + make_runner(db=DB.ZillizCloud, case_id=CaseType.Performance768D100M, k=16_385), + ) + + def test_reuse_key_distinguishes_physical_db_targets(): assert_not_reusable( make_runner(db_config=TurboPufferConfig(api_key="key", region="aws-us-east-1", namespace="namespace_a")), diff --git a/tests/test_cloud_payload_case.py b/tests/test_cloud_payload_case.py index d3d97fb60..0aa76ed5a 100644 --- a/tests/test_cloud_payload_case.py +++ b/tests/test_cloud_payload_case.py @@ -1,6 +1,7 @@ import pytest from vectordb_bench import config +from vectordb_bench.backend.assembler import Assembler from vectordb_bench.backend.cases import CaseType, CloudPayloadSearchCase from vectordb_bench.backend.clients import DB from vectordb_bench.backend.clients.api import EmptyDBCaseConfig @@ -64,6 +65,52 @@ def test_case_config_builds_cloud_payload_case_from_custom_case(): assert case.payload_profile == PayloadProfile.VECTOR +def test_case_config_preserves_legacy_payload_profile(): + case_config = CaseConfig( + case_id=CaseType.CloudPayloadSearchCase, + custom_case={"payload_profile": "vector"}, + ) + + assert case_config.payload_profile is None + assert case_config.case.payload_profile == PayloadProfile.VECTOR + + +def test_case_config_accepts_matching_top_level_and_legacy_payload_profiles(): + case_config = CaseConfig( + case_id=CaseType.CloudPayloadSearchCase, + custom_case={"payload_profile": "vector"}, + payload_profile=PayloadProfile.VECTOR, + ) + + assert case_config.payload_profile == PayloadProfile.VECTOR + assert case_config.case.payload_profile == PayloadProfile.VECTOR + + +def test_case_config_rejects_conflicting_payload_profiles(): + with pytest.raises(ValueError, match="conflicts with custom_case"): + CaseConfig( + case_id=CaseType.CloudPayloadSearchCase, + custom_case={"payload_profile": "ids_only"}, + payload_profile=PayloadProfile.VECTOR, + ) + + +def test_assembler_preserves_top_level_performance_payload_profile(): + task = TaskConfig( + db=DB.Test, + db_config=DB.Test.config_cls(), + db_case_config=EmptyDBCaseConfig(), + case_config=CaseConfig( + case_id=CaseType.Performance768D100M, + payload_profile=PayloadProfile.VECTOR, + ), + ) + + runner = Assembler.assemble("run-id", task, DatasetSource.S3) + + assert runner.ca.payload_profile == PayloadProfile.VECTOR + + def test_case_runner_reuse_key_distinguishes_scalar_label_schema_requirement(): ids_only_case = CloudPayloadSearchCase( dataset_with_size_type=DatasetWithSizeType.CohereSmall.value, @@ -179,3 +226,46 @@ def test_search_runners_fail_fast_for_unsupported_payload_profile(): k=3, payload_profile=PayloadProfile.VECTOR, ) + + +def test_case_runner_rejects_unsupported_payload_before_db_init(monkeypatch: pytest.MonkeyPatch): + events = [] + case_config = CaseConfig( + case_id=CaseType.Performance768D100M, + payload_profile=PayloadProfile.VECTOR, + ) + task = TaskConfig( + db=DB.Test, + db_config=DB.Test.config_cls(), + db_case_config=EmptyDBCaseConfig(), + case_config=case_config, + ) + runner = CaseRunner( + run_id="run-id", + config=task, + ca=case_config.case, + status=RunningStatus.PENDING, + dataset_source=DatasetSource.S3, + ) + + monkeypatch.setattr( + type(runner.ca.dataset), + "resolve_search_files", + lambda self, **kwargs: events.append("resolve"), + ) + monkeypatch.setattr( + type(runner.ca.dataset), + "prepare", + lambda self, *args, **kwargs: events.append("prepare"), + ) + + def fake_init_db(self, drop_old=True): + events.append("init_db") + self.db = FakeDB() + + monkeypatch.setattr(CaseRunner, "init_db", fake_init_db) + + with pytest.raises(NotImplementedError, match="payload_profile=vector"): + runner._pre_run(drop_old=False) + + assert events == [] diff --git a/tests/test_cloud_payload_search.py b/tests/test_cloud_payload_search.py index cda8f63ba..4a2602f49 100644 --- a/tests/test_cloud_payload_search.py +++ b/tests/test_cloud_payload_search.py @@ -42,6 +42,7 @@ def test_scalar_label_payload_profile_requires_scalar_label_materialization_with def test_dataset_prepare_loads_separated_scalar_labels_for_scalar_payload(monkeypatch): dataset = Dataset.LAION.manager(100_000_000) dataset.data.with_remote_resource = False + dataset.data.with_gt = False loaded_scalar_labels = object() def fake_read_file(file_name): @@ -118,11 +119,13 @@ def test_serial_search_runner_passes_tenant_and_skips_recall(): measure_recall=False, ) - recall, ndcg, p99, p95 = runner.search((runner.test_data, runner.ground_truth)) + recall, ndcg, p99, p95, p50, recall_at = runner.search((runner.test_data, runner.ground_truth)) assert recall == 0 assert ndcg == 0 assert p99 >= 0 assert p95 >= 0 + assert p50 >= 0 + assert recall_at == {} assert set(db.tenants).issubset({"tenant_0000", "tenant_0001"}) assert db.tenants diff --git a/tests/test_dataset.py b/tests/test_dataset.py index d4ccb283d..3e5d469a7 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1,12 +1,190 @@ -from vectordb_bench.backend.dataset import Dataset import logging +import pickle + +import polars as pl import pytest from pydantic import ValidationError -from vectordb_bench.backend.data_source import DatasetSource +from vectordb_bench import config +from vectordb_bench.backend import dataset as dataset_module +from vectordb_bench.backend.clients import MetricType +from vectordb_bench.backend.data_source import DatasetSource +from vectordb_bench.backend.dataset import CustomDataset, Dataset, DatasetManager +from vectordb_bench.backend.filter import LabelFilter, NewIntFilter, non_filter log = logging.getLogger("vectordb_bench") + +@pytest.mark.parametrize( + ("k", "test_file", "gt_file", "width", "query_count"), + [ + (1_000, "test.parquet", "neighbors.parquet", 1_000, 1_000), + (1_001, "test_nq200.parquet", "neighbors_top100k_nq200.parquet", 100_000, 200), + (100_000, "test_nq200.parquet", "neighbors_top100k_nq200.parquet", 100_000, 200), + (100_001, "test_nq200.parquet", "neighbors_top1m_nq200.parquet", 1_000_000, 200), + (1_000_000, "test_nq200.parquet", "neighbors_top1m_nq200.parquet", 1_000_000, 200), + ], +) +def test_laion_artifact_selection(k, test_file, gt_file, width, query_count): + dataset = Dataset.LAION.manager(100_000_000) + assert hasattr(dataset, "resolve_search_files") + + files = dataset.resolve_search_files(k=k, filters=non_filter) + + assert files.test_file == test_file + assert files.gt_file == gt_file + assert files.width == width + assert files.query_count == query_count + + +@pytest.mark.parametrize("k", [0, -1, 1_000_001]) +def test_laion_artifact_selection_rejects_unsupported_k(k): + dataset = Dataset.LAION.manager(100_000_000) + assert hasattr(dataset, "resolve_search_files") + + with pytest.raises(ValueError, match="LAION"): + dataset.resolve_search_files(k=k, filters=non_filter) + + +def _laion_int_filter(filter_rate: float) -> NewIntFilter: + return NewIntFilter(filter_rate=filter_rate, int_field="id", int_value=int(100_000_000 * filter_rate)) + + +def test_laion_integer_filter_keeps_standard_artifacts_at_k_1000(): + dataset = Dataset.LAION.manager(100_000_000) + files = dataset.resolve_search_files(k=1_000, filters=_laion_int_filter(0.5)) + + assert files.test_file == "test.parquet" + assert files.gt_file == "neighbors_int_50p.parquet" + + +@pytest.mark.parametrize( + ("filter_rate", "max_k"), + [ + (0.5, 1_000_000), + (0.6, 1_000_000), + (0.7, 1_000_000), + (0.8, 1_000_000), + (0.9, 1_000_000), + (0.95, 1_000_000), + (0.98, 1_000_000), + (0.99, 1_000_000), + (0.995, 500_000), + (0.998, 200_000), + (0.999, 100_000), + ], +) +def test_laion_integer_filter_selects_published_maximum_gt(filter_rate, max_k): + dataset = Dataset.LAION.manager(100_000_000) + filters = _laion_int_filter(filter_rate) + files = dataset.resolve_search_files(k=max_k, filters=filters) + width_suffix = f"{max_k // 1_000_000}m" if max_k >= 1_000_000 else f"{max_k // 1_000}k" + + assert files.test_file == "test_nq200.parquet" + assert files.gt_file == f"neighbors_{filters.int_rate}_top{width_suffix}_nq200.parquet" + assert files.width == max_k + assert files.query_count == 200 + + +def test_laion_integer_filter_selects_top100k_for_smaller_large_topk(): + dataset = Dataset.LAION.manager(100_000_000) + files = dataset.resolve_search_files(k=1_001, filters=_laion_int_filter(0.995)) + + assert files.gt_file == "neighbors_int_99.5p_top100k_nq200.parquet" + assert files.width == 100_000 + + +@pytest.mark.parametrize( + ("filters", "k", "error"), + [ + (_laion_int_filter(0.75), 1_000, "supported filter rates"), + (_laion_int_filter(0.999), 100_001, "supports K up to 100,000"), + (LabelFilter(label_percentage=0.5), 1_001, "integer filters"), + ], +) +def test_laion_filtered_artifact_selection_rejects_unpublished_combinations(filters, k, error): + dataset = Dataset.LAION.manager(100_000_000) + + with pytest.raises(ValueError, match=error): + dataset.resolve_search_files(k=k, filters=filters) + + +def test_dataset_prepare_keeps_ground_truth_path_based(tmp_path, monkeypatch): + assert hasattr(dataset_module, "ParquetGroundTruth") + monkeypatch.setattr(config, "DATASET_LOCAL_DIR", tmp_path) + dataset = _custom_dataset_manager() + dataset.data.with_remote_resource = False + dataset.data_dir.mkdir(parents=True) + _write_vector_fixture(dataset.data_dir) + + dataset.prepare(with_train_files=False, k=4) + + assert isinstance(dataset.gt_data, dataset_module.ParquetGroundTruth) + assert dataset.gt_data.row_count == 2 + assert dataset.gt_data.width == 4 + restored = pickle.loads(pickle.dumps(dataset.gt_data)) + assert [row.tolist() for row in restored.iter_rows()] == [[1, 2, 3, 4], [5, 6, 7, 8]] + + +def test_parquet_ground_truth_rejects_query_id_mismatch(tmp_path): + assert hasattr(dataset_module, "ParquetGroundTruth") + gt_path = tmp_path / "neighbors.parquet" + pl.DataFrame({"id": [11, 20], "neighbors_id": [[1, 2], [3, 4]]}).write_parquet(gt_path) + + with pytest.raises(ValueError, match="query IDs"): + dataset_module.ParquetGroundTruth.from_file( + gt_path, + id_field="id", + neighbors_field="neighbors_id", + expected_query_ids=[10, 20], + minimum_width=2, + ) + + +def test_parquet_ground_truth_rejects_narrow_row(tmp_path): + assert hasattr(dataset_module, "ParquetGroundTruth") + gt_path = tmp_path / "neighbors.parquet" + pl.DataFrame({"id": [10, 20], "neighbors_id": [[1, 2], [3]]}).write_parquet(gt_path) + + with pytest.raises(ValueError, match="width"): + dataset_module.ParquetGroundTruth.from_file( + gt_path, + id_field="id", + neighbors_field="neighbors_id", + expected_query_ids=[10, 20], + minimum_width=2, + ) + + +def _custom_dataset_manager() -> DatasetManager: + data = CustomDataset( + name="local", + size=8, + dim=2, + metric_type=MetricType.L2, + use_shuffled=False, + with_gt=True, + dir="large_topk_fixture", + file_num=1, + ) + return DatasetManager(data=data) + + +def _write_vector_fixture(data_dir): + pl.DataFrame( + { + "id": [10, 20], + "emb": [[0.1, 0.2], [0.3, 0.4]], + } + ).write_parquet(data_dir / "test.parquet") + pl.DataFrame( + { + "id": [10, 20], + "neighbors_id": [[1, 2, 3, 4], [5, 6, 7, 8]], + } + ).write_parquet(data_dir / "neighbors.parquet") + + class TestDataSet: def test_iter_dataset(self): for ds in Dataset: @@ -29,6 +207,7 @@ def test_iter_cohere(self): cohere_10m.prepare() import time + before = time.time() for i in cohere_10m: log.debug(i.head(1)) @@ -40,9 +219,11 @@ def test_iter_cohere(self): def test_iter_laion(self): laion_100m = Dataset.LAION.manager(100_000_000) from vectordb_bench.backend.data_source import DatasetSource + laion_100m.prepare(source=DatasetSource.AliyunOSS) import time + before = time.time() for i in laion_100m: log.debug(i.head(1)) @@ -74,4 +255,3 @@ def test_download_small(self): files=files, local_ds_root=openai_50k.data_dir, ) - diff --git a/tests/test_frontend_run_settings.py b/tests/test_frontend_run_settings.py index 8e4b2a1eb..f3d00cd45 100644 --- a/tests/test_frontend_run_settings.py +++ b/tests/test_frontend_run_settings.py @@ -4,13 +4,20 @@ from vectordb_bench.backend.cases import CaseType from vectordb_bench.backend.clients import DB +from vectordb_bench.backend.clients.api import EmptyDBCaseConfig +from vectordb_bench.backend.dataset import DatasetWithSizeType from vectordb_bench.frontend.components.run_test import generateTasks from vectordb_bench.frontend.components.run_test.runSettings import ( DEFAULT_STREAMING_INSERT_RATE, validate_streaming_insert_rates, ) +from vectordb_bench.frontend.components.run_test.submitTask import ( + advancedSettings, + apply_run_settings, + get_max_search_k, +) from vectordb_bench.frontend.config.dbCaseConfigs import custom_streaming_config -from vectordb_bench.models import CaseConfig, CaseConfigParamType +from vectordb_bench.models import CaseConfig, CaseConfigParamType, TaskConfig def streaming_case(insert_rate: int | None = None) -> CaseConfig: @@ -18,6 +25,115 @@ def streaming_case(insert_rate: int | None = None) -> CaseConfig: return CaseConfig(case_id=CaseType.StreamingPerformanceCase, custom_case=custom_case) +def laion_task(filter_rate: float | None = None) -> TaskConfig: + if filter_rate is None: + case_config = CaseConfig(case_id=CaseType.Performance768D100M) + else: + case_config = CaseConfig( + case_id=CaseType.NewIntFilterPerformanceCase, + custom_case={ + "dataset_with_size_type": DatasetWithSizeType.LAIONLarge, + "filter_rate": filter_rate, + }, + ) + return TaskConfig( + db=DB.Test, + db_config=DB.Test.config_cls(), + db_case_config=EmptyDBCaseConfig(), + case_config=case_config, + ) + + +def fts_task() -> TaskConfig: + return TaskConfig( + db=DB.Test, + db_config=DB.Test.config_cls(), + db_case_config=EmptyDBCaseConfig(), + case_config=CaseConfig(case_id=CaseType.FTSBm25Performance), + ) + + +class SettingsContainer: + def __init__(self): + self.number_inputs = {} + + def columns(self, _widths): + return [self, self] + + def checkbox(self, _label, *, value): + return value + + def caption(self, _text): + return None + + def number_input(self, label, **kwargs): + self.number_inputs[label] = kwargs + return kwargs["value"] + + def text_input(self, _label, *, value, **_kwargs): + return value + + +def test_global_k_control_uses_most_restrictive_selected_case(): + tasks = [laion_task(), laion_task(0.995), laion_task(0.999), fts_task()] + container = SettingsContainer() + + assert get_max_search_k([laion_task()]) == 1_000_000 + assert get_max_search_k([laion_task(0.995)]) == 500_000 + assert get_max_search_k([laion_task(0.998)]) == 200_000 + assert get_max_search_k(tasks) == 100_000 + advancedSettings(container, tasks) + + assert container.number_inputs["k"]["max_value"] == 100_000 + + +def test_global_k_control_leaves_fts_uncapped(): + container = SettingsContainer() + + assert get_max_search_k([fts_task()]) is None + advancedSettings(container, [fts_task()]) + + assert "max_value" not in container.number_inputs["k"] + + +def test_apply_run_settings_validates_all_cases_before_mutating_tasks(): + tasks = [laion_task(), laion_task(0.999)] + + with pytest.raises(ValueError, match="supports K up to 100,000"): + apply_run_settings( + tasks, + k=100_001, + concurrencies=[40, 60], + concurrency_duration=120, + load_concurrency=5, + ) + + assert [task.case_config.k for task in tasks] == [100, 100] + + +def test_apply_run_settings_applies_valid_global_settings(): + tasks = [laion_task(0.999), fts_task()] + + apply_run_settings( + tasks, + k=100_000, + concurrencies=[40, 60], + concurrency_duration=120, + load_concurrency=5, + ) + + assert [task.case_config.k for task in tasks] == [100_000, 100_000] + assert [task.case_config.concurrency_search_config.num_concurrency for task in tasks] == [ + [40, 60], + [40, 60], + ] + assert [task.case_config.concurrency_search_config.concurrency_duration for task in tasks] == [ + 120, + 120, + ] + assert [task.load_concurrency for task in tasks] == [5, 5] + + @pytest.mark.parametrize( ("insert_rate", "batch_size", "expected_message"), [ diff --git a/tests/test_large_topk_case.py b/tests/test_large_topk_case.py new file mode 100644 index 000000000..2f72ff74a --- /dev/null +++ b/tests/test_large_topk_case.py @@ -0,0 +1,284 @@ +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from vectordb_bench.backend.cases import CaseLabel, CaseType +from vectordb_bench.backend.clients import DB +from vectordb_bench.backend.clients.api import EmptyDBCaseConfig +from vectordb_bench.backend.data_source import DatasetSource +from vectordb_bench.backend.dataset import DatasetManager, DatasetWithSizeType +from vectordb_bench.backend.runner.mp_runner import MultiProcessingSearchRunner +from vectordb_bench.backend.runner.serial_runner import SerialSearchRunner +from vectordb_bench.backend.task_runner import CaseRunner, RunningStatus +from vectordb_bench.models import CaseConfig, TaskConfig, TaskStage + + +class SearchProbeDB: + name = "SearchProbeDB" + + def __init__(self, results=None): + self.results = list(results or []) + self.init_calls = 0 + self.search_calls = 0 + + def supports_payload_profile(self, payload_profile): + return True + + @contextmanager + def init(self): + self.init_calls += 1 + yield + + def prepare_filter(self, filters): + return None + + def search_embedding(self, query, k=100, payload_profile=None, tenant=None): + self.search_calls += 1 + if self.results: + return self.results.pop(0) + return [] + + +def _large_topk_property_runner(db: DB, k: int) -> CaseRunner: + return CaseRunner.model_construct( + config=SimpleNamespace(db=db, case_config=SimpleNamespace(k=k)), + ca=SimpleNamespace(label=CaseLabel.Performance), + ) + + +@pytest.mark.parametrize("db", [DB.Milvus, DB.ZillizCloud]) +def test_large_topk_selects_collection_mode_and_logs_requested_k(db, monkeypatch): + runner = _large_topk_property_runner(db, 100_000) + messages = [] + monkeypatch.setattr( + "vectordb_bench.backend.task_runner.log.info", + lambda message, *args: messages.append(message % args), + ) + + properties = runner._collection_properties(log_selection=True) + + assert properties == {"query_mode": "large_topk"} + assert db.value in messages[0] + assert "requested K=100000" in messages[0] + assert "query_mode=large_topk" in messages[0] + + +@pytest.mark.parametrize( + ("db", "k"), + [ + # Milvus and Zilliz Cloud stay in default mode at or below the TopK limit. + (DB.ZillizCloud, 16_384), + (DB.Milvus, 16_384), + (DB.Milvus, 100), + # Backends without a Large TopK collection mode are never reconfigured. + (DB.Pinecone, 100_000), + ], +) +def test_large_topk_collection_mode_does_not_change_other_workloads(db, k): + runner = _large_topk_property_runner(db, k) + + assert runner._collection_properties() == {} + + +def test_serial_runner_rejects_query_count_mismatch_before_db_init(): + db = SearchProbeDB() + runner = SerialSearchRunner( + db=db, + test_data=[[0.1], [0.2]], + ground_truth=[[1, 2, 3, 4]], + k=4, + ) + + with pytest.raises(ValueError, match="query count"): + runner.search((runner.test_data, runner.ground_truth)) + + assert db.init_calls == 0 + assert db.search_calls == 0 + + +def test_serial_runner_rejects_narrow_ground_truth_before_db_init(): + db = SearchProbeDB() + runner = SerialSearchRunner( + db=db, + test_data=[[0.1]], + ground_truth=[[1, 2]], + k=4, + ) + + with pytest.raises(ValueError, match="width"): + runner.search((runner.test_data, runner.ground_truth)) + + assert db.init_calls == 0 + assert db.search_calls == 0 + + +def test_serial_runner_reports_p50_and_prefix_correct_recall(): + ground_truth = list(range(1_000)) + results = [*range(50), *range(100, 150), *range(50, 100), *range(150, 1_000)] + db = SearchProbeDB(results=[results]) + runner = SerialSearchRunner( + db=db, + test_data=[[0.1]], + ground_truth=[ground_truth], + k=1_000, + ) + + recall, ndcg, p99, p95, p50, recall_at = runner.search((runner.test_data, runner.ground_truth)) + + assert recall == 1.0 + assert ndcg == 1.0 + assert p99 >= p50 >= 0 + assert p95 >= p50 + assert recall_at == {100: 0.5, 1_000: 1.0} + + +def test_concurrent_latency_aggregation_includes_p50(): + runner = MultiProcessingSearchRunner(db=SearchProbeDB(), test_data=[[0.1]]) + results = [ + (10, 0, {"p99": 0.9, "p95": 0.8, "p50": 0.3, "avg": 0.4, "count": 10}), + (20, 0, {"p99": 0.7, "p95": 0.6, "p50": 0.2, "avg": 0.3, "count": 20}), + ] + + assert runner._aggregate_latency_stats(results) == pytest.approx((0.9, 0.8, 0.3, 1 / 3)) + + +def test_concurrent_latency_aggregation_handles_empty_success_window(): + runner = MultiProcessingSearchRunner(db=SearchProbeDB(), test_data=[[0.1]]) + + assert runner._aggregate_latency_stats([]) == (0, 0, 0, 0) + assert hasattr(runner, "_latency_summary") + assert runner._latency_summary([]) == (0, 0, 0, 0) + + +@pytest.mark.parametrize( + "case_config", + [ + CaseConfig(case_id=CaseType.Performance768D100M, k=1_000_001), + CaseConfig( + case_id=CaseType.NewIntFilterPerformanceCase, + custom_case={"dataset_with_size_type": DatasetWithSizeType.LAIONLarge, "filter_rate": 0.999}, + k=100_001, + ), + ], +) +def test_case_runner_rejects_unsupported_laion_k_before_db_init(monkeypatch, case_config): + runner = CaseRunner( + run_id="large-topk", + config=TaskConfig( + db=DB.Test, + db_config=DB.Test.config_cls(), + db_case_config=EmptyDBCaseConfig(), + case_config=case_config, + ), + ca=case_config.case, + status=RunningStatus.PENDING, + dataset_source=DatasetSource.S3, + ) + init_called = False + + def fake_init_db(self, drop_old=True): + nonlocal init_called + init_called = True + + def fake_prepare(self, *args, **kwargs): + return self.resolve_search_files(k=case_config.k, filters=kwargs["filters"]) + + monkeypatch.setattr(CaseRunner, "init_db", fake_init_db) + monkeypatch.setattr(DatasetManager, "prepare", fake_prepare) + + with pytest.raises(ValueError, match="LAION"): + runner._pre_run(drop_old=False) + + assert init_called is False + + +def test_cloud_cold_latency_keeps_standard_query_artifacts_for_large_k(monkeypatch): + case_config = CaseConfig( + case_id=CaseType.CloudColdLatencyCase, + custom_case={"query_count": 1_000}, + k=1_001, + ) + runner = CaseRunner( + run_id="large-topk", + config=TaskConfig( + db=DB.Test, + db_config=DB.Test.config_cls(), + db_case_config=EmptyDBCaseConfig(), + case_config=case_config, + stages=[TaskStage.SEARCH_SERIAL], + ), + ca=case_config.case, + status=RunningStatus.PENDING, + dataset_source=DatasetSource.S3, + ) + selected_files = None + + def fake_init_db(self, drop_old=True): + return None + + def fake_prepare(self, *args, **kwargs): + nonlocal selected_files + selected_files = self.resolve_search_files(k=kwargs["k"], filters=kwargs["filters"]) + return True + + monkeypatch.setattr(CaseRunner, "init_db", fake_init_db) + monkeypatch.setattr(DatasetManager, "prepare", fake_prepare) + + runner._pre_run(drop_old=False) + + assert selected_files is not None + assert selected_files.test_file == "test.parquet" + assert selected_files.query_count == 1_000 + assert runner.config.case_config.k == 1_001 + + +def test_case_runner_propagates_large_topk_metrics(monkeypatch): + case_config = CaseConfig(case_id=CaseType.Performance768D100M, k=1_000) + runner = CaseRunner( + run_id="large-topk", + config=TaskConfig( + db=DB.Test, + db_config=DB.Test.config_cls(), + db_case_config=EmptyDBCaseConfig(), + 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.8, 0.7, 0.9, 0.8, 0.5, {100: 0.9, 1_000: 0.8}), 0.1 + + class ConcurrentRunner: + def run(self): + return 12.0, [1], [12.0], [0.9], [0.8], [0.6], [0.5] + + 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) + runner.ca.dataset.gt_data = SimpleNamespace( + path=Path("neighbors_top100k_nq200.parquet"), + row_count=200, + width=100_000, + ) + + metrics = runner._run_perf_case(drop_old=False) + + assert metrics.serial_latency_p50 == 0.5 + assert metrics.conc_latency_p50_list == [0.5] + assert metrics.recall_at == {100: 0.9, 1_000: 0.8} + assert metrics.additional_parameters["ground_truth"] == { + "file": "neighbors_top100k_nq200.parquet", + "query_count": 200, + "width": 100_000, + } diff --git a/tests/test_large_topk_cli.py b/tests/test_large_topk_cli.py new file mode 100644 index 000000000..d935ff951 --- /dev/null +++ b/tests/test_large_topk_cli.py @@ -0,0 +1,93 @@ +import pytest +from click.testing import CliRunner +from pydantic import ValidationError +from pytest import MonkeyPatch + +from vectordb_bench.backend.cases import CaseType +from vectordb_bench.backend.clients.test import cli as test_cli +from vectordb_bench.backend.dataset import DatasetWithSizeType +from vectordb_bench.backend.payload import PayloadProfile +from vectordb_bench.cli import cli as common_cli +from vectordb_bench.models import CaseConfig + + +def invoke_test_command(monkeypatch: MonkeyPatch, args: list[str]): + captured = {} + + def fake_run(tasks, task_label): + captured["task"] = tasks[0] + captured["task_label"] = task_label + + monkeypatch.setattr(common_cli.benchmark_runner, "run", fake_run) + monkeypatch.setattr(common_cli.benchmark_runner, "has_running", lambda: False) + result = CliRunner().invoke(test_cli.Test, args) + return result, captured + + +def test_case_config_rejects_non_positive_k(): + with pytest.raises(ValidationError, match="positive"): + CaseConfig(case_id=CaseType.Performance768D100M, k=0) + + +def test_cli_rejects_non_positive_k(): + result = CliRunner().invoke(test_cli.Test, ["--k", "0", "--dry-run"]) + + assert result.exit_code != 0 + assert "range" in result.output + + +def test_cli_help_describes_laion_large_topk_limit(): + result = CliRunner().invoke(test_cli.Test, ["--help"]) + + assert result.exit_code == 0, result.output + assert "LAION" in result.output + assert "1,000,000" in result.output + + +def test_cli_builds_laion_large_topk_integer_filter_case(monkeypatch: MonkeyPatch): + result, captured = invoke_test_command( + monkeypatch, + [ + "--case-type", + "NewIntFilterPerformanceCase", + "--dataset-with-size-type", + DatasetWithSizeType.LAIONLarge.value, + "--filter-rate", + "0.99", + "--k", + "1000000", + ], + ) + + assert result.exit_code == 0, result.output + case_config = captured["task"].case_config + assert case_config.k == 1_000_000 + assert case_config.case.dataset.data.name == "LAION" + assert case_config.case.filters.filter_rate == 0.99 + + +def test_cli_applies_vector_payload_to_standard_performance_case(monkeypatch: MonkeyPatch): + result, captured = invoke_test_command( + monkeypatch, + [ + "--case-type", + "Performance768D100M", + "--payload-profile", + "vector", + ], + ) + + assert result.exit_code == 0, result.output + case_config = captured["task"].case_config + assert case_config.payload_profile == PayloadProfile.VECTOR + assert case_config.case.payload_profile == PayloadProfile.VECTOR + + +def test_cli_does_not_set_top_level_payload_for_capacity_case(monkeypatch: MonkeyPatch): + result, captured = invoke_test_command( + monkeypatch, + ["--case-type", "CapacityDim128"], + ) + + assert result.exit_code == 0, result.output + assert captured["task"].case_config.payload_profile is None diff --git a/tests/test_large_topk_frontend.py b/tests/test_large_topk_frontend.py new file mode 100644 index 000000000..7afd3cba2 --- /dev/null +++ b/tests/test_large_topk_frontend.py @@ -0,0 +1,232 @@ +import pytest +from pydantic import ValidationError + +from vectordb_bench.backend.cases import CaseType +from vectordb_bench.backend.clients import DB +from vectordb_bench.backend.clients.api import EmptyDBCaseConfig +from vectordb_bench.backend.dataset import DatasetWithSizeType +from vectordb_bench.backend.payload import PayloadProfile +from vectordb_bench.frontend.components.check_results import charts, data +from vectordb_bench.frontend.components.qps_recall import data as qps_recall_data +from vectordb_bench.frontend.components.run_test.caseSelector import payloadProfileSetting +from vectordb_bench.frontend.components.tables import data as table_data +from vectordb_bench.frontend.config.dbCaseConfigs import ( + UI_CASE_CLUSTERS, + UICaseItem, + generate_normal_cases, + get_payload_profile_options, +) +from vectordb_bench.metric import Metric +from vectordb_bench.models import CaseConfig, CaseResult, TaskConfig + + +def _ui_cluster(label: str): + return next(cluster for cluster in UI_CASE_CLUSTERS if cluster.label == label) + + +def test_unfiltered_laion_ui_uses_global_top_k_setting(): + cluster = _ui_cluster("Search Performance Test") + item = next( + item + for item in cluster.uiCaseItems + if item.cases and item.cases[0].case_id == CaseType.Performance768D100M + ).model_copy(deep=True) + + assert item.extra_custom_case_config_inputs == [] + item.payload_profiles = [PayloadProfile.IDS_ONLY, PayloadProfile.VECTOR] + cases = item.get_cases() + + assert [case.k for case in cases] == [100, 100] + assert [case.payload_profile for case in cases] == [PayloadProfile.IDS_ONLY, PayloadProfile.VECTOR] + + +def test_filtered_laion_ui_groups_use_published_top_k_limits(): + cluster = _ui_cluster("New-Int-Filter Search Performance Test") + items = [ + item.model_copy(deep=True) + for item in cluster.uiCaseItems + if item.cases + and item.cases[0].custom_case.get("dataset_with_size_type") == DatasetWithSizeType.LAIONLarge + ] + expected_groups = { + "Large LAION Int-Filter - K up to 1,000,000": [0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 0.98, 0.99], + "Large LAION Int-Filter - K up to 500,000": [0.995], + "Large LAION Int-Filter - K up to 200,000": [0.998], + "Large LAION Int-Filter - K up to 100,000": [0.999], + } + + assert len(items) == 4 + for item in items: + rates = [case.custom_case["filter_rate"] for case in item.cases] + assert rates == expected_groups[item.label] + assert item.extra_custom_case_config_inputs == [] + + +def test_performance_ui_case_expands_selected_payload_profiles(): + item = UICaseItem(cases=generate_normal_cases(CaseType.Performance768D100M)) + item.payload_profiles = [PayloadProfile.IDS_ONLY, PayloadProfile.VECTOR] + + cases = item.get_cases() + + assert [case.payload_profile for case in cases] == [ + PayloadProfile.IDS_ONLY, + PayloadProfile.VECTOR, + ] + assert all(case.case_id == CaseType.Performance768D100M for case in cases) + + +def test_capacity_ui_case_does_not_expand_payload_profiles(): + item = UICaseItem(cases=generate_normal_cases(CaseType.CapacityDim128)) + item.payload_profiles = [PayloadProfile.IDS_ONLY, PayloadProfile.VECTOR] + + cases = item.get_cases() + + assert len(cases) == 1 + assert cases[0].payload_profile is None + + +def test_ui_case_expansion_preserves_payload_conflict_validation(): + item = UICaseItem( + cases=[ + CaseConfig( + case_id=CaseType.CloudPayloadSearchCase, + custom_case={"payload_profile": "vector"}, + ) + ] + ) + item.payload_profiles = [PayloadProfile.IDS_ONLY] + + with pytest.raises(ValidationError, match="conflicts with custom_case"): + item.get_cases() + + +def test_payload_profile_options_require_only_supported_backends(): + assert get_payload_profile_options([DB.Milvus]) == [ + PayloadProfile.IDS_ONLY, + PayloadProfile.VECTOR, + ] + assert get_payload_profile_options([DB.Milvus, DB.ZillizCloud]) == [ + PayloadProfile.IDS_ONLY, + PayloadProfile.VECTOR, + ] + assert get_payload_profile_options([DB.Milvus, DB.Test]) == [PayloadProfile.IDS_ONLY] + assert get_payload_profile_options([]) == [PayloadProfile.IDS_ONLY] + + +def test_payload_profile_setting_records_frontend_selection(): + class FakeContainer: + def __init__(self): + self.options = [] + + def multiselect(self, label, options, default, format_func, key): + assert label == "Return scenario" + assert default == [PayloadProfile.IDS_ONLY] + assert format_func(PayloadProfile.VECTOR) == "Vector payload" + assert key + self.options = options + return options + + def error(self, message): + raise AssertionError(message) + + item = UICaseItem(cases=generate_normal_cases(CaseType.Performance768D100M)) + container = FakeContainer() + + payloadProfileSetting(container, item, [DB.Milvus]) + + assert container.options == [PayloadProfile.IDS_ONLY, PayloadProfile.VECTOR] + assert item.payload_profiles == [PayloadProfile.IDS_ONLY, PayloadProfile.VECTOR] + + +def test_merge_tasks_keeps_results_with_different_k_separate(): + merged, failed = data.mergeTasks( + [ + _case_result(k=100, qps=10), + _case_result(k=1_000_000, qps=5), + ] + ) + + assert failed == {} + assert len(merged) == 2 + assert {item["k"] for item in merged} == {100, 1_000_000} + assert len({item["case_name"] for item in merged}) == 2 + + +def test_merge_tasks_keeps_payload_profiles_separate_for_same_k(): + merged, failed = data.mergeTasks( + [ + _case_result(k=1_000_000, qps=10, payload_profile=PayloadProfile.IDS_ONLY), + _case_result(k=1_000_000, qps=5, payload_profile=PayloadProfile.VECTOR), + ] + ) + + assert failed == {} + assert len(merged) == 2 + assert {item["payload_profile"] for item in merged} == {"ids_only", "vector"} + assert len({item["case_name"] for item in merged}) == 2 + + +def test_build_recall_at_chart_data_normalizes_and_sorts_cutoffs(): + assert hasattr(charts, "buildRecallAtChartData") + chart_data = charts.buildRecallAtChartData( + [ + { + "db": "milvus", + "db_name": "milvus-flat", + "recall_at": {"1000": 0.8, 100: 0.9}, + } + ] + ) + + assert chart_data == [ + {"k": 100, "recall": 0.9, "db": "milvus", "db_name": "milvus-flat"}, + {"k": 1_000, "recall": 0.8, "db": "milvus", "db_name": "milvus-flat"}, + ] + + +def test_qps_recall_data_uses_k_aware_case_name(): + task = _case_result(k=1_000_000, qps=5) + case_name = data.getCaseResultName(task) + + chart_data, failed = qps_recall_data.getChartData( + [task], + dbNames=[task.task_config.db_name], + caseNames=[case_name], + ) + + assert failed == {} + assert chart_data[0]["case_name"] == case_name + assert chart_data[0]["k"] == 1_000_000 + + +def test_results_table_uses_k_aware_case_name(): + rows = table_data.formatData( + [ + _case_result(k=100, qps=10), + _case_result(k=1_000_000, qps=5), + ] + ) + + assert [row["k"] for row in rows] == [100, 1_000_000] + assert len({row["case_name"] for row in rows}) == 2 + + +def _case_result( + *, + k: int, + qps: float, + payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY, +) -> CaseResult: + return CaseResult( + task_config=TaskConfig( + db=DB.Test, + db_config=DB.Test.config_cls(db_label="same-db"), + db_case_config=EmptyDBCaseConfig(), + case_config=CaseConfig( + case_id=CaseType.Performance768D100M, + k=k, + payload_profile=payload_profile, + ), + ), + metrics=Metric(qps=qps, payload_profile=payload_profile.value), + ) diff --git a/tests/test_milvus.py b/tests/test_milvus.py index 2d6163276..d08faf0a2 100644 --- a/tests/test_milvus.py +++ b/tests/test_milvus.py @@ -13,7 +13,7 @@ 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 +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.payload import PayloadProfile from vectordb_bench.interface import BenchMarkRunner @@ -22,6 +22,69 @@ log = logging.getLogger(__name__) +def test_milvus_vector_payload_requests_vector_field_and_returns_ids(): + captured = {} + + def search(**kwargs): + captured.update(kwargs) + return [[{"pk": 1, "vector": [0.1, 0.2]}]] + + db = object.__new__(Milvus) + db.client = SimpleNamespace(search=search) + db.collection_name = "test_collection" + db._vector_field = "vector" + db._primary_field = "pk" + db._scalar_label_field = "label" + db.case_config = SimpleNamespace(search_param=lambda: {"metric_type": "COSINE"}) + db.expr = "" + + result = db.search_embedding([0.1, 0.2], k=3, payload_profile=PayloadProfile.VECTOR) + + assert result == [1] + assert captured["output_fields"] == ["vector"] + + +def _fake_milvus_client(monkeypatch, *, collection_exists=False, properties=None): + client = MagicMock() + client.has_collection.return_value = collection_exists + client.describe_collection.return_value = {"properties": properties or {}} + 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) + return client + + +def _create_milvus_with_collection_properties(monkeypatch, *, collection_exists=False, properties=None): + client = _fake_milvus_client( + monkeypatch, + collection_exists=collection_exists, + properties=properties, + ) + Milvus( + dim=2, + db_config={"uri": "http://example.invalid"}, + db_case_config=SimpleNamespace( + index_param=lambda: {"index_type": "AUTOINDEX", "metric_type": "COSINE", "params": {}}, + ), + collection_properties={"query_mode": "large_topk"}, + ) + return client + + +def test_milvus_creates_collection_properties_before_index(monkeypatch): + client = _create_milvus_with_collection_properties(monkeypatch) + + assert client.create_collection.call_args.kwargs["properties"] == {"query_mode": "large_topk"} + method_names = [method_call[0] for method_call in client.method_calls] + assert method_names.index("create_collection") < method_names.index("create_index") + + +def test_milvus_rejects_existing_collection_with_incompatible_properties(monkeypatch): + with pytest.raises(ValueError, match="incompatible collection properties"): + _create_milvus_with_collection_properties(monkeypatch, collection_exists=True) + + class TestMilvusOptimize: def _milvus( self, @@ -539,3 +602,19 @@ def test_milvus_fts_filter_index_is_conditional(monkeypatch: pytest.MonkeyPatch) db._build_index_params() fields = {call.kwargs["field_name"] for call in params.add_index.call_args_list} assert fields == expected_fields + + +def test_milvus_fts_schema_enables_analyzer_without_text_match(monkeypatch: pytest.MonkeyPatch) -> None: + client = MagicMock() + client.has_collection.return_value = False + schema = MagicMock() + client_type = MagicMock(return_value=client) + client_type.create_schema.return_value = schema + client_type.prepare_index_params.return_value = MagicMock() + monkeypatch.setattr("vectordb_bench.backend.clients.milvus.milvus.MilvusClient", client_type) + + Milvus(dim=0, db_config={}, db_case_config=MilvusFtsConfig(), collection_name="test_collection") + + text_field = next(call for call in schema.add_field.call_args_list if call.args[0] == "text") + assert text_field.kwargs["enable_analyzer"] is True + assert "enable_match" not in text_field.kwargs diff --git a/tests/test_models.py b/tests/test_models.py index d68dd6afb..554c807a0 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,21 +1,58 @@ -import pytest +import json import logging -from vectordb_bench.models import ( - TaskConfig, CaseConfig, - CaseResult, TestResult, - Metric, CaseType -) -from vectordb_bench.backend.clients import ( - DB, - IndexType -) +from pathlib import Path -from vectordb_bench import config +import pytest +from pydantic import ValidationError +from vectordb_bench import config +from vectordb_bench.backend.clients import DB, IndexType +from vectordb_bench.backend.clients.api import EmptyDBCaseConfig +from vectordb_bench.backend.payload import PayloadProfile +from vectordb_bench.models import CaseConfig, CaseResult, CaseType, Metric, TaskConfig, TestResult +from vectordb_bench.restful.format_res import format_results +from vectordb_bench.results import getLeaderboardDataV2 log = logging.getLogger("vectordb_bench") +def test_performance_case_config_applies_top_level_payload_without_mutating_custom_case(): + custom_case = {} + case_config = CaseConfig( + case_id=CaseType.Performance768D100M, + custom_case=custom_case, + payload_profile=PayloadProfile.VECTOR, + ) + + assert case_config.case.payload_profile == PayloadProfile.VECTOR + assert custom_case == {} + + +def test_performance_case_config_payload_round_trip_and_hash_identity(): + ids_only = CaseConfig( + case_id=CaseType.Performance768D100M, + payload_profile=PayloadProfile.IDS_ONLY, + ) + vector = CaseConfig( + case_id=CaseType.Performance768D100M, + payload_profile=PayloadProfile.VECTOR, + ) + + restored = CaseConfig.model_validate(vector.model_dump(mode="json")) + + assert restored.payload_profile == PayloadProfile.VECTOR + assert restored.case.payload_profile == PayloadProfile.VECTOR + assert hash(ids_only) != hash(vector) + + +def test_case_config_rejects_payload_for_non_performance_case(): + with pytest.raises(ValidationError, match="only supported for PerformanceCase"): + CaseConfig( + case_id=CaseType.CapacityDim128, + payload_profile=PayloadProfile.VECTOR, + ) + + class TestModels: @pytest.mark.skip("runs locally") def test_test_result(self): @@ -33,7 +70,7 @@ def test_test_result(self): test_result.flush() with pytest.raises(ValueError): - result = TestResult.read_file('nosuchfile.json') + result = TestResult.read_file("nosuchfile.json") def test_test_result_read_write(self): result_dir = config.RESULTS_LOCAL_DIR @@ -68,3 +105,119 @@ def test_test_result_display(self): log.info(json_file) res = TestResult.read_file(json_file) res.display() + + +def test_old_result_defaults_large_topk_metrics(tmp_path: Path): + test_result = _large_topk_test_result(Metric()) + payload = test_result.model_dump_for_output() + case_config = payload["results"][0]["task_config"]["case_config"] + metrics = payload["results"][0]["metrics"] + case_config.pop("payload_profile", None) + metrics.pop("serial_latency_p50", None) + metrics.pop("conc_latency_p50_list", None) + metrics.pop("recall_at", None) + metrics.pop("payload_profile", None) + metrics.pop("payload_estimated_bytes_per_query", None) + result_file = tmp_path / "old-result.json" + result_file.write_text(json.dumps(payload), encoding="utf-8") + + loaded = TestResult.read_file(result_file) + + assert loaded.results[0].task_config.case_config.payload_profile is None + assert loaded.results[0].task_config.case_config.case.payload_profile == PayloadProfile.IDS_ONLY + assert loaded.results[0].metrics.serial_latency_p50 == 0 + assert loaded.results[0].metrics.conc_latency_p50_list == [] + assert loaded.results[0].metrics.recall_at == {} + assert loaded.results[0].metrics.payload_profile == "ids_only" + assert loaded.results[0].metrics.payload_estimated_bytes_per_query == 0 + + +def test_large_topk_metrics_round_trip_and_convert_serial_p50(tmp_path: Path): + test_result = _large_topk_test_result( + Metric( + serial_latency_p50=0.25, + conc_latency_p50_list=[0.3], + recall_at={100: 0.9, 1_000: 0.8}, + payload_profile="vector", + payload_estimated_bytes_per_query=3_092_000_000, + ), + payload_profile=PayloadProfile.VECTOR, + ) + result_file = tmp_path / "large-topk-result.json" + result_file.write_text(json.dumps(test_result.model_dump_for_output()), encoding="utf-8") + + loaded = TestResult.read_file(result_file, trans_unit=True) + + metrics = loaded.results[0].metrics + assert loaded.results[0].task_config.case_config.payload_profile == PayloadProfile.VECTOR + assert metrics.serial_latency_p50 == 250 + assert metrics.conc_latency_p50_list == [0.3] + assert metrics.recall_at == {100: 0.9, 1_000: 0.8} + assert metrics.payload_profile == "vector" + assert metrics.payload_estimated_bytes_per_query == 3_092_000_000 + + +def test_rest_formatter_exports_large_topk_metrics(): + test_result = _large_topk_test_result( + Metric( + serial_latency_p50=0.25, + conc_latency_p50_list=[0.3], + recall_at={100: 0.9}, + payload_profile="vector", + payload_estimated_bytes_per_query=3_092_000_000, + ), + payload_profile=PayloadProfile.VECTOR, + ) + + formatted = format_results([test_result], task_label="large-topk")[0] + + assert formatted["serial_latency_p50"] == 0.25 + assert formatted["conc_latency_p50_list"] == [0.3] + assert formatted["recall_at"] == {100: 0.9} + assert formatted["payload_profile"] == "vector" + assert formatted["payload_estimated_bytes_per_query"] == 3_092_000_000 + + +def test_legacy_leaderboard_exports_payload_profile(monkeypatch: pytest.MonkeyPatch): + captured = {} + result = _large_topk_test_result( + Metric(qps=1, recall=1, payload_profile="vector"), + payload_profile=PayloadProfile.VECTOR, + ).results[0] + + monkeypatch.setattr(getLeaderboardDataV2, "get_standard_2025_results", lambda: [result]) + monkeypatch.setattr( + getLeaderboardDataV2, + "save_to_json", + lambda data, file_name: captured.setdefault(str(file_name), data), + ) + + getLeaderboardDataV2.main() + + performance_rows = next(rows for rows in captured.values() if rows) + assert performance_rows[0]["payload_profile"] == "vector" + + +def _large_topk_test_result( + metric: Metric, + payload_profile: PayloadProfile | None = None, +) -> TestResult: + return TestResult( + run_id="large-topk", + task_label="large-topk", + results=[ + CaseResult( + task_config=TaskConfig( + db=DB.Test, + db_config=DB.Test.config_cls(), + db_case_config=EmptyDBCaseConfig(), + case_config=CaseConfig( + case_id=CaseType.Performance768D100M, + k=1_000_000, + payload_profile=payload_profile, + ), + ), + metrics=metric, + ) + ], + ) diff --git a/tests/test_multitenant_case.py b/tests/test_multitenant_case.py index 45939c32a..ccdf69412 100644 --- a/tests/test_multitenant_case.py +++ b/tests/test_multitenant_case.py @@ -126,6 +126,9 @@ def test_search_only_zilliz_multitenant_validates_existing_partition_key_schema( calls: list[tuple[str, object]] = [] class ExistingCollectionDB: + def supports_payload_profile(self, payload_profile: PayloadProfile) -> bool: + return payload_profile == PayloadProfile.IDS_ONLY + def supports_multitenant(self) -> bool: return True diff --git a/tests/test_utils.py b/tests/test_utils.py index df3fa6ffe..69a2f7953 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -2,34 +2,47 @@ import logging from vectordb_bench.backend import utils -from vectordb_bench.metric import calc_recall +from vectordb_bench import metric +from vectordb_bench.metric import calc_ndcg, calc_recall, get_ideal_dcg log = logging.getLogger(__name__) + +class NoLinearContains(list): + def __contains__(self, item): + raise AssertionError("recall must not linearly scan ground truth") + + class TestUtils: - @pytest.mark.parametrize("testcases", [ - (1, '1'), - (10, '10'), - (100, '100'), - (1000, '1K'), - (2000, '2K'), - (30_000, '30K'), - (400_000, '400K'), - (5_000_000, '5M'), - (60_000_000, '60M'), - (1_000_000_000, '1B'), - (1_000_000_000_000, '1000B'), - ]) + @pytest.mark.parametrize( + "testcases", + [ + (1, "1"), + (10, "10"), + (100, "100"), + (1000, "1K"), + (2000, "2K"), + (30_000, "30K"), + (400_000, "400K"), + (5_000_000, "5M"), + (60_000_000, "60M"), + (1_000_000_000, "1B"), + (1_000_000_000_000, "1000B"), + ], + ) def test_numerize(self, testcases): t_in, expected = testcases assert expected == utils.numerize(t_in) - @pytest.mark.parametrize("got_expected", [ - ([1, 3, 5, 7, 9, 10], 1.0), - ([11, 12, 13, 14, 15, 16], 0.0), - ([1, 3, 5, 11, 12, 13], 0.5), - ([1, 3, 5], 0.5), - ]) + @pytest.mark.parametrize( + "got_expected", + [ + ([1, 3, 5, 7, 9, 10], 1.0), + ([11, 12, 13, 14, 15, 16], 0.0), + ([1, 3, 5, 11, 12, 13], 0.5), + ([1, 3, 5], 0.5), + ], + ) def test_recall(self, got_expected): got, expected = got_expected ground_truth = [1, 3, 5, 7, 9, 10] @@ -37,14 +50,46 @@ def test_recall(self, got_expected): log.info(f"recall: {res}, expected: {expected}") assert res == expected + def test_recall_does_not_count_duplicate_backend_ids_twice(self): + assert calc_recall(4, [1, 2, 3, 4], [1, 1, 2, 9]) == 0.5 + + def test_recall_does_not_use_linear_list_membership(self): + ground_truth = NoLinearContains([1, 2, 3, 4]) + + assert calc_recall(4, ground_truth, [1, 2, 8, 9]) == 0.5 + + def test_ndcg_preserves_rank_based_semantics_and_deduplicates_results(self): + ground_truth = [10, 20, 30, 40] + got = [30, 10, 30, 99] + expected = (1 / 2 + 1) / get_ideal_dcg(4) + + assert calc_ndcg(ground_truth, got, get_ideal_dcg(4)) == pytest.approx(expected) + + def test_calc_vector_metrics_reports_prefix_recall(self): + assert hasattr(metric, "calc_vector_metrics") + + recall, ndcg, recall_at = metric.calc_vector_metrics( + 4, + [10, 20, 30, 40], + [10, 30, 20, 99], + recall_cutoffs=(2, 4, 8), + ) + + assert recall == 0.75 + assert ndcg == pytest.approx(calc_ndcg([10, 20, 30, 40], [10, 30, 20, 99], get_ideal_dcg(4))) + assert recall_at == {2: 0.5, 4: 0.75} + class TestGetFiles: - @pytest.mark.parametrize("train_count", [ - 1, - 10, - 50, - 100, - ]) + @pytest.mark.parametrize( + "train_count", + [ + 1, + 10, + 50, + 100, + ], + ) def test_train_count(self, train_count): files = utils.compose_train_files(train_count, True) log.info(files) diff --git a/vectordb_bench/backend/assembler.py b/vectordb_bench/backend/assembler.py index fe7f0ddc6..48c77f518 100644 --- a/vectordb_bench/backend/assembler.py +++ b/vectordb_bench/backend/assembler.py @@ -22,9 +22,7 @@ def __init__(self, db_name: str, filter_type: FilterOp): class Assembler: @classmethod def assemble(cls, run_id: str, task: TaskConfig, source: DatasetSource) -> CaseRunner: - c_cls = task.case_config.case_id.case_cls - - c = c_cls(task.case_config.custom_case) + c = task.case_config.case if c.label == CaseLabel.FullTextSearchPerformance and not task.db.init_cls.supports_full_text_search(): msg = f"{task.db.value} does not support full-text search" raise ValueError(msg) diff --git a/vectordb_bench/backend/clients/api.py b/vectordb_bench/backend/clients/api.py index f6e6dfbc2..3f82afd70 100644 --- a/vectordb_bench/backend/clients/api.py +++ b/vectordb_bench/backend/clients/api.py @@ -244,7 +244,8 @@ def need_normalize_cosine(self) -> bool: """Wheather this database need to normalize dataset to support COSINE""" return False - def supports_payload_profile(self, payload_profile: PayloadProfile) -> bool: + @staticmethod + def supports_payload_profile(payload_profile: PayloadProfile) -> bool: return payload_profile == PayloadProfile.IDS_ONLY def has_text_field(self) -> bool: diff --git a/vectordb_bench/backend/clients/milvus/milvus.py b/vectordb_bench/backend/clients/milvus/milvus.py index 42a949f4e..7c639a977 100644 --- a/vectordb_bench/backend/clients/milvus/milvus.py +++ b/vectordb_bench/backend/clients/milvus/milvus.py @@ -35,7 +35,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: PLR0915 + def __init__( # noqa: PLR0912, PLR0915 self, dim: int, db_config: dict, @@ -52,6 +52,7 @@ def __init__( # noqa: PLR0915 self.case_config = db_case_config self.collection_name = collection_name self.with_scalar_labels = with_scalar_labels + collection_properties = kwargs.get("collection_properties", {}) self._scalar_label_field = "label" self._scalar_payload_label_field = self._scalar_label_field @@ -115,7 +116,6 @@ def __init__( # noqa: PLR0915 DataType.VARCHAR, max_length=65535, enable_analyzer=True, - enable_match=True, analyzer_params=analyzer_params, ) schema.add_field(self._sparse_field, DataType.SPARSE_FLOAT_VECTOR) @@ -159,17 +159,31 @@ def __init__( # noqa: PLR0915 log.info(f"{self.name} create collection: {self.collection_name}") index_params = self._build_index_params() + create_kwargs = {} + if collection_properties: + # Large TopK collection properties must be applied before the vector index is created. + create_kwargs["properties"] = collection_properties client.create_collection( collection_name=self.collection_name, schema=schema, num_shards=self.db_config.get("num_shards", 1), consistency_level="Session", + **create_kwargs, ) client.create_index(self.collection_name, index_params) client.load_collection( 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) client.close() @@ -489,7 +503,8 @@ def prepare_filter(self, filters: Filter): msg = f"Not support Filter for Milvus - {filters}" raise ValueError(msg) - def supports_payload_profile(self, payload_profile: PayloadProfile) -> bool: + @staticmethod + def supports_payload_profile(payload_profile: PayloadProfile) -> bool: return payload_profile in { PayloadProfile.IDS_ONLY, PayloadProfile.VECTOR, diff --git a/vectordb_bench/backend/clients/pinecone/pinecone.py b/vectordb_bench/backend/clients/pinecone/pinecone.py index f123ae231..1414001f1 100644 --- a/vectordb_bench/backend/clients/pinecone/pinecone.py +++ b/vectordb_bench/backend/clients/pinecone/pinecone.py @@ -92,7 +92,8 @@ def __getstate__(self): def optimize(self, data_size: int | None = None): pass - def supports_payload_profile(self, payload_profile: PayloadProfile) -> bool: + @staticmethod + def supports_payload_profile(payload_profile: PayloadProfile) -> bool: return payload_profile in { PayloadProfile.IDS_ONLY, PayloadProfile.SCALAR_LABEL, diff --git a/vectordb_bench/backend/clients/turbopuffer/turbopuffer.py b/vectordb_bench/backend/clients/turbopuffer/turbopuffer.py index 918112ede..0b87462d4 100644 --- a/vectordb_bench/backend/clients/turbopuffer/turbopuffer.py +++ b/vectordb_bench/backend/clients/turbopuffer/turbopuffer.py @@ -330,7 +330,8 @@ def insert_embeddings( return 0, e return len(embeddings), None - def supports_payload_profile(self, payload_profile: PayloadProfile) -> bool: + @staticmethod + def supports_payload_profile(payload_profile: PayloadProfile) -> bool: return payload_profile in { PayloadProfile.IDS_ONLY, PayloadProfile.SCALAR_LABEL, diff --git a/vectordb_bench/backend/dataset.py b/vectordb_bench/backend/dataset.py index f797f5bed..bf81c7cac 100644 --- a/vectordb_bench/backend/dataset.py +++ b/vectordb_bench/backend/dataset.py @@ -30,7 +30,7 @@ from . import utils from .clients import MetricType from .data_source import DatasetReader, DatasetSource -from .filter import Filter, FilterOp, non_filter +from .filter import Filter, FilterOp, NewIntFilter, non_filter log = logging.getLogger(__name__) DEFAULT_INSERT_BATCH_SIZE = config.DEFAULT_INSERT_BATCH_SIZE @@ -158,6 +158,125 @@ class LAION(BaseDataset): } +@dataclass(frozen=True) +class SearchDatasetFiles: + test_file: str + gt_file: str + width: int | None = None + query_count: int | None = None + + +LAION_SEARCH_DATASET_FILES = ( + (1_000, SearchDatasetFiles("test.parquet", "neighbors.parquet", width=1_000, query_count=1_000)), + ( + 100_000, + SearchDatasetFiles( + "test_nq200.parquet", + "neighbors_top100k_nq200.parquet", + width=100_000, + query_count=200, + ), + ), + ( + 1_000_000, + SearchDatasetFiles( + "test_nq200.parquet", + "neighbors_top1m_nq200.parquet", + width=1_000_000, + query_count=200, + ), + ), +) + +# Published widths are capped by the population left after applying each ID threshold. +LAION_INT_FILTER_SEARCH_WIDTHS: dict[float, tuple[int, ...]] = { + 0.5: (100_000, 1_000_000), + 0.6: (100_000, 1_000_000), + 0.7: (100_000, 1_000_000), + 0.8: (100_000, 1_000_000), + 0.9: (100_000, 1_000_000), + 0.95: (100_000, 1_000_000), + 0.98: (100_000, 1_000_000), + 0.99: (100_000, 1_000_000), + 0.995: (100_000, 500_000), + 0.998: (100_000, 200_000), + 0.999: (100_000,), +} + + +@dataclass(frozen=True) +class ParquetGroundTruth: + path: pathlib.Path + neighbors_field: str + row_count: int + width: int + + @classmethod + def from_file( + cls, + path: pathlib.Path, + *, + id_field: str, + neighbors_field: str, + expected_query_ids: typing.Sequence[Any], + minimum_width: int, + expected_width: int | None = None, + ) -> "ParquetGroundTruth": + if not path.exists(): + msg = f"No such file: {path}" + raise FileNotFoundError(msg) + + parquet_file = ParquetFile(path, memory_map=True, pre_buffer=False) + schema_names = parquet_file.schema_arrow.names + missing_fields = [field for field in (id_field, neighbors_field) if field not in schema_names] + if missing_fields: + msg = f"Ground truth file {path} is missing fields: {missing_fields}" + raise ValueError(msg) + + query_ids = parquet_file.read(columns=[id_field]).column(0).to_pylist() + if query_ids != list(expected_query_ids): + msg = f"Ground truth query IDs in {path} do not match the selected query file" + raise ValueError(msg) + + row_count = parquet_file.metadata.num_rows + minimum_observed_width = None + observed_rows = 0 + for batch in parquet_file.iter_batches(batch_size=1, columns=[neighbors_field]): + for row in batch.column(0): + if not row.is_valid: + msg = f"Ground truth file {path} contains a null neighbors row" + raise ValueError(msg) + width = len(row.values) + observed_rows += 1 + minimum_observed_width = width if minimum_observed_width is None else min(minimum_observed_width, width) + if expected_width is not None and width != expected_width: + msg = f"Ground truth width {width} in {path} does not match expected width {expected_width}" + raise ValueError(msg) + if width < minimum_width: + msg = f"Ground truth width {width} in {path} is smaller than requested K={minimum_width}" + raise ValueError(msg) + + if observed_rows != row_count or minimum_observed_width is None: + msg = f"Ground truth row count in {path} is invalid: expected {row_count}, read {observed_rows}" + raise ValueError(msg) + + return cls( + path=path, + neighbors_field=neighbors_field, + row_count=row_count, + width=minimum_observed_width, + ) + + def __len__(self) -> int: + return self.row_count + + def iter_rows(self) -> Iterator[Any]: + parquet_file = ParquetFile(self.path, memory_map=True, pre_buffer=False) + for batch in parquet_file.iter_batches(batch_size=1, columns=[self.neighbors_field]): + for row in batch.column(0): + yield row.values.to_numpy(zero_copy_only=False) # noqa: PD011 + + class GIST(BaseDataset): name: str = "GIST" dim: int = 960 @@ -319,7 +438,8 @@ class DatasetManager(BaseModel): data: BaseDataset test_data: list[list[float]] | None = None - gt_data: list[list[int]] | None = None + gt_data: ParquetGroundTruth | list[list[int]] | None = None + search_files: SearchDatasetFiles | None = None scalar_labels: pl.DataFrame | None = None train_files: list[str] = [] reader: DatasetReader | None = None @@ -363,6 +483,7 @@ def prepare( filters: Filter = non_filter, with_train_files: bool = True, with_scalar_labels: bool = False, + k: int | None = None, ) -> bool: """Download the dataset from DatasetSource url = f"{source}/{self.data.dir_name}" @@ -371,15 +492,18 @@ def prepare( source(DatasetSource): S3 or AliyunOSS, default as S3 filters(Filter): combined with dataset's with_gt to compose the correct ground_truth file + k(int | None): requested search depth used to select and validate ground truth Returns: bool: whether the dataset is successfully prepared """ + requested_k = config.K_DEFAULT if k is None else k self.train_files = self.data.train_files if with_train_files else [] gt_file, test_file = None, None if self.data.with_gt: - gt_file, test_file = filters.groundtruth_file, self.data.test_file + self.search_files = self.resolve_search_files(k=requested_k, filters=filters) + gt_file, test_file = self.search_files.gt_file, self.search_files.test_file if self.data.with_remote_resource: download_files = [file for file in self.train_files] @@ -400,13 +524,83 @@ def prepare( self.scalar_labels = self._read_file(self.data.scalar_labels_file) if gt_file is not None and test_file is not None: - self.test_data = self._read_file(test_file)[self.data.test_vector_field].to_list() - self.gt_data = self._read_file(gt_file)[self.data.gt_neighbors_field].to_list() + test_frame = self._read_file(test_file) + if self.search_files.query_count is not None and len(test_frame) != self.search_files.query_count: + msg = ( + f"Query row count {len(test_frame)} in {test_file} does not match " + f"expected count {self.search_files.query_count}" + ) + raise ValueError(msg) + query_ids = test_frame[self.data.test_id_field].to_list() + self.test_data = test_frame[self.data.test_vector_field].to_list() + self.gt_data = ParquetGroundTruth.from_file( + pathlib.Path(self.data_dir, gt_file), + id_field=self.data.gt_id_field, + neighbors_field=self.data.gt_neighbors_field, + expected_query_ids=query_ids, + minimum_width=requested_k, + expected_width=self.search_files.width, + ) log.debug(f"{self.data.name}: available train files {self.train_files}") return True + def max_search_k(self, filters: Filter = non_filter) -> int | None: + if not isinstance(self.data, LAION): + return None + if isinstance(filters, NewIntFilter): + widths = LAION_INT_FILTER_SEARCH_WIDTHS.get(filters.filter_rate) + return widths[-1] if widths is not None else None + if filters.type == FilterOp.NonFilter: + return LAION_SEARCH_DATASET_FILES[-1][0] + return LAION_SEARCH_DATASET_FILES[0][0] + + def resolve_search_files(self, *, k: int, filters: Filter = non_filter) -> SearchDatasetFiles: + if k <= 0: + msg = f"{self.data.name} search K must be positive, got {k}" + raise ValueError(msg) + + if isinstance(self.data, LAION): + max_k = LAION_SEARCH_DATASET_FILES[-1][0] + if k > max_k: + msg = f"LAION supports K up to {max_k:,}, got {k:,}" + raise ValueError(msg) + + if isinstance(filters, NewIntFilter): + widths = LAION_INT_FILTER_SEARCH_WIDTHS.get(filters.filter_rate) + if widths is None: + supported_rates = ", ".join(f"{rate * 100:g}%" for rate in LAION_INT_FILTER_SEARCH_WIDTHS) + msg = f"LAION supported filter rates are: {supported_rates}; got {filters.filter_rate * 100:g}%" + raise ValueError(msg) + if k <= LAION_SEARCH_DATASET_FILES[0][0]: + return SearchDatasetFiles(self.data.test_file, filters.groundtruth_file) + for width in widths: + if k <= width: + width_suffix = f"{width // 1_000_000}m" if width >= 1_000_000 else f"{width // 1_000}k" + return SearchDatasetFiles( + "test_nq200.parquet", + f"neighbors_{filters.int_rate}_top{width_suffix}_nq200.parquet", + width=width, + query_count=200, + ) + msg = ( + f"LAION integer filter {filters.filter_rate * 100:g}% supports K up to " + f"{widths[-1]:,}, got {k:,}" + ) + raise ValueError(msg) + + if filters.type != FilterOp.NonFilter: + if k > LAION_SEARCH_DATASET_FILES[0][0]: + msg = "LAION large-TopK ground truth is published only for integer filters" + raise ValueError(msg) + return SearchDatasetFiles(self.data.test_file, filters.groundtruth_file) + for upper_bound, files in LAION_SEARCH_DATASET_FILES: + if k <= upper_bound: + return files + + return SearchDatasetFiles(self.data.test_file, filters.groundtruth_file) + def _read_file(self, file_name: str) -> pl.DataFrame: """read one file from disk into memory""" log.info(f"Read the entire file into memory: {file_name}") diff --git a/vectordb_bench/backend/runner/mp_runner.py b/vectordb_bench/backend/runner/mp_runner.py index b5a7fe17d..51a5422c6 100644 --- a/vectordb_bench/backend/runner/mp_runner.py +++ b/vectordb_bench/backend/runner/mp_runner.py @@ -177,6 +177,7 @@ def _run_all_concurrencies_mem_efficient(self): conc_latency_p99_list = [] conc_latency_p95_list = [] conc_latency_avg_list = [] + conc_latency_p50_list = [] try: for conc in self.concurrencies: with mp.Manager() as m: @@ -197,9 +198,7 @@ def _run_all_concurrencies_mem_efficient(self): start = time.perf_counter() all_count = sum([r.result()[0] for r in future_iter]) latencies = sum([r.result()[2] for r in future_iter], start=[]) - latency_p99 = np.percentile(latencies, 99) - latency_p95 = np.percentile(latencies, 95) - latency_avg = np.mean(latencies) + latency_p99, latency_p95, latency_p50, latency_avg = self._latency_summary(latencies) cost = time.perf_counter() - start qps = round(all_count / cost, 4) @@ -208,6 +207,7 @@ def _run_all_concurrencies_mem_efficient(self): conc_latency_p99_list.append(latency_p99) conc_latency_p95_list.append(latency_p95) conc_latency_avg_list.append(latency_avg) + conc_latency_p50_list.append(latency_p50) log.info(f"End search in concurrency {conc}: dur={cost}s, total_count={all_count}, qps={qps}") if qps > max_qps: @@ -233,6 +233,7 @@ def _run_all_concurrencies_mem_efficient(self): conc_latency_p99_list, conc_latency_p95_list, conc_latency_avg_list, + conc_latency_p50_list, ) def _wait_for_queue_fill(self, q: Queue, size: int): @@ -254,30 +255,42 @@ def run(self) -> float: def stop(self) -> None: pass - def _aggregate_latency_stats(self, res: list) -> tuple[float, float, float]: + @staticmethod + def _latency_summary(latencies: list[float]) -> tuple[float, float, float, float]: + if not latencies: + return 0, 0, 0, 0 + return ( + float(np.percentile(latencies, 99)), + float(np.percentile(latencies, 95)), + float(np.percentile(latencies, 50)), + float(np.mean(latencies)), + ) + + def _aggregate_latency_stats(self, res: list) -> tuple[float, float, float, float]: """Aggregate latency stats from worker processes. Returns: - tuple: (p99, p95, avg) latencies in seconds + tuple: (p99, p95, p50, avg) latencies in seconds """ latency_stats_list = [r[2] for r in res if r[2] and r[2].get("count", 0) > 0] if not latency_stats_list: - return 0, 0, 0 + return 0, 0, 0, 0 total_query_count = sum(stats["count"] for stats in latency_stats_list) if total_query_count == 0: - return 0, 0, 0 + return 0, 0, 0, 0 # Use max for conservative percentile estimate latency_p99 = max(stats["p99"] for stats in latency_stats_list) latency_p95 = max(stats["p95"] for stats in latency_stats_list) + latency_p50 = max(stats["p50"] for stats in latency_stats_list) # Weighted average latency_avg = sum(stats["avg"] * stats["count"] for stats in latency_stats_list) / total_query_count - return latency_p99, latency_p95, latency_avg + return latency_p99, latency_p95, latency_p50, latency_avg def run_by_dur(self, duration: int) -> tuple[float, float]: """ @@ -287,7 +300,7 @@ def run_by_dur(self, duration: int) -> tuple[float, float]: """ return self._run_by_dur(duration) - def _run_by_dur(self, duration: int) -> tuple[float, float, list, list, list, list, list]: + def _run_by_dur(self, duration: int) -> tuple[float, float, list, list, list, list, list, list]: """ Returns: float: largest qps @@ -297,6 +310,7 @@ def _run_by_dur(self, duration: int) -> tuple[float, float, list, list, list, li list: p99 latencies at each concurrency list: p95 latencies at each concurrency list: avg latencies at each concurrency + list: p50 latencies at each concurrency """ max_qps = 0 conc_num_list = [] @@ -304,6 +318,7 @@ def _run_by_dur(self, duration: int) -> tuple[float, float, list, list, list, li conc_latency_p99_list = [] conc_latency_p95_list = [] conc_latency_avg_list = [] + conc_latency_p50_list = [] try: for conc in self.concurrencies: with mp.Manager() as m: @@ -329,23 +344,26 @@ def _run_by_dur(self, duration: int) -> tuple[float, float, list, list, list, li res = [r.result() for r in future_iter] all_success_count = sum([r[0] for r in res]) all_failed_count = sum([r[1] for r in res]) - failed_rate = all_failed_count / (all_failed_count + all_success_count) + total_count = all_failed_count + all_success_count + failed_rate = all_failed_count / total_count if total_count > 0 else 0 cost = time.perf_counter() - start qps = round(all_success_count / cost, 4) - latency_p99, latency_p95, latency_avg = self._aggregate_latency_stats(res) + latency_p99, latency_p95, latency_p50, latency_avg = self._aggregate_latency_stats(res) conc_num_list.append(conc) conc_qps_list.append(qps) conc_latency_p99_list.append(latency_p99) conc_latency_p95_list.append(latency_p95) conc_latency_avg_list.append(latency_avg) + conc_latency_p50_list.append(latency_p50) log.info( f"End search in concurrency {conc}: dur={cost}s, failed_rate={failed_rate}, " f"all_success_count={all_success_count}, all_failed_count={all_failed_count}, qps={qps}, " - f"p99={latency_p99:.4f}s, p95={latency_p95:.4f}s, avg={latency_avg:.4f}s", + f"p99={latency_p99:.4f}s, p95={latency_p95:.4f}s, " + f"p50={latency_p50:.4f}s, avg={latency_avg:.4f}s", ) if qps > max_qps: max_qps = qps @@ -371,6 +389,7 @@ def _run_by_dur(self, duration: int) -> tuple[float, float, list, list, list, li conc_latency_p99_list, conc_latency_p95_list, conc_latency_avg_list, + conc_latency_p50_list, ) def search_by_dur( @@ -380,7 +399,7 @@ def search_by_dur( Returns: int: successful requests count int: failed requests count - dict: latency statistics with p99, p95, avg, count (computed via HDR Histogram) + dict: latency statistics with p99, p95, p50, avg, count (computed via HDR Histogram) """ # sync all process q.put(1) @@ -403,7 +422,7 @@ def search_by_dur( self._search_once(test_data[idx]) success_count += 1 latency_us = int((time.perf_counter() - s) * US_TO_SECONDS) - histogram.record_value(min(latency_us, HDR_HISTOGRAM_MAX_US)) + histogram.record_value(max(HDR_HISTOGRAM_MIN_US, min(latency_us, HDR_HISTOGRAM_MAX_US))) except Exception as e: failed_cnt += 1 # reduce log @@ -431,6 +450,7 @@ def search_by_dur( latency_stats = { "p99": histogram.get_value_at_percentile(99) / US_TO_SECONDS, "p95": histogram.get_value_at_percentile(95) / US_TO_SECONDS, + "p50": histogram.get_value_at_percentile(50) / US_TO_SECONDS, "avg": histogram.get_mean_value() / US_TO_SECONDS, "count": histogram.get_total_count(), } diff --git a/vectordb_bench/backend/runner/read_write_runner.py b/vectordb_bench/backend/runner/read_write_runner.py index 8293128fe..31b5fa5be 100644 --- a/vectordb_bench/backend/runner/read_write_runner.py +++ b/vectordb_bench/backend/runner/read_write_runner.py @@ -102,10 +102,10 @@ def run_search(self, perc: int): log.info("Search after write - Serial search start") test_time = round(time.perf_counter(), 4) res, ssearch_dur = self.serial_search_runner.run() - recall, ndcg, p99_latency, p95_latency = res + recall, ndcg, p99_latency, p95_latency, p50_latency, _ = res log.info( f"Search after write - Serial search - recall={recall}, ndcg={ndcg}, " - f"p99={p99_latency}, p95={p95_latency}, dur={ssearch_dur:.4f}", + f"p99={p99_latency}, p95={p95_latency}, p50={p50_latency}, dur={ssearch_dur:.4f}", ) log.info( f"Search after wirte - Conc search start, dur for each conc={self.read_dur_after_write}", @@ -266,10 +266,11 @@ def wait_next_target(start: int, target_batch: int) -> bool: log.info(f"[{target_batch}/{total_batch}] Serial search - {perc}% start") res, ssearch_dur = self.serial_search_runner.run() ssearch_dur = round(ssearch_dur, 4) - recall, ndcg, p99_latency, p95_latency = res + recall, ndcg, p99_latency, p95_latency, p50_latency, _ = res log.info( f"[{target_batch}/{total_batch}] Serial search - {perc}% done, " - f"recall={recall}, ndcg={ndcg}, p99={p99_latency}, p95={p95_latency}, dur={ssearch_dur}" + f"recall={recall}, ndcg={ndcg}, p99={p99_latency}, p95={p95_latency}, " + f"p50={p50_latency}, dur={ssearch_dur}" ) each_conc_search_dur = self.get_each_conc_search_dur( diff --git a/vectordb_bench/backend/runner/serial_runner.py b/vectordb_bench/backend/runner/serial_runner.py index 3aba25c1c..09aef08cc 100644 --- a/vectordb_bench/backend/runner/serial_runner.py +++ b/vectordb_bench/backend/runner/serial_runner.py @@ -5,6 +5,7 @@ import random import time import traceback +from collections.abc import Iterator import numpy as np @@ -14,7 +15,7 @@ from vectordb_bench.backend.workload import WorkloadKind from ... import config -from ...metric import calc_mrr_fts, calc_ndcg, calc_ndcg_fts, calc_recall, calc_recall_fts, get_ideal_dcg +from ...metric import calc_mrr_fts, calc_ndcg_fts, calc_recall_fts, calc_vector_metrics from ...models import LoadTimeoutError from .. import utils from ..clients import api @@ -210,19 +211,49 @@ def _get_db_search_res( return results - def search(self, args: tuple[list, list[list[int]] | list[dict[str, int]]]) -> tuple[float, ...]: + def _validate_ground_truth(self, test_data: list, ground_truth: object | None) -> None: + if not self.measure_recall or ground_truth is None: + return + if len(test_data) != len(ground_truth): + msg = f"Search query count {len(test_data)} does not match ground truth row count {len(ground_truth)}" + raise ValueError(msg) + if self._use_fts_metrics: + return + + source_width = getattr(ground_truth, "width", None) + if source_width is not None: + if source_width < self.k: + msg = f"Ground truth width {source_width} is smaller than requested K={self.k}" + raise ValueError(msg) + return + + for row_idx, row in enumerate(ground_truth): + if len(row) < self.k: + msg = f"Ground truth width {len(row)} at row {row_idx} is smaller than requested K={self.k}" + raise ValueError(msg) + + @staticmethod + def _iter_ground_truth(ground_truth: object) -> Iterator: + if hasattr(ground_truth, "iter_rows"): + return ground_truth.iter_rows() + return iter(ground_truth) + + def search(self, args: tuple[list, object]) -> tuple[float, ...]: log.info(f"{mp.current_process().name:14} start search the entire test_data to get recall and latency") + test_data, ground_truth = args + self._validate_ground_truth(test_data, ground_truth) + ground_truth_iter = self._iter_ground_truth(ground_truth) if ground_truth is not None else None + with self.db.init(): self.db.prepare_filter(self.filters) - test_data, ground_truth = args - ideal_dcg = None if self._use_fts_metrics else get_ideal_dcg(self.k) log.debug(f"test dataset size: {len(test_data)}") log.debug(f"ground truth size: {len(ground_truth) if ground_truth is not None else 0}") latencies, recalls, ndcgs, mrrs = [], [], [], [] + recall_at_samples = {} tenant_rng = random.Random(0) - for idx, emb in enumerate(test_data): + for emb in test_data: tenant = ( self.tenant_labels[tenant_rng.randrange(len(self.tenant_labels))] if self.workload_kind == WorkloadKind.VECTOR and self.tenant_labels @@ -238,14 +269,17 @@ def search(self, args: tuple[list, list[list[int]] | list[dict[str, int]]]) -> t latencies.append(time.perf_counter() - s) if self.measure_recall and ground_truth is not None: - gt = ground_truth[idx] + gt = next(ground_truth_iter) if self._use_fts_metrics: recalls.append(calc_recall_fts(self.k, gt, results)) ndcgs.append(calc_ndcg_fts(self.k, gt, results)) mrrs.append(calc_mrr_fts(self.k, gt, results)) else: - recalls.append(calc_recall(self.k, gt[: self.k], results)) - ndcgs.append(calc_ndcg(gt[: self.k], results, ideal_dcg)) + recall, ndcg, recall_at = calc_vector_metrics(self.k, gt, results) + recalls.append(recall) + ndcgs.append(ndcg) + for cutoff, value in recall_at.items(): + recall_at_samples.setdefault(cutoff, []).append(value) else: recalls.append(0) ndcgs.append(0) @@ -264,6 +298,7 @@ def search(self, args: tuple[list, list[list[int]] | list[dict[str, int]]]) -> t cost = round(np.sum(latencies), 4) p99 = round(np.percentile(latencies, 99), 4) p95 = round(np.percentile(latencies, 95), 4) + p50 = round(np.percentile(latencies, 50), 4) if self._use_fts_metrics: avg_mrr = round(np.mean(mrrs), 4) log.info( @@ -279,6 +314,7 @@ def search(self, args: tuple[list, list[list[int]] | list[dict[str, int]]]) -> t ) return (avg_recall, avg_ndcg, avg_mrr, p99, p95) + avg_recall_at = {cutoff: round(np.mean(values), 4) for cutoff, values in recall_at_samples.items()} log.info( f"{mp.current_process().name:14} search entire test_data: " f"cost={cost}s, " @@ -287,9 +323,11 @@ def search(self, args: tuple[list, list[list[int]] | list[dict[str, int]]]) -> t f"avg_ndcg={avg_ndcg}, " f"avg_latency={avg_latency}, " f"p99={p99}, " - f"p95={p95}" + f"p95={p95}, " + f"p50={p50}, " + f"recall_at={avg_recall_at}" ) - return (avg_recall, avg_ndcg, p99, p95) + return (avg_recall, avg_ndcg, p99, p95, p50, avg_recall_at) def _run_in_subprocess(self) -> tuple[float, ...]: with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor: diff --git a/vectordb_bench/backend/task_runner.py b/vectordb_bench/backend/task_runner.py index 4b9314e9d..ea32a0953 100644 --- a/vectordb_bench/backend/task_runner.py +++ b/vectordb_bench/backend/task_runner.py @@ -8,6 +8,7 @@ import numpy as np +from .. import config from ..base import BaseModel from ..metric import Metric from ..models import PerformanceTimeoutError, TaskConfig, TaskStage @@ -28,6 +29,9 @@ from .workload import WorkloadKind log = logging.getLogger(__name__) +# Milvus and Zilliz Cloud cap topK at this value unless the collection opts into Large TopK mode. +MILVUS_DEFAULT_TOPK_LIMIT = 16_384 +LARGE_TOPK_QUERY_MODE_DBS = frozenset({DB.Milvus, DB.ZillizCloud}) class RunningStatus(Enum): @@ -83,6 +87,7 @@ def load_reuse_key(self) -> tuple | None: self._collection_name_hash_key(), self._dataset_hash_key(), self.config.insert_batch_size, + self._hashable_value(self._collection_properties()), self.ca.with_scalar_labels, self.ca.is_multitenant, self._multitenant_routing_hash_key(), @@ -181,6 +186,25 @@ def workload_kind(self) -> WorkloadKind: def is_fts(self) -> bool: return self.workload_kind == WorkloadKind.FULL_TEXT + def _collection_properties(self, *, log_selection: bool = False) -> dict[str, str]: + requested_k = self.config.case_config.k or config.K_DEFAULT + if ( + self.config.db not in LARGE_TOPK_QUERY_MODE_DBS + or self.ca.label != CaseLabel.Performance + or requested_k <= MILVUS_DEFAULT_TOPK_LIMIT + ): + return {} + + # Large TopK mode must be applied at collection creation, before the vector index is created. + if log_selection: + log.info( + "%s requested K=%d exceeds the default TopK limit %d; using query_mode=large_topk", + self.config.db.value, + requested_k, + MILVUS_DEFAULT_TOPK_LIMIT, + ) + return {"query_mode": "large_topk"} + def init_db(self, drop_old: bool = True) -> None: db_cls = self.config.db.init_cls # Compose a compact, case-unique collection/table name for Doris to avoid cross-case interference @@ -205,6 +229,9 @@ def init_db(self, drop_old: bool = True) -> None: extra_db_kwargs["fts_filter_enabled"] = self.ca.filters.type != FilterOp.NonFilter if self.config.db is DB.AWSOpenSearch: extra_db_kwargs["insert_batch_size"] = self.config.insert_batch_size + collection_properties = self._collection_properties(log_selection=True) + if collection_properties: + extra_db_kwargs["collection_properties"] = collection_properties self.db = db_cls( dim=getattr(self.ca.dataset.data, "dim", 0), @@ -215,9 +242,22 @@ def init_db(self, drop_old: bool = True) -> None: **extra_db_kwargs, ) + def _validate_vector_payload_profile(self) -> None: + if self.ca.label != CaseLabel.Performance or self.is_fts: + return + if not self.config.db.init_cls.supports_payload_profile(self.ca.payload_profile): + msg = f"{self.config.db_name} does not support payload_profile={self.ca.payload_profile.value}" + raise NotImplementedError(msg) + def _pre_run(self, drop_old: bool = True): try: self._validate_cloud_cold_latency_config(drop_old) + requested_k = self.config.case_config.k or config.K_DEFAULT + ground_truth_k = ( + requested_k + if self.ca.label == CaseLabel.Performance and getattr(self.ca, "measure_recall", True) + else config.K_DEFAULT + ) creates_multitenant_collection = ( TaskStage.DROP_OLD in self.config.stages or TaskStage.LOAD in self.config.stages ) @@ -238,6 +278,9 @@ def _pre_run(self, drop_old: bool = True): self.init_db(drop_old) return + self._validate_vector_payload_profile() + if self.ca.dataset.data.with_gt: + self.ca.dataset.resolve_search_files(k=ground_truth_k, filters=self.ca.filters) self.init_db(drop_old) if self.ca.is_multitenant and self.db is not None: if not self.db.supports_multitenant(): @@ -251,6 +294,7 @@ def _pre_run(self, drop_old: bool = True): filters=self.ca.filters, with_train_files=TaskStage.LOAD in self.config.stages, with_scalar_labels=self.ca.with_scalar_labels, + k=ground_truth_k, ) except ModuleNotFoundError as e: log.warning(f"pre run case error: please install client for db: {self.config.db}, error={e}") @@ -363,6 +407,7 @@ def _run_perf_case(self, drop_old: bool = True) -> Metric: m.conc_latency_p99_list, m.conc_latency_p95_list, m.conc_latency_avg_list, + m.conc_latency_p50_list, ) = search_results if TaskStage.SEARCH_SERIAL in self.config.stages: cooldown = self.config.case_config.concurrency_search_config.serial_cooldown @@ -375,7 +420,21 @@ def _run_perf_case(self, drop_old: bool = True) -> Metric: if self.is_fts: m.recall, m.ndcg, m.mrr, m.serial_latency_p99, m.serial_latency_p95 = search_results else: - m.recall, m.ndcg, m.serial_latency_p99, m.serial_latency_p95 = search_results + ( + m.recall, + m.ndcg, + m.serial_latency_p99, + m.serial_latency_p95, + m.serial_latency_p50, + m.recall_at, + ) = search_results + gt_data = getattr(self.ca.dataset, "gt_data", None) + if gt_data is not None and hasattr(gt_data, "path"): + m.additional_parameters["ground_truth"] = { + "file": gt_data.path.name, + "query_count": gt_data.row_count, + "width": gt_data.width, + } 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( @@ -533,8 +592,8 @@ def _serial_search(self) -> tuple[float, ...]: calculate the recall, serial_latency_p99, serial_latency_p95 Returns: - tuple[float, ...]: vector cases return recall, ndcg, p99, p95; - FTS cases return recall, p99, p95. + tuple[float, ...]: vector cases return recall, ndcg, p99, p95, p50, recall_at; + FTS cases return recall, p99, p95, p50. """ try: if self.serial_search_runner is None: diff --git a/vectordb_bench/cli/cli.py b/vectordb_bench/cli/cli.py index fbc8cd0a2..8baa524fa 100644 --- a/vectordb_bench/cli/cli.py +++ b/vectordb_bench/cli/cli.py @@ -18,7 +18,7 @@ from yaml import load from .. import config -from ..backend.cases import FTS_FILTER_RATES +from ..backend.cases import FTS_FILTER_RATES, PerformanceCase, type2case from ..backend.clients import DB from ..backend.clients.api import IndexType, MetricType from ..backend.dataset import DatasetWithSizeType, FtsDatasetWithSizeType @@ -329,6 +329,13 @@ def apply_fts_cli_db_case_params( return db_case_config.model_copy(update=updates) +def get_case_payload_profile(parameters: dict[str, Any]) -> PayloadProfile | None: + case_type = CaseType[parameters["case_type"]] + if not issubclass(type2case[case_type], PerformanceCase): + return None + return PayloadProfile(parameters["payload_profile"]) + + def select_cli_db_case_config( db: DB, db_case_config: DBCaseConfig, @@ -483,10 +490,10 @@ class CommonTypedDict(TypedDict): int, click.option( "--k", - type=int, + type=click.IntRange(min=1), default=config.K_DEFAULT, show_default=True, - help="K value for number of nearest neighbors to search", + help="Number of nearest neighbors. LAION 100M selects tiered GT automatically up to 1,000,000.", ), ] concurrency_duration: Annotated[ @@ -669,7 +676,7 @@ class CommonTypedDict(TypedDict): click.option( "--payload-profile", type=click.Choice([profile.value for profile in PayloadProfile]), - help="Response payload profile for payload and FTS cases", + help="Response payload profile for vector performance, cloud payload, and FTS cases", default="ids_only", show_default=True, ), @@ -967,6 +974,7 @@ def run( db_case_config=select_cli_db_case_config(db, db_case_config, parameters["case_type"], parameters), case_config=CaseConfig( case_id=CaseType[parameters["case_type"]], + payload_profile=get_case_payload_profile(parameters), k=parameters["k"], concurrency_search_config=ConcurrencySearchConfig( concurrency_duration=parameters["concurrency_duration"], diff --git a/vectordb_bench/frontend/components/check_results/charts.py b/vectordb_bench/frontend/components/check_results/charts.py index d36cc5fc6..b88ebda4e 100644 --- a/vectordb_bench/frontend/components/check_results/charts.py +++ b/vectordb_bench/frontend/components/check_results/charts.py @@ -45,6 +45,43 @@ def drawChart(data, st, key_prefix: str): key = f"{key_prefix}-{metric}" drawMetricChart(data, metric, container, key=key) + drawRecallAtChart(data, st.container(), key=f"{key_prefix}-recall-at") + + +def buildRecallAtChartData(data): + chart_data = [] + for case_data in data: + for cutoff, recall in case_data.get("recall_at", {}).items(): + chart_data.append( + { + "k": int(cutoff), + "recall": recall, + "db": case_data["db"], + "db_name": case_data["db_name"], + } + ) + return sorted(chart_data, key=lambda item: (item["k"], item["db_name"])) + + +def drawRecallAtChart(data, st, key: str): + chart_data = buildRecallAtChartData(data) + if len(chart_data) == 0: + return + + fig = px.line( + chart_data, + x="k", + y="recall", + color="db_name", + markers=True, + hover_data={"db": True, "db_name": True}, + title="Recall by K", + ) + fig.update_xaxes(type="log", title_text="K") + fig.update_yaxes(range=[0, 1], title_text="Recall") + fig.update_layout(showlegend=True, legend_title_text="") + st.plotly_chart(fig, width="stretch", key=key) + def getLabelToShapeMap(data): labelIndexMap = {} diff --git a/vectordb_bench/frontend/components/check_results/data.py b/vectordb_bench/frontend/components/check_results/data.py index 289453f7a..cbe615ed7 100644 --- a/vectordb_bench/frontend/components/check_results/data.py +++ b/vectordb_bench/frontend/components/check_results/data.py @@ -1,9 +1,30 @@ from collections import defaultdict from dataclasses import asdict + +from vectordb_bench import config +from vectordb_bench.backend.cases import CaseType, PerformanceCase +from vectordb_bench.backend.payload import PayloadProfile from vectordb_bench.metric import QPS_METRIC, isLowerIsBetterMetric from vectordb_bench.models import CaseResult, ResultLabel +def getCaseResultName(task: CaseResult) -> str: + case_config = task.task_config.case_config + case = case_config.case + details = [] + if case_config.k is not None and case_config.k != config.K_DEFAULT: + details.append(f"K={case_config.k:,}") + if ( + isinstance(case, PerformanceCase) + and case.case_id != CaseType.CloudPayloadSearchCase + and case.payload_profile != PayloadProfile.IDS_ONLY + ): + details.append(f"Payload={case.payload_profile.value}") + if not details: + return case.name + return f"{case.name} ({', '.join(details)})" + + def getChartData( tasks: list[CaseResult], dbNames: list[str], @@ -20,9 +41,7 @@ def getFilterTasks( caseNames: list[str], ) -> list[CaseResult]: filterTasks = [ - task - for task in tasks - if task.task_config.db_name in dbNames and task.task_config.case_config.case_name in caseNames + task for task in tasks if task.task_config.db_name in dbNames and getCaseResultName(task) in caseNames ] return filterTasks @@ -35,13 +54,15 @@ def mergeTasks(tasks: list[CaseResult]): db_label = task.task_config.db_config.db_label or "" version = task.task_config.db_config.version or "" case = task.task_config.case_config.case - case_name = case.name + case_name = getCaseResultName(task) + k = task.task_config.case_config.k dataset_name = case.dataset.data.full_name filter_rate = case.filter_rate - dbCaseMetricsMap[db_name][case.name] = { + dbCaseMetricsMap[db_name][case_name] = { "db": db, "db_label": db_label, "version": version, + "k": k, "dataset_name": dataset_name, "filter_rate": filter_rate, "metrics": mergeMetrics( @@ -62,6 +83,7 @@ def mergeTasks(tasks: list[CaseResult]): db = metricInfo["db"] db_label = metricInfo["db_label"] version = metricInfo["version"] + k = metricInfo["k"] label = metricInfo["label"] dataset_name = metricInfo["dataset_name"] filter_rate = metricInfo["filter_rate"] @@ -74,6 +96,7 @@ def mergeTasks(tasks: list[CaseResult]): "dataset_name": dataset_name, "filter_rate": filter_rate, "version": version, + "k": k, "case_name": case_name, "metricsSet": set(metrics.keys()), **metrics, diff --git a/vectordb_bench/frontend/components/check_results/filters.py b/vectordb_bench/frontend/components/check_results/filters.py index e42a63fd6..310b1ffd1 100644 --- a/vectordb_bench/frontend/components/check_results/filters.py +++ b/vectordb_bench/frontend/components/check_results/filters.py @@ -1,7 +1,7 @@ -from vectordb_bench.backend.cases import Case, CaseLabel +from vectordb_bench.backend.cases import CaseLabel from vectordb_bench.backend.dataset import DatasetWithSizeType from vectordb_bench.backend.filter import FilterOp -from vectordb_bench.frontend.components.check_results.data import getChartData +from vectordb_bench.frontend.components.check_results.data import getCaseResultName, getChartData from vectordb_bench.frontend.components.check_results.expanderStyle import ( initSidebarExanderStyle, ) @@ -63,8 +63,6 @@ def getShowDbsAndCases(st, result: list[CaseResult], filter_type: FilterOp) -> t case_results = [res for res in result if res.task_config.case_config.case.filters.type == filter_type] allDbNames = list(set({res.task_config.db_name for res in case_results})) allDbNames.sort() - allCases: list[Case] = [res.task_config.case_config.case for res in case_results] - # DB Filter dbFilterContainer = st.container() showDBNames = filterView( @@ -76,14 +74,29 @@ def getShowDbsAndCases(st, result: list[CaseResult], filter_type: FilterOp) -> t showCaseNames = [] # Handle FTS cases separately - fts_cases = [case for case in allCases if case.label == CaseLabel.FullTextSearchPerformance] - non_fts_cases = [case for case in allCases if case.label != CaseLabel.FullTextSearchPerformance] + fts_case_results = [ + result + for result in case_results + if result.task_config.case_config.case.label == CaseLabel.FullTextSearchPerformance + ] + non_fts_case_results = [ + result + for result in case_results + if result.task_config.case_config.case.label != CaseLabel.FullTextSearchPerformance + ] if filter_type == FilterOp.NonFilter: - allCaseNameSet = set({case.name for case in allCases}) - allCaseNames = [case_name for case_name in CASE_NAME_ORDER if case_name in allCaseNameSet] + [ - case_name for case_name in allCaseNameSet if case_name not in CASE_NAME_ORDER - ] + display_to_base = { + getCaseResultName(result): result.task_config.case_config.case.name for result in case_results + } + case_order = {case_name: idx for idx, case_name in enumerate(CASE_NAME_ORDER)} + allCaseNames = sorted( + display_to_base, + key=lambda display_name: ( + case_order.get(display_to_base[display_name], len(case_order)), + display_name, + ), + ) # Case Filter caseFilterContainer = st.container() @@ -105,9 +118,15 @@ def getShowDbsAndCases(st, result: list[CaseResult], filter_type: FilterOp) -> t optionLables=[v.value for v in datasetWithSizeTypes], ) datasets = [dataset_with_size_type.get_manager() for dataset_with_size_type in showDatasetWithSizeTypes] - showCaseNames = list(set([case.name for case in non_fts_cases if case.dataset in datasets])) + showCaseNames = list( + { + getCaseResultName(result) + for result in non_fts_case_results + if result.task_config.case_config.case.dataset in datasets + } + ) # Add FTS cases - fts_case_names = [case.name for case in fts_cases] + fts_case_names = [getCaseResultName(result) for result in fts_case_results] showCaseNames.extend(fts_case_names) return showDBNames, showCaseNames diff --git a/vectordb_bench/frontend/components/concurrent/charts.py b/vectordb_bench/frontend/components/concurrent/charts.py index 5369d5912..7b4f0ced2 100644 --- a/vectordb_bench/frontend/components/concurrent/charts.py +++ b/vectordb_bench/frontend/components/concurrent/charts.py @@ -25,6 +25,11 @@ def drawChartsByCase(allData, showCaseNames: list[str], st, latency_type: str): if "conc_latency_p95_list" in caseData and 0 <= i < len(caseData["conc_latency_p95_list"]) else 0 ), + "latency_p50": ( + caseData["conc_latency_p50_list"][i] * 1000 + if "conc_latency_p50_list" in caseData and 0 <= i < len(caseData["conc_latency_p50_list"]) + else 0 + ), "latency_avg": ( caseData["conc_latency_avg_list"][i] * 1000 if 0 <= i < len(caseData["conc_latency_avg_list"]) diff --git a/vectordb_bench/frontend/components/qps_recall/charts.py b/vectordb_bench/frontend/components/qps_recall/charts.py index a44c51a36..30281a626 100644 --- a/vectordb_bench/frontend/components/qps_recall/charts.py +++ b/vectordb_bench/frontend/components/qps_recall/charts.py @@ -69,7 +69,7 @@ def drawlinechart(st, data: list[object], metric, key: str): new_data, new_remain_data = drawBestperformance(data, y, group) unique_db_names = list(set(item["db_name"] for item in new_data + new_remain_data)) - colors = plt.cm.get_cmap("tab10", len(unique_db_names)) + colors = plt.get_cmap("tab10", len(unique_db_names)) color_map = { db: f"rgb({int(colors(i)[0] * 255)}, {int(colors(i)[1] * 255)}, {int(colors(i)[2] * 255)})" diff --git a/vectordb_bench/frontend/components/qps_recall/data.py b/vectordb_bench/frontend/components/qps_recall/data.py index b4cbcb1b5..eb363cd5c 100644 --- a/vectordb_bench/frontend/components/qps_recall/data.py +++ b/vectordb_bench/frontend/components/qps_recall/data.py @@ -1,7 +1,7 @@ from collections import defaultdict from dataclasses import asdict from vectordb_bench.backend.filter import FilterOp -from vectordb_bench.frontend.components.check_results.data import getFilterTasks +from vectordb_bench.frontend.components.check_results.data import getCaseResultName, getFilterTasks from vectordb_bench.frontend.components.check_results.filters import getShowDbsAndCases, getshownResults from vectordb_bench.models import CaseResult, ResultLabel, TestResult @@ -33,7 +33,8 @@ def getChartData( db_label = task.task_config.db_config.db_label or "" version = task.task_config.db_config.version or "" case = task.task_config.case_config.case - case_name = case.name + case_name = getCaseResultName(task) + k = task.task_config.case_config.k dataset_name = case.dataset.data.full_name filter_rate = case.filter_rate metrics = asdict(task.metrics) @@ -47,6 +48,7 @@ def getChartData( "dataset_name": dataset_name, "filter_rate": filter_rate, "version": version, + "k": k, "case_name": case_name, "metricsSet": set(metrics.keys()), **metrics, diff --git a/vectordb_bench/frontend/components/run_test/caseSelector.py b/vectordb_bench/frontend/components/run_test/caseSelector.py index 685a38ec2..3414a3f16 100644 --- a/vectordb_bench/frontend/components/run_test/caseSelector.py +++ b/vectordb_bench/frontend/components/run_test/caseSelector.py @@ -1,6 +1,9 @@ +from collections import defaultdict +from typing import Any + from vectordb_bench.backend.clients import DB +from vectordb_bench.backend.payload import PayloadProfile from vectordb_bench.frontend.components.run_test.inputWidget import inputWidget -from collections import defaultdict from vectordb_bench.frontend.config.dbCaseConfigs import ( UI_CASE_CLUSTERS, UICaseItem, @@ -8,6 +11,7 @@ get_case_config_inputs, get_custom_case_cluter, get_custom_streaming_case_cluster, + get_payload_profile_options, get_selectable_case_items, ) from vectordb_bench.frontend.config.styles import ( @@ -15,10 +19,34 @@ CHECKBOX_INDENT, DB_CASE_CONFIG_SETTING_COLUMNS, ) - from vectordb_bench.frontend.utils import addHorizontalLine from vectordb_bench.models import CaseConfig +PAYLOAD_PROFILE_LABELS = { + PayloadProfile.IDS_ONLY: "IDs only", + PayloadProfile.VECTOR: "Vector payload", +} + + +def payloadProfileSetting(container: Any, ui_case_item: UICaseItem, active_dbs: list[DB]) -> None: + if not ui_case_item.supports_payload_profiles: + return + + options = get_payload_profile_options(active_dbs) + selected = [profile for profile in ui_case_item.payload_profiles if profile in options] + if not selected: + selected = [PayloadProfile.IDS_ONLY] + backend_key = "-".join(sorted(db.name for db in active_dbs)) or "none" + ui_case_item.payload_profiles = container.multiselect( + "Return scenario", + options=options, + default=selected, + format_func=PAYLOAD_PROFILE_LABELS.__getitem__, + key=f"payload-profile-{ui_case_item.label}-{backend_key}", + ) + if not ui_case_item.payload_profiles: + container.error("Select at least one return scenario.") + def caseSelector(st, activedDbList: list[DB]): st.markdown( @@ -66,6 +94,7 @@ def caseItemCheckbox(st, dbToCaseClusterConfigs, uiCaseItem: UICaseItem, actived caseConfigSetting(st.container(), uiCaseItem) if selected: + payloadProfileSetting(st.container(), uiCaseItem, activedDbList) dbCaseConfigSetting(st.container(), dbToCaseClusterConfigs, uiCaseItem, activedDbList) return uiCaseItem.get_cases() if selected else [] diff --git a/vectordb_bench/frontend/components/run_test/submitTask.py b/vectordb_bench/frontend/components/run_test/submitTask.py index 93fd280c3..efe9b9ee9 100644 --- a/vectordb_bench/frontend/components/run_test/submitTask.py +++ b/vectordb_bench/frontend/components/run_test/submitTask.py @@ -3,6 +3,7 @@ import streamlit as st from vectordb_bench import config +from vectordb_bench.backend.dataset import DatasetManager from vectordb_bench.frontend.config import styles from vectordb_bench.interface import benchmark_runner from vectordb_bench.models import TaskConfig @@ -35,7 +36,39 @@ def taskLabelInput(container): return cols[0].text_input("task_label", defaultTaskLabel, label_visibility="collapsed") -def advancedSettings(container): +def get_max_search_k(tasks: list[TaskConfig]) -> int | None: + limits = [] + for task in tasks: + case = task.case_config.case + if not isinstance(case.dataset, DatasetManager): + continue + max_k = case.dataset.max_search_k(case.filters) + if max_k is not None: + limits.append(max_k) + return min(limits, default=None) + + +def apply_run_settings( + tasks: list[TaskConfig], + *, + k: int, + concurrencies: list[int], + concurrency_duration: int, + load_concurrency: int, +) -> None: + for task in tasks: + case = task.case_config.case + if isinstance(case.dataset, DatasetManager) and case.dataset.data.with_gt: + case.dataset.resolve_search_files(k=k, filters=case.filters) + + for task in tasks: + task.case_config.k = k + task.case_config.concurrency_search_config.num_concurrency = concurrencies + task.case_config.concurrency_search_config.concurrency_duration = concurrency_duration + task.load_concurrency = load_concurrency + + +def advancedSettings(container, tasks: list[TaskConfig]): cols = container.columns([1, 2]) index_already_exists = cols[0].checkbox("Index already exists", value=False) cols[1].caption("if selected, inserting and building will be skipped.") @@ -45,7 +78,16 @@ def advancedSettings(container): cols[1].caption("if selected, the dataset will be downloaded from Aliyun OSS shanghai, default AWS S3 aws-us-west.") cols = container.columns([1, 2]) - k = cols[0].number_input("k", min_value=1, value=100, label_visibility="collapsed") + max_k = get_max_search_k(tasks) + k_config = dict( + min_value=1, + value=config.K_DEFAULT, + label_visibility="collapsed", + key=f"search-k-{max_k or 'unbounded'}", + ) + if max_k is not None: + k_config["max_value"] = max_k + k = cols[0].number_input("k", **k_config) cols[1].caption("K value for number of nearest neighbors to search") cols = container.columns([1, 2]) @@ -69,23 +111,24 @@ def advancedSettings(container): def controlPanel(container, tasks: list[TaskConfig], taskLabel, isAllValid): index_already_exists, use_aliyun, k, concurrentInput, concurrency_duration, load_concurrency = advancedSettings( - container + container, tasks ) def runHandler(): - benchmark_runner.set_drop_old(not index_already_exists) - try: concurrentInput_list = [int(item.strip()) for item in concurrentInput.split(",")] - except ValueError: - container.write("please input correct number") + apply_run_settings( + tasks, + k=k, + concurrencies=concurrentInput_list, + concurrency_duration=concurrency_duration, + load_concurrency=load_concurrency, + ) + except ValueError as exc: + container.write(str(exc)) return None - for task in tasks: - task.case_config.k = k - task.case_config.concurrency_search_config.num_concurrency = concurrentInput_list - task.case_config.concurrency_search_config.concurrency_duration = concurrency_duration - task.load_concurrency = load_concurrency + benchmark_runner.set_drop_old(not index_already_exists) benchmark_runner.set_download_address(use_aliyun) benchmark_runner.run(tasks, taskLabel) diff --git a/vectordb_bench/frontend/components/tables/data.py b/vectordb_bench/frontend/components/tables/data.py index 5a229e6e4..69534c0c1 100644 --- a/vectordb_bench/frontend/components/tables/data.py +++ b/vectordb_bench/frontend/components/tables/data.py @@ -1,7 +1,10 @@ from dataclasses import asdict + +import pandas as pd + +from vectordb_bench.frontend.components.check_results.data import getCaseResultName from vectordb_bench.interface import benchmark_runner from vectordb_bench.models import CaseResult, ResultLabel -import pandas as pd def getNewResults(): @@ -32,7 +35,8 @@ def formatData(caseResults: list[CaseResult]): { "db": db, "db_label": db_label, - "case_name": case.name, + "case_name": getCaseResultName(caseResult), + "k": case_config.k, "dataset": dataset, "filter_rate": filter_rate, **metrics, diff --git a/vectordb_bench/frontend/config/dbCaseConfigs.py b/vectordb_bench/frontend/config/dbCaseConfigs.py index 84ed76f8e..95f47586b 100644 --- a/vectordb_bench/frontend/config/dbCaseConfigs.py +++ b/vectordb_bench/frontend/config/dbCaseConfigs.py @@ -1,18 +1,32 @@ -from enum import IntEnum, Enum import typing -from pydantic import BaseModel -from vectordb_bench.backend.cases import CaseLabel, CaseType +from enum import Enum, IntEnum + +from pydantic import BaseModel, Field + +from vectordb_bench.backend.cases import CaseLabel, CaseType, PerformanceCase from vectordb_bench.backend.clients import DB from vectordb_bench.backend.clients.api import IndexType, MetricType, SQType -from vectordb_bench.backend.dataset import DatasetWithSizeType, FtsDatasetWithSizeType +from vectordb_bench.backend.dataset import ( + LAION_INT_FILTER_SEARCH_WIDTHS, + DatasetWithSizeType, + FtsDatasetWithSizeType, +) +from vectordb_bench.backend.payload import PayloadProfile from vectordb_bench.frontend.components.custom.getCustomConfig import get_custom_configs - from vectordb_bench.models import CaseConfig, CaseConfigParamType MAX_STREAMLIT_INT = (1 << 53) - 1 DB_LIST = [d for d in DB if d != DB.Test] FTS_SUPPORTED_DBS = {DB.Milvus, DB.ZillizCloud, DB.ElasticCloud, DB.OSSOpenSearch, DB.Vespa, DB.TurboPuffer} +VECTOR_PAYLOAD_SUPPORTED_DBS = {DB.Milvus, DB.ZillizCloud} + + +def get_payload_profile_options(active_dbs: list[DB]) -> list[PayloadProfile]: + profiles = [PayloadProfile.IDS_ONLY] + if active_dbs and all(db in VECTOR_PAYLOAD_SUPPORTED_DBS for db in active_dbs): + profiles.append(PayloadProfile.VECTOR) + return profiles class Delimiter(Enum): @@ -56,6 +70,7 @@ class UICaseItem(BaseModel): supportedDbs: list[DB] | None = None extra_custom_case_config_inputs: list[ConfigInput] = [] tmp_custom_config: dict = dict() + payload_profiles: list[PayloadProfile] = Field(default_factory=lambda: [PayloadProfile.IDS_ONLY]) def __init__( self, @@ -91,20 +106,29 @@ def __init__( def __hash__(self) -> int: return hash(self.key if self.key else self.label) + @property + def supports_payload_profiles(self) -> bool: + return bool(self.cases) and all(isinstance(case.case, PerformanceCase) for case in self.cases) + def get_cases(self) -> list[CaseConfig]: - # return self.cases - if len(self.extra_custom_case_config_inputs) == 0: - return self.cases - cases = [ - CaseConfig( - case_id=c.case_id, - k=c.k, - concurrency_search_config=c.concurrency_search_config, - custom_case={**c.custom_case, **self.tmp_custom_config}, - ) - for c in self.cases + cases = self.cases + if self.extra_custom_case_config_inputs: + cases = [ + CaseConfig( + case_id=case.case_id, + k=case.k, + concurrency_search_config=case.concurrency_search_config, + custom_case={**(case.custom_case or {}), **self.tmp_custom_config}, + ) + for case in cases + ] + if not self.supports_payload_profiles: + return cases + return [ + CaseConfig.model_validate({**case.model_dump(), "payload_profile": payload_profile}) + for case in cases + for payload_profile in self.payload_profiles ] - return cases def supports_dbs(self, dbs: list[DB]) -> bool: if self.supportedDbs is None: @@ -289,8 +313,12 @@ def generate_label_filter_cases(dataset_with_size_type: DatasetWithSizeType) -> ] -def generate_int_filter_cases(dataset_with_size_type: DatasetWithSizeType) -> list[CaseConfig]: - filter_rates = dataset_with_size_type.get_manager().data.scalar_int_rates +def generate_int_filter_cases( + dataset_with_size_type: DatasetWithSizeType, + filter_rates: typing.Iterable[float] | None = None, +) -> list[CaseConfig]: + if filter_rates is None: + filter_rates = dataset_with_size_type.get_manager().data.scalar_int_rates return [ CaseConfig( case_id=CaseType.NewIntFilterPerformanceCase, @@ -300,6 +328,21 @@ def generate_int_filter_cases(dataset_with_size_type: DatasetWithSizeType) -> li ] +def generate_laion_large_topk_filter_items() -> list[UICaseItem]: + rates_by_max_k: dict[int, list[float]] = {} + for filter_rate, widths in LAION_INT_FILTER_SEARCH_WIDTHS.items(): + rates_by_max_k.setdefault(widths[-1], []).append(filter_rate) + + return [ + UICaseItem( + label=f"Large LAION Int-Filter - K up to {max_k:,}", + description="Filter rates: " + ", ".join(f"{rate * 100:g}%" for rate in filter_rates), + cases=generate_int_filter_cases(DatasetWithSizeType.LAIONLarge, filter_rates), + ) + for max_k, filter_rates in rates_by_max_k.items() + ] + + UI_CASE_CLUSTERS: list[UICaseItemCluster] = [ UICaseItemCluster( label="Search Performance Test", @@ -332,7 +375,8 @@ def generate_int_filter_cases(dataset_with_size_type: DatasetWithSizeType) -> li ), UICaseItemCluster( label="New-Int-Filter Search Performance Test", - uiCaseItems=[ + uiCaseItems=generate_laion_large_topk_filter_items() + + [ UICaseItem( label=f"Int-Filter Search Performance Test - {dataset_with_size_type.value}", description=( diff --git a/vectordb_bench/frontend/pages/concurrent.py b/vectordb_bench/frontend/pages/concurrent.py index 20b273044..7911a5ef9 100644 --- a/vectordb_bench/frontend/pages/concurrent.py +++ b/vectordb_bench/frontend/pages/concurrent.py @@ -60,7 +60,10 @@ def check_conc_data(res: TestResult): getResults(resultesContainer, "vectordb_bench_concurrent") # main - latency_type = st.radio("Latency Type", options=["latency_p99", "latency_p95", "latency_avg"]) + latency_type = st.radio( + "Latency Type", + options=["latency_p99", "latency_p95", "latency_p50", "latency_avg"], + ) drawChartsByCase(shownData, showCaseNames, st.container(), latency_type=latency_type) # footer diff --git a/vectordb_bench/metric.py b/vectordb_bench/metric.py index 540cefd28..2f8c8e89f 100644 --- a/vectordb_bench/metric.py +++ b/vectordb_bench/metric.py @@ -1,10 +1,15 @@ import logging +from collections.abc import Iterable, Sequence from dataclasses import dataclass, field +from functools import cache +from itertools import islice import numpy as np log = logging.getLogger(__name__) +RECALL_CUTOFFS = (100, 1_000, 10_000, 100_000, 1_000_000) + @dataclass class Metric: @@ -22,13 +27,16 @@ class Metric: qps: float = 0.0 serial_latency_p99: float = 0.0 serial_latency_p95: float = 0.0 + serial_latency_p50: float = 0.0 recall: float = 0.0 + recall_at: dict[int, float] = field(default_factory=dict) ndcg: float = 0.0 mrr: float = 0.0 conc_num_list: list[int] = field(default_factory=list) conc_qps_list: list[float] = field(default_factory=list) conc_latency_p99_list: list[float] = field(default_factory=list) conc_latency_p95_list: list[float] = field(default_factory=list) + conc_latency_p50_list: list[float] = field(default_factory=list) conc_latency_avg_list: list[float] = field(default_factory=list) payload_profile: str = "ids_only" payload_estimated_bytes_per_query: int = 0 @@ -63,6 +71,7 @@ class Metric: LOAD_DURATION_METRIC = "load_duration" SERIAL_LATENCY_P99_METRIC = "serial_latency_p99" SERIAL_LATENCY_P95_METRIC = "serial_latency_p95" +SERIAL_LATENCY_P50_METRIC = "serial_latency_p50" MAX_LOAD_COUNT_METRIC = "max_load_count" QPS_METRIC = "qps" RECALL_METRIC = "recall" @@ -71,6 +80,7 @@ class Metric: LOAD_DURATION_METRIC: "s", SERIAL_LATENCY_P99_METRIC: "ms", SERIAL_LATENCY_P95_METRIC: "ms", + SERIAL_LATENCY_P50_METRIC: "ms", MAX_LOAD_COUNT_METRIC: "K", QURIES_PER_DOLLAR_METRIC: "K", } @@ -79,6 +89,7 @@ class Metric: LOAD_DURATION_METRIC, SERIAL_LATENCY_P99_METRIC, SERIAL_LATENCY_P95_METRIC, + SERIAL_LATENCY_P50_METRIC, ] metric_order = [ @@ -87,6 +98,7 @@ class Metric: LOAD_DURATION_METRIC, SERIAL_LATENCY_P99_METRIC, SERIAL_LATENCY_P95_METRIC, + SERIAL_LATENCY_P50_METRIC, MAX_LOAD_COUNT_METRIC, ] @@ -95,33 +107,75 @@ def isLowerIsBetterMetric(metric: str) -> bool: return metric in lower_is_better_metrics -def calc_recall(count: int, ground_truth: list[int], got: list[int]) -> float: - recalls = np.zeros(count) - for i, result in enumerate(got): - if result in ground_truth: - recalls[i] = 1 +def calc_recall(count: int, ground_truth: Iterable[int], got: Iterable[int]) -> float: + if count <= 0: + return 0.0 - return np.mean(recalls) + ground_truth_ids = set(ground_truth) + hits = {result for result in islice(got, count) if result in ground_truth_ids} + return len(hits) / count +@cache def get_ideal_dcg(k: int): - ideal_dcg = 0 - for i in range(k): - ideal_dcg += 1 / np.log2(i + 2) + if k <= 0: + return 0.0 + + ranks = np.arange(2, k + 2, dtype=np.float64) + return float(np.sum(1 / np.log2(ranks))) - return ideal_dcg +def _build_ground_truth_ranks(ground_truth: Iterable[int]) -> dict[int, int]: + ranks = {} + for rank, neighbor_id in enumerate(ground_truth): + ranks.setdefault(neighbor_id, rank) + return ranks -def calc_ndcg(ground_truth: list[int], got: list[int], ideal_dcg: float) -> float: - dcg = 0 - ground_truth = list(ground_truth) + +def _calc_ndcg_from_ranks(ground_truth_ranks: dict[int, int], got: Iterable[int], ideal_dcg: float) -> float: + if ideal_dcg <= 0: + return 0.0 + + dcg = 0.0 for got_id in set(got): - if got_id in ground_truth: - idx = ground_truth.index(got_id) - dcg += 1 / np.log2(idx + 2) + rank = ground_truth_ranks.get(got_id) + if rank is not None: + dcg += 1 / np.log2(rank + 2) return dcg / ideal_dcg +def calc_ndcg(ground_truth: Iterable[int], got: Iterable[int], ideal_dcg: float) -> float: + return _calc_ndcg_from_ranks(_build_ground_truth_ranks(ground_truth), got, ideal_dcg) + + +def calc_vector_metrics( + count: int, + ground_truth: Iterable[int], + got: Sequence[int], + recall_cutoffs: Iterable[int] = RECALL_CUTOFFS, +) -> tuple[float, float, dict[int, float]]: + if count <= 0: + return 0.0, 0.0, {} + + ground_truth_ranks = _build_ground_truth_ranks(islice(ground_truth, count)) + unique_results = set(islice(got, count)) + recall = len(unique_results.intersection(ground_truth_ranks)) / count + ndcg = _calc_ndcg_from_ranks(ground_truth_ranks, unique_results, get_ideal_dcg(count)) + + recall_at = {} + for cutoff in recall_cutoffs: + if cutoff <= 0 or cutoff > count: + continue + hits = { + result + for result in islice(got, cutoff) + if (rank := ground_truth_ranks.get(result)) is not None and rank < cutoff + } + recall_at[cutoff] = len(hits) / cutoff + + return recall, ndcg, recall_at + + def _positive_fts_qrels(ground_truth: dict[str, int] | list[int] | list[str]) -> dict[str, int]: if isinstance(ground_truth, dict): return {str(doc_id): int(rel) for doc_id, rel in ground_truth.items() if int(rel) > 0} diff --git a/vectordb_bench/models.py b/vectordb_bench/models.py index 1c4924ed1..90ce5445d 100644 --- a/vectordb_bench/models.py +++ b/vectordb_bench/models.py @@ -6,13 +6,12 @@ from typing import Any, ClassVar, Self import ujson -from pydantic import PositiveInt, model_validator +from pydantic import PositiveInt, field_validator, model_validator -from vectordb_bench.backend.cases import type2case from vectordb_bench.backend.dataset import DatasetWithSizeMap from . import config -from .backend.cases import Case, CaseType +from .backend.cases import Case, CaseType, PerformanceCase, type2case from .backend.clients import ( DB, DBCaseConfig, @@ -20,6 +19,7 @@ EmptyDBCaseConfig, ) from .backend.clients.api import IndexType +from .backend.payload import PayloadProfile from .base import BaseModel from .metric import Metric @@ -213,9 +213,34 @@ class CaseConfig(BaseModel): case_id: CaseType custom_case: dict | None = None + payload_profile: PayloadProfile | None = None k: int | None = config.K_DEFAULT concurrency_search_config: ConcurrencySearchConfig = ConcurrencySearchConfig() + @field_validator("k") + @classmethod + def validate_k(cls, value: int | None) -> int | None: + if value is not None and value <= 0: + msg = f"K must be positive, got {value}" + raise ValueError(msg) + return value + + @model_validator(mode="after") + def validate_payload_profile(self) -> Self: + if self.payload_profile is None: + return self + + case_cls = type2case[self.case_id] + if not issubclass(case_cls, PerformanceCase): + msg = "Top-level payload_profile is only supported for PerformanceCase cases" + raise ValueError(msg) # noqa: TRY004 + + legacy_profile = (self.custom_case or {}).get("payload_profile") + if legacy_profile is not None and PayloadProfile(legacy_profile) != self.payload_profile: + msg = "Top-level payload_profile conflicts with custom_case payload_profile" + raise ValueError(msg) + return self + ''' @property def k(self): @@ -233,7 +258,10 @@ def __hash__(self) -> int: @property def case(self) -> Case: - return self.case_id.case_cls(self.custom_case) + custom_case = dict(self.custom_case or {}) + if self.payload_profile is not None: + custom_case["payload_profile"] = self.payload_profile + return self.case_id.case_cls(custom_case or None) @property def case_name(self) -> str: @@ -340,7 +368,7 @@ def _redact_sensitive_fields(cls, value: Any) -> Any: return { key: ( "**********" - if key.lower() in cls.sensitive_output_fields and item + if isinstance(key, str) and key.lower() in cls.sensitive_output_fields and item else cls._redact_sensitive_fields(item) ) for key, item in value.items() @@ -451,7 +479,7 @@ def get_case_config(case_config: CaseConfig) -> dict[CaseConfig]: return case_config @classmethod - def read_file(cls, full_path: pathlib.Path, trans_unit: bool = False) -> Self: + def read_file(cls, full_path: pathlib.Path, trans_unit: bool = False) -> Self: # noqa: PLR0912 if not full_path.exists(): msg = f"No such file: {full_path}" raise ValueError(msg) @@ -519,6 +547,12 @@ def read_file(cls, full_path: pathlib.Path, trans_unit: bool = False) -> Self: ) elif "serial_latency_p99" in metrics: metrics["serial_latency_p95"] = 0.0 + + if "serial_latency_p50" in metrics: + cur_latency_p50 = metrics["serial_latency_p50"] + metrics["serial_latency_p50"] = ( + cur_latency_p50 * 1000 if cur_latency_p50 > 0 else cur_latency_p50 + ) return TestResult.model_validate(test_result) def display(self, dbs: list[DB] | None = None): diff --git a/vectordb_bench/restful/format_res.py b/vectordb_bench/restful/format_res.py index c4af0d579..76900b05f 100644 --- a/vectordb_bench/restful/format_res.py +++ b/vectordb_bench/restful/format_res.py @@ -32,13 +32,18 @@ class FormatResult(BaseModel): qps: float = 0 serial_latency_p99: float = 0 serial_latency_p95: float = 0 + serial_latency_p50: float = 0 recall: float = 0 + recall_at: dict[int, float] = {} ndcg: float = 0 mrr: float = 0 + payload_profile: str = "ids_only" + payload_estimated_bytes_per_query: int = 0 conc_num_list: list[int] = [] conc_qps_list: list[float] = [] conc_latency_p99_list: list[float] = [] conc_latency_p95_list: list[float] = [] + conc_latency_p50_list: list[float] = [] conc_latency_avg_list: list[float] = [] diff --git a/vectordb_bench/results/getLeaderboardDataV2.py b/vectordb_bench/results/getLeaderboardDataV2.py index 188d9876f..96fe8fa71 100644 --- a/vectordb_bench/results/getLeaderboardDataV2.py +++ b/vectordb_bench/results/getLeaderboardDataV2.py @@ -50,6 +50,7 @@ def main(): "latency": round(latency, 4), "recall": round(recall, 4), "filter_ratio": round(filter_ratio, 3), + "payload_profile": metrics.payload_profile, } ) else: