feat: support configurable NQ for concurrent Milvus searches

README.md

@@ -138,6 +138,32 @@ vectordbbench milvusautoindex \
   <other options>
 ```
 
+Use `--nq` (default: `1`) to set the number of query vectors per concurrent search request.
+Milvus and Zilliz Cloud support batching for single-tenant vector performance cases,
+including filtered searches. This example reuses a collection containing the selected
+dataset and a matching HNSW index:
+
+```shell
+vectordbbench milvushnsw \
+  --uri http://localhost:19530 --collection-name VDBBench \
+  --case-type Performance768D1M --m 16 --ef-construction 128 --ef-search 128 \
+  --k 10 --nq 2 \
+  --skip-drop-old --skip-load
+```
+
+YAML uses `nq: 2`. The value is saved in `task_config.case_config.nq` in result JSON.
+Result views distinguish non-default NQ values in case names so different batch sizes stay separate.
+Omitting `--num-concurrency` uses the configured concurrency list. Each request contains
+exactly `nq` vectors, wrapping through the query dataset as needed.
+
+- `qps` and `conc_qps_list` count successful query vectors per second. Successful request throughput
+  is QPS divided by `nq`. Failed batches or responses missing a query's result list contribute zero.
+- `conc_latency_*` measures the whole successful batch request.
+- `serial_latency_*` and recall use NQ=1 and concurrency=1.
+
+Other backends and streaming, full-text, cold-latency and multi-tenant cases reject
+`nq > 1` before data loading.
+
 Use `--note` or `--note-file` to preserve deployment, resource, client, network, and constraint context in each result JSON under `task_config.db_config.note`. The options are mutually exclusive. Prefer `--note-file` for structured or multiline context, and never include credentials, tokens, or sensitive connection details.
 
 ```shell

tests/test_case_runner_reuse.py

@@ -1,3 +1,6 @@
+from unittest.mock import Mock
+
+import pytest
 from pydantic import SecretStr
 
 from vectordb_bench import config
@@ -66,6 +69,27 @@ def make_runner(
     )
 
 
+def test_nq_rejects_unsupported_cases_before_data_loading(monkeypatch):
+    runner = make_runner(db=DB.Test)
+    runner.config.case_config.nq = 2
+    dataset = Mock()
+    init_db = Mock()
+    monkeypatch.setattr(runner.ca, "dataset", dataset)
+    monkeypatch.setattr(CaseRunner, "init_db", init_db)
+
+    with pytest.raises(NotImplementedError, match="nq > 1"):
+        runner._pre_run()
+
+    assert dataset.mock_calls == []
+    init_db.assert_not_called()
+
+    for case_id in (CaseType.FTSBm25Performance, CaseType.CloudMultiTenantSearchCase):
+        runner = make_runner(db=DB.ZillizCloud, case_id=case_id)
+        runner.config.case_config.nq = 2
+        with pytest.raises(ValueError, match="nq > 1"):
+            runner._validate_nq()
+
+
 def assert_not_reusable(left: CaseRunner, right: CaseRunner) -> None:
     assert left != right
     assert right != left

tests/test_large_topk_frontend.py

@@ -152,6 +152,17 @@ def test_merge_tasks_keeps_results_with_different_k_separate():
     assert len({item["case_name"] for item in merged}) == 2
 
 
+def test_results_keep_different_nq_separate():
+    tasks = [_case_result(k=100, qps=10), _case_result(k=100, qps=20, nq=2)]
+    base_name = tasks[0].task_config.case_config.case.name
+    names = [base_name, f"{base_name} (NQ=2)"]
+
+    assert [data.getCaseResultName(task) for task in tasks] == names
+    merged, failed = data.mergeTasks(tasks)
+    assert failed == {}
+    assert {row["case_name"]: row["qps"] for row in merged} == {names[0]: 10, names[1]: 20}
+
+
 def test_merge_tasks_keeps_payload_profiles_separate_for_same_k():
     merged, failed = data.mergeTasks(
         [
@@ -215,6 +226,7 @@ def _case_result(
     *,
     k: int,
     qps: float,
+    nq: int = 1,
     payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY,
 ) -> CaseResult:
     return CaseResult(
@@ -225,6 +237,7 @@ def _case_result(
             case_config=CaseConfig(
                 case_id=CaseType.Performance768D100M,
                 k=k,
+                nq=nq,
                 payload_profile=payload_profile,
             ),
         ),

tests/test_milvus.py

@@ -4,6 +4,7 @@
 """
 
 import logging
+from contextlib import nullcontext
 from types import SimpleNamespace
 from unittest.mock import MagicMock, call
 
@@ -16,6 +17,7 @@
 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.backend.runner.mp_runner import MultiProcessingSearchRunner
 from vectordb_bench.interface import BenchMarkRunner
 from vectordb_bench.models import CaseConfig, TaskConfig
 
@@ -44,6 +46,47 @@ def search(**kwargs):
     assert captured["output_fields"] == ["vector"]
 
 
+def test_milvus_batch_search_counts_successful_queries(monkeypatch):
+    queries = [[0.0], [1.0], [2.0]]
+    # Complete zero-hit batch, incomplete batch, SDK failure.
+    responses = [[[], []], [[]], RuntimeError("search failed")]
+    db = object.__new__(Milvus)
+    db.client = SimpleNamespace(search=MagicMock(side_effect=responses))
+    db.collection_name = "test_collection"
+    db._vector_field = "vector"
+    db._primary_field = "pk"
+    db.case_config = SimpleNamespace(search_param=lambda: {"metric_type": "COSINE"})
+    db.expr = ""
+    db.init = lambda: nullcontext()
+    db.prepare_filter = MagicMock()
+    runner = MultiProcessingSearchRunner(db, queries, nq=2, duration=3)
+    monkeypatch.setattr("vectordb_bench.backend.runner.mp_runner.random.randint", lambda _a, _b: 0)
+    # Each SDK request takes one simulated second, regardless of clock reads.
+    clock = MagicMock()
+
+    def perf_counter():
+        assert clock.call_count < 100, "Search loop did not finish"
+        return float(db.client.search.call_count)
+
+    clock.side_effect = perf_counter
+    monkeypatch.setattr("vectordb_bench.backend.runner.mp_runner.time", SimpleNamespace(perf_counter=clock))
+
+    count, _, latencies = runner.search(queries, MagicMock(), MagicMock())
+    assert count == 2
+    assert latencies == [1.0]
+    assert db.client.search.call_count == 3
+    expected_batches = [queries[:2], [queries[2], queries[0]], queries[1:]]
+    assert [c.kwargs["data"] for c in db.client.search.call_args_list] == expected_batches
+
+    db.client.search.reset_mock()
+    db.client.search.side_effect = responses
+    count, failed_count, stats = runner.search_by_dur(3, queries, MagicMock(), MagicMock())
+    assert (count, failed_count) == (2, 4)
+    assert stats["count"] == 1
+    assert stats["avg"] == pytest.approx(1.0, rel=0.01)
+    assert [c.kwargs["data"] for c in db.client.search.call_args_list] == expected_batches
+
+
 def _fake_milvus_client(monkeypatch, *, collection_exists=False, properties=None):
     client = MagicMock()
     client.has_collection.return_value = collection_exists

tests/test_milvus_zilliz_cli.py

@@ -4,6 +4,7 @@
 
 from vectordb_bench.backend.clients.milvus import cli as milvus_cli
 from vectordb_bench.backend.clients.zilliz_cloud import cli as zilliz_cli
+from vectordb_bench.cli import cli as common_cli
 
 
 def test_milvus_cli_builds_shared_connection_config() -> None:
@@ -131,3 +132,34 @@ def fake_run(**kwargs):
 
     assert result.exit_code == 0, result.output
     assert captured["db_case_config"].use_partition_key is True
+
+
+def test_milvus_autoindex_cli_nq(monkeypatch: MonkeyPatch) -> None:
+    captured = {}
+
+    def fake_run(tasks, task_label):
+        captured["task"] = tasks[0]
+
+    monkeypatch.setattr(common_cli.benchmark_runner, "run", fake_run)
+    monkeypatch.setattr(common_cli.benchmark_runner, "has_running", lambda: False)
+
+    result = CliRunner().invoke(
+        milvus_cli.MilvusAutoIndex,
+        ["--uri", "http://localhost:19530"],
+    )
+    assert result.exit_code == 0, result.output
+    assert captured["task"].case_config.nq == 1
+
+    result = CliRunner().invoke(
+        milvus_cli.MilvusAutoIndex,
+        ["--uri", "http://localhost:19530", "--nq", "2"],
+    )
+    assert result.exit_code == 0, result.output
+    assert captured["task"].case_config.nq == 2
+
+    result = CliRunner().invoke(
+        milvus_cli.MilvusAutoIndex,
+        ["--uri", "http://localhost:19530", "--nq", "0", "--dry-run"],
+    )
+    assert result.exit_code == 2, result.output
+    assert "Invalid value for '--nq'" in result.output

vectordb_bench/backend/clients/api.py

@@ -183,6 +183,7 @@ class VectorDB(ABC):
     "The filtering types supported by the VectorDB Client, default only non-filter"
     supported_filter_types: list[FilterOp] = [FilterOp.NonFilter]
     name: str = ""
+    supports_batch_search: bool = False
 
     # Whether the client can share a single connection across threads.
     # If False, concurrent runners will deep-copy the instance and call
@@ -377,6 +378,22 @@ def search_embedding(
         """
         raise NotImplementedError
 
+    def search_embeddings(
+        self,
+        queries: list[list[float]],
+        k: int = 100,
+        payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY,
+        tenant: str | None = None,
+    ) -> list[list[int]]:
+        """Search a batch in one request, returning one ordered result list per query.
+
+        Implementations must raise on incomplete or failed batch responses. Backends
+        supporting this API must set supports_batch_search=True; sequential calls to
+        search_embedding are not a batch implementation.
+        """
+        msg = f"{self.name} does not support batch vector search"
+        raise NotImplementedError(msg)
+
     @abstractmethod
     def optimize(self, data_size: int | None = None):
         """optimize will be called between insertion and search in performance cases.

vectordb_bench/backend/clients/milvus/milvus.py

@@ -22,6 +22,7 @@
 
 
 class Milvus(VectorDB):
+    supports_batch_search = True
     supported_filter_types: list[FilterOp] = [
         FilterOp.NonFilter,
         FilterOp.NumGE,
@@ -532,6 +533,15 @@ def search_embedding(
         tenant: str | None = None,
     ) -> list[int]:
         """Perform a search on a query embedding and return results."""
+        return self.search_embeddings([query], k, payload_profile=payload_profile, tenant=tenant)[0]
+
+    def search_embeddings(
+        self,
+        queries: list[list[float]],
+        k: int = 100,
+        payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY,
+        tenant: str | None = None,
+    ) -> list[list[int]]:
         assert self.client is not None
 
         output_fields = None
@@ -548,16 +558,18 @@ def search_embedding(
 
         search_kwargs = {
             "collection_name": self.collection_name,
-            "data": [query],
+            "data": queries,
             "anns_field": self._vector_field,
             "search_params": self.case_config.search_param(),
             "limit": k,
             "filter": expr,
             "output_fields": output_fields,
         }
         res = self.client.search(**search_kwargs)
+        if len(res) != len(queries):
+            raise RuntimeError("Milvus batch search returned an incomplete response")
 
-        return [result[self._primary_field] for result in res[0]]
+        return [[result[self._primary_field] for result in hits] for hits in res]
 
     def search_documents(
         self,

vectordb_bench/backend/runner/mp_runner.py

@@ -49,8 +49,14 @@ def __init__(
         payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY,
         tenant_labels: list[str] | None = None,
         workload_kind: WorkloadKind = WorkloadKind.VECTOR,
+        nq: int = 1,
     ):
         self.db = db
+        if nq < 1:
+            raise ValueError("nq must be positive")
+        if nq > 1 and (workload_kind != WorkloadKind.VECTOR or not db.supports_batch_search):
+            raise NotImplementedError("nq > 1 requires a backend with batch vector search support")
+        self.nq = nq
         self.k = k
         self.filters = filters
         self.workload_kind = workload_kind
@@ -93,16 +99,19 @@ def __setstate__(self, state: dict):
             msg = f"Unsupported search workload: {self.workload_kind}"
             raise NotImplementedError(msg)
 
-    def _search_embedding(self, emb: list[float], tenant: str | None = None) -> list[int]:
+    def _search_embedding(
+        self, emb: list[float] | list[list[float]], tenant: str | None = None
+    ) -> list[int] | list[list[int]]:
+        search = self.db.search_embedding if self.nq == 1 else self.db.search_embeddings
         if tenant is None:
             if self.payload_profile == PayloadProfile.IDS_ONLY:
-                return self.db.search_embedding(emb, self.k)
-            return self.db.search_embedding(emb, self.k, payload_profile=self.payload_profile)
+                return search(emb, self.k)
+            return search(emb, self.k, payload_profile=self.payload_profile)
         if self.payload_profile == PayloadProfile.IDS_ONLY:
-            return self.db.search_embedding(emb, self.k, tenant=tenant)
-        return self.db.search_embedding(emb, self.k, payload_profile=self.payload_profile, tenant=tenant)
+            return search(emb, self.k, tenant=tenant)
+        return search(emb, self.k, payload_profile=self.payload_profile, tenant=tenant)
 
-    def _search_once(self, query: list[float] | str, tenant_rng: random.Random | None = None):
+    def _search_once(self, query: list[float] | list[list[float]] | str, tenant_rng: random.Random | None = None):
         if self.workload_kind == WorkloadKind.FULL_TEXT:
             if self.payload_profile == PayloadProfile.IDS_ONLY:
                 return self._search_func(query, self.k)
@@ -112,7 +121,16 @@ def _search_once(self, query: list[float] | str, tenant_rng: random.Random | Non
             if tenant_rng is not None and self.tenant_labels
             else None
         )
-        return self._search_func(query, tenant=tenant)
+        results = self._search_func(query, tenant=tenant)
+        if self.nq > 1 and len(results) != len(query):
+            raise RuntimeError("Batch search must return one result list per query")
+        return results
+
+    def _query_at(self, test_data: list, idx: int):
+        if self.nq == 1:
+            return test_data[idx]
+        # Wrap around to keep every throughput request at exactly nq vectors.
+        return [test_data[(idx + offset) % len(test_data)] for offset in range(self.nq)]
 
     def search(
         self,
@@ -141,14 +159,14 @@ def search(
             while time.perf_counter() < start_time + self.duration:
                 s = time.perf_counter()
                 try:
-                    self._search_once(test_data[idx], tenant_rng=tenant_rng)
-                    count += 1
+                    self._search_once(self._query_at(test_data, idx), tenant_rng=tenant_rng)
+                    count += self.nq
                     latencies.append(time.perf_counter() - s)
                 except Exception as e:
                     log.warning(f"VectorDB search_embedding error: {e}")
 
                 # loop through the test data
-                idx = idx + 1 if idx < num - 1 else 0
+                idx = (idx + self.nq) % num
 
                 if count % 500 == 0:
                     log.debug(
@@ -397,8 +415,8 @@ def search_by_dur(
     ) -> tuple[int, int, dict]:
         """
         Returns:
-            int: successful requests count
-            int: failed requests count
+            int: successful query vectors count (requests for full-text search)
+            int: failed query vectors count (requests for full-text search)
             dict: latency statistics with p99, p95, p50, avg, count (computed via HDR Histogram)
         """
         # sync all process
@@ -419,20 +437,20 @@ def search_by_dur(
             while time.perf_counter() < start_time + dur:
                 s = time.perf_counter()
                 try:
-                    self._search_once(test_data[idx])
-                    success_count += 1
+                    self._search_once(self._query_at(test_data, idx))
+                    success_count += self.nq
                     latency_us = int((time.perf_counter() - s) * US_TO_SECONDS)
                     histogram.record_value(max(HDR_HISTOGRAM_MIN_US, min(latency_us, HDR_HISTOGRAM_MAX_US)))
                 except Exception as e:
-                    failed_cnt += 1
+                    failed_cnt += self.nq
                     # reduce log
-                    if failed_cnt <= 3:
+                    if failed_cnt <= 3 * self.nq:
                         log.warning(f"VectorDB search_embedding error: {e}")
                     else:
                         log.debug(f"VectorDB search_embedding error: {e}")
 
                 # loop through the test data
-                idx = idx + 1 if idx < num - 1 else 0
+                idx = (idx + self.nq) % num
 
                 if success_count % 500 == 0:
                     log.debug(