498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
751 lines
37 KiB
Python
751 lines
37 KiB
Python
"""
|
||
Large TopK Feature E2E Tests
|
||
Feature: collection-level property `query_mode: large_topk`
|
||
- Backend index auto-switches to IVF RBQ2
|
||
- Supports topk up to 1M level
|
||
- alter/drop property requires dropping vector index first
|
||
|
||
Test Plan: tests/python_client/docs/test-plan-large-topk.md
|
||
Issue: https://github.com/milvus-io/milvus/issues/48725
|
||
"""
|
||
import pytest
|
||
from base.client_v2_base import TestMilvusClientV2Base
|
||
from common import common_func as cf
|
||
from common import common_type as ct
|
||
from common.common_type import CaseLabel, CheckTasks
|
||
from pymilvus import AnnSearchRequest, DataType, MilvusException, RRFRanker
|
||
|
||
prefix = "large_topk"
|
||
default_nb = 3000 # > 1024 to trigger IVF index build
|
||
default_dim = ct.default_dim # 128
|
||
default_nq = ct.default_nq # 2
|
||
default_limit = ct.default_limit # 10
|
||
vec_field = ct.default_float_vec_field_name # "float_vector"
|
||
large_topk_first = 16385 # first topk above the normal 16384 limit
|
||
large_topk_total = 21000 # total rows in col_large_topk (> large_topk_first + default_nb for headroom)
|
||
large_topk_max = 1_000_000 # maximum supported topk in large_topk mode
|
||
|
||
|
||
@pytest.mark.xdist_group("TestLargeTopkShared")
|
||
class TestLargeTopkShared(TestMilvusClientV2Base):
|
||
"""
|
||
L0 + L1 read-only tests. Two shared collections are prepared once:
|
||
- col_large_topk: query_mode=large_topk, 6×default_nb vectors, FLAT index
|
||
- col_normal: no query_mode set, default_nb vectors, FLAT index
|
||
All tests are read-only; no data modification or index changes.
|
||
"""
|
||
|
||
def setup_class(self):
|
||
super().setup_class(self)
|
||
self.col_large_topk = "TestLargeTopkSharedLargeTopk" + cf.gen_unique_str("_")
|
||
self.col_normal = "TestLargeTopkSharedNormal" + cf.gen_unique_str("_")
|
||
|
||
@pytest.fixture(scope="class", autouse=True)
|
||
def prepare_collections(self, request):
|
||
client = self._client()
|
||
|
||
def _create(col_name, enable_large_topk, batches=1):
|
||
schema = self.create_schema(client)[0]
|
||
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
||
schema.add_field(vec_field, DataType.FLOAT_VECTOR, dim=default_dim)
|
||
query_mode_props = {"query_mode": "large_topk"} if enable_large_topk else None
|
||
self.create_collection(client, col_name, schema=schema,
|
||
properties=query_mode_props, force_teardown=False)
|
||
index_params = self.prepare_index_params(client)[0]
|
||
# FLAT: 100% recall, simplifies assertions
|
||
index_params.add_index(vec_field, index_type="FLAT", metric_type="L2")
|
||
self.create_index(client, col_name, index_params)
|
||
self.load_collection(client, col_name)
|
||
rows = [{vec_field: cf.gen_vectors(1, default_dim)[0]} for _ in range(default_nb)]
|
||
for _ in range(batches):
|
||
self.insert(client, col_name, rows)
|
||
self.flush(client, col_name)
|
||
|
||
# large_topk_total rows total; topk tests use large_topk_total - default_nb for ~1 batch headroom
|
||
_create(self.col_large_topk, enable_large_topk=True, batches=large_topk_total // default_nb)
|
||
_create(self.col_normal, enable_large_topk=False)
|
||
|
||
def teardown():
|
||
client.drop_collection(self.col_large_topk)
|
||
client.drop_collection(self.col_normal)
|
||
|
||
request.addfinalizer(teardown)
|
||
|
||
@pytest.mark.tags(CaseLabel.L0)
|
||
def test_create_with_large_topk_property(self):
|
||
"""
|
||
target: verify query_mode=large_topk property is correctly set at create time
|
||
method:
|
||
1. describe_collection and check properties dict contains query_mode=large_topk
|
||
2. search with limit=100, check nq/limit/metric
|
||
expected: property present; search returns 100 results with ascending L2 distances
|
||
"""
|
||
client = self._client()
|
||
desc = client.describe_collection(self.col_large_topk)
|
||
props = desc.get("properties", {})
|
||
assert props.get("query_mode") == "large_topk", f"property not set: {props}"
|
||
|
||
vectors = cf.gen_vectors(default_nq, default_dim)
|
||
self.search(client, self.col_large_topk, data=vectors,
|
||
anns_field=vec_field, limit=100,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": default_nq,
|
||
"limit": 100,
|
||
"metric": "L2"})
|
||
|
||
@pytest.mark.tags(CaseLabel.L1)
|
||
@pytest.mark.parametrize("topk", [1, 100, 16384])
|
||
def test_search_various_topk(self, topk):
|
||
"""
|
||
target: verify search with various topk values all return correct counts
|
||
method: search col_with_prop with topk in [1, 100, 16384], check nq/limit/metric
|
||
expected: each search returns exactly min(topk, default_nb) hits per query
|
||
"""
|
||
client = self._client()
|
||
vectors = cf.gen_vectors(default_nq, default_dim)
|
||
expected_limit = min(topk, large_topk_total)
|
||
self.search(client, self.col_large_topk, data=vectors,
|
||
anns_field=vec_field, limit=topk,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": default_nq,
|
||
"limit": expected_limit,
|
||
"metric": "L2"})
|
||
|
||
@pytest.mark.tags(CaseLabel.L1)
|
||
def test_search_result_consistency(self):
|
||
"""
|
||
target: verify repeated searches return identical results (IDs and distances)
|
||
method: search same query vector twice, compare id list and distance list
|
||
expected: ids and distances are identical across two calls
|
||
"""
|
||
client = self._client()
|
||
vectors = cf.gen_vectors(1, default_dim)
|
||
|
||
res1 = client.search(self.col_large_topk, data=vectors,
|
||
limit=50, anns_field=vec_field)
|
||
res2 = client.search(self.col_large_topk, data=vectors,
|
||
limit=50, anns_field=vec_field)
|
||
|
||
ids1 = [r["id"] for r in res1[0]]
|
||
ids2 = [r["id"] for r in res2[0]]
|
||
dist1 = [r["distance"] for r in res1[0]]
|
||
dist2 = [r["distance"] for r in res2[0]]
|
||
assert ids1 == ids2, f"Inconsistent IDs: {ids1[:5]} vs {ids2[:5]}"
|
||
assert dist1 == dist2, f"Inconsistent distances: {dist1[:5]} vs {dist2[:5]}"
|
||
|
||
@pytest.mark.tags(CaseLabel.L0)
|
||
@pytest.mark.parametrize("topk", [large_topk_first, large_topk_total - default_nb])
|
||
def test_large_topk_above_normal_limit(self, topk):
|
||
"""
|
||
target: verify query_mode=large_topk allows topk above 16384 without error (core MVP)
|
||
method: search col_large_topk (large_topk_total vectors) with topk in [large_topk_first, large_topk_total - default_nb]
|
||
expected: search completes without exception, returns results with ascending L2 distances.
|
||
Note: exact result count is not asserted — large_topk forces IVF index which
|
||
does not guarantee 100% recall, so returned count may be < topk.
|
||
"""
|
||
client = self._client()
|
||
results = client.search(self.col_large_topk, data=cf.gen_vectors(default_nq, default_dim),
|
||
anns_field=vec_field, limit=topk)
|
||
for hits in results:
|
||
assert len(hits) > 0, "Expected non-empty results"
|
||
distances = [h["distance"] for h in hits]
|
||
assert distances == sorted(distances), "L2 distances should be ascending"
|
||
|
||
@pytest.mark.tags(CaseLabel.L2)
|
||
def test_topk_without_large_topk_property(self):
|
||
"""
|
||
target: verify topk>16384 is rejected when query_mode=large_topk is not set
|
||
method:
|
||
1. search col_without_prop with limit=16384 — should succeed
|
||
2. search col_without_prop with limit=large_topk_first — should raise MilvusException
|
||
expected: limit=16384 OK; limit=large_topk_first raises error with message about invalid topk
|
||
"""
|
||
client = self._client()
|
||
vectors = cf.gen_vectors(default_nq, default_dim)
|
||
|
||
# Normal topk limit works fine
|
||
self.search(client, self.col_normal, data=vectors,
|
||
anns_field=vec_field, limit=16384,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": default_nq,
|
||
"limit": default_nb,
|
||
"metric": "L2"})
|
||
|
||
# Above limit must be rejected
|
||
error = {ct.err_code: 65535,
|
||
ct.err_msg: f"topk [{large_topk_first}] is invalid, it should be in range [1, 16384]"}
|
||
self.search(client, self.col_normal, data=vectors,
|
||
anns_field=vec_field, limit=large_topk_first,
|
||
check_task=CheckTasks.err_res,
|
||
check_items=error)
|
||
|
||
# Note: search_iterator and query_iterator are NOT affected by query_mode=large_topk.
|
||
# The SDK enforces batch_size <= 16384 client-side (ParamError, unrelated to large_topk).
|
||
# The iterator `limit` (total result count) uses internal pagination with batch_size <= 16384,
|
||
# so per-request topk never exceeds 16384. No large_topk-specific iterator tests are needed.
|
||
|
||
@pytest.mark.tags(CaseLabel.L1)
|
||
def test_query_large_limit(self):
|
||
"""
|
||
target: verify query() with limit > 16384 works when query_mode=large_topk is set
|
||
method: query col_large_topk with limit=large_topk_first, verify results returned
|
||
expected: returns large_topk_first results without error
|
||
"""
|
||
client = self._client()
|
||
res = client.query(
|
||
self.col_large_topk,
|
||
filter="",
|
||
output_fields=["id"],
|
||
limit=large_topk_first,
|
||
)
|
||
assert len(res) == large_topk_first, \
|
||
f"Expected {large_topk_first} results, got {len(res)}"
|
||
|
||
@pytest.mark.tags(CaseLabel.L2)
|
||
def test_query_without_property_fails(self):
|
||
"""
|
||
target: verify query() with limit > 16384 is rejected when query_mode=large_topk is NOT set
|
||
method: query col_normal with limit=large_topk_first
|
||
expected: MilvusException with invalid topk message
|
||
"""
|
||
client = self._client()
|
||
with pytest.raises(MilvusException) as exc_info:
|
||
client.query(
|
||
self.col_normal,
|
||
filter="",
|
||
output_fields=["id"],
|
||
limit=large_topk_first,
|
||
)
|
||
assert str(large_topk_first) in str(exc_info.value), \
|
||
f"Expected topk error, got: {exc_info.value}"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Independent Tests (each test owns its own collection)
|
||
# ---------------------------------------------------------------------------
|
||
class TestLargeTopkIndependent(TestMilvusClientV2Base):
|
||
"""
|
||
Tests that require modifying index or property — each test gets its own collection.
|
||
force_teardown=True (default) ensures cleanup even on failure.
|
||
"""
|
||
|
||
def _setup_col(self, client, enable_large_topk=True, nb=default_nb):
|
||
"""Create collection with optional query_mode=large_topk, FLAT index, insert nb rows."""
|
||
col = cf.gen_collection_name_by_testcase_name(module_index=2)
|
||
schema = self.create_schema(client)[0]
|
||
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
||
schema.add_field(vec_field, DataType.FLOAT_VECTOR, dim=default_dim)
|
||
query_mode_props = {"query_mode": "large_topk"} if enable_large_topk else None
|
||
self.create_collection(client, col, schema=schema,
|
||
properties=query_mode_props, force_teardown=True)
|
||
index_params = self.prepare_index_params(client)[0]
|
||
index_params.add_index(vec_field, index_type="FLAT", metric_type="L2")
|
||
self.create_index(client, col, index_params)
|
||
self.load_collection(client, col)
|
||
if nb > 0:
|
||
rows = [{vec_field: cf.gen_vectors(1, default_dim)[0]} for _ in range(nb)]
|
||
self.insert(client, col, rows)
|
||
self.flush(client, col)
|
||
return col
|
||
|
||
@pytest.mark.tags(CaseLabel.L1)
|
||
def test_alter_collection_add_property(self):
|
||
"""
|
||
target: verify alter_collection_properties correctly adds query_mode=large_topk
|
||
method:
|
||
1. create collection without property, build FLAT index, insert data
|
||
2. release → drop_index → alter_collection_properties
|
||
3. describe_collection to verify property
|
||
4. rebuild index, load, search with limit=100
|
||
expected: property set; search returns 100 results with ascending L2 distances
|
||
"""
|
||
client = self._client()
|
||
col = self._setup_col(client, enable_large_topk=False)
|
||
|
||
self.release_collection(client, col)
|
||
self.drop_index(client, col, vec_field)
|
||
|
||
self.alter_collection_properties(client, col,
|
||
properties={"query_mode": "large_topk"})
|
||
|
||
desc = client.describe_collection(col)
|
||
assert desc.get("properties", {}).get("query_mode") == "large_topk"
|
||
|
||
index_params = self.prepare_index_params(client)[0]
|
||
index_params.add_index(vec_field, index_type="FLAT", metric_type="L2")
|
||
self.create_index(client, col, index_params)
|
||
self.load_collection(client, col)
|
||
|
||
self.search(client, col, data=cf.gen_vectors(default_nq, default_dim),
|
||
anns_field=vec_field, limit=100,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": default_nq, "limit": 100, "metric": "L2"})
|
||
|
||
@pytest.mark.tags(CaseLabel.L1)
|
||
def test_drop_collection_property(self):
|
||
"""
|
||
target: verify drop_collection_properties removes query_mode and restores normal behavior
|
||
method:
|
||
1. create collection with property, build FLAT index, insert data
|
||
2. release → drop_index → drop_collection_properties
|
||
3. describe_collection to verify property absent
|
||
4. rebuild index, load, search with limit=default_limit
|
||
expected: property absent; normal search returns default_limit results
|
||
"""
|
||
client = self._client()
|
||
col = self._setup_col(client, enable_large_topk=True)
|
||
|
||
self.release_collection(client, col)
|
||
self.drop_index(client, col, vec_field)
|
||
|
||
self.drop_collection_properties(client, col, property_keys=["query_mode"])
|
||
|
||
desc = client.describe_collection(col)
|
||
assert "query_mode" not in desc.get("properties", {}), \
|
||
f"property still present: {desc.get('properties')}"
|
||
|
||
index_params = self.prepare_index_params(client)[0]
|
||
index_params.add_index(vec_field, index_type="FLAT", metric_type="L2")
|
||
self.create_index(client, col, index_params)
|
||
self.load_collection(client, col)
|
||
|
||
self.search(client, col, data=cf.gen_vectors(default_nq, default_dim),
|
||
anns_field=vec_field, limit=default_limit,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": default_nq,
|
||
"limit": default_limit,
|
||
"metric": "L2"})
|
||
|
||
@pytest.mark.tags(CaseLabel.L1)
|
||
def test_alter_property_without_dropping_index_fails(self):
|
||
"""
|
||
target: verify alter_collection_properties is rejected when vector index exists
|
||
method: create collection with index, call alter_collection_properties directly
|
||
expected: MilvusException with error code 702 and message containing "vector index"
|
||
"""
|
||
client = self._client()
|
||
col = self._setup_col(client, enable_large_topk=False, nb=0)
|
||
|
||
error = {ct.err_code: 702,
|
||
ct.err_msg: "can not alter query_mode if the collection already has a vector index"}
|
||
self.alter_collection_properties(client, col,
|
||
properties={"query_mode": "large_topk"},
|
||
check_task=CheckTasks.err_res,
|
||
check_items=error)
|
||
|
||
@pytest.mark.tags(CaseLabel.L1)
|
||
def test_drop_property_without_dropping_index_fails(self):
|
||
"""
|
||
target: verify drop_collection_properties is rejected when vector index exists
|
||
method: create collection with large_topk property and index, call drop directly
|
||
expected: MilvusException with error code 702 and message containing "vector index"
|
||
"""
|
||
client = self._client()
|
||
col = self._setup_col(client, enable_large_topk=True, nb=0)
|
||
|
||
error = {ct.err_code: 702,
|
||
ct.err_msg: "can not alter query_mode if the collection already has a vector index"}
|
||
self.drop_collection_properties(client, col,
|
||
property_keys=["query_mode"],
|
||
check_task=CheckTasks.err_res,
|
||
check_items=error)
|
||
|
||
@pytest.mark.tags(CaseLabel.L2)
|
||
def test_empty_collection_search(self):
|
||
"""
|
||
target: verify search on empty large_topk collection returns 0 results
|
||
method: create collection with property, build FLAT index, load, search without inserting
|
||
expected: 0 results returned, no exception
|
||
"""
|
||
client = self._client()
|
||
col = self._setup_col(client, enable_large_topk=True, nb=0)
|
||
|
||
res = client.search(col, data=cf.gen_vectors(default_nq, default_dim),
|
||
limit=default_limit, anns_field=vec_field)
|
||
for hits in res:
|
||
assert len(hits) == 0, f"Empty collection should return 0 results, got {len(hits)}"
|
||
|
||
@pytest.mark.tags(CaseLabel.L2)
|
||
def test_add_then_drop_property_roundtrip(self):
|
||
"""
|
||
target: verify adding then dropping query_mode property restores normal behavior
|
||
method:
|
||
1. create plain collection, build FLAT index, insert data
|
||
2. drop index → alter_collection_properties (add) → rebuild index → search
|
||
3. drop index → drop_collection_properties → rebuild index → search
|
||
4. verify property absent after final drop
|
||
expected: both searches return default_limit results; property absent after drop
|
||
"""
|
||
client = self._client()
|
||
col = self._setup_col(client, enable_large_topk=False)
|
||
vectors = cf.gen_vectors(default_nq, default_dim)
|
||
|
||
# Phase 1: add property
|
||
self.release_collection(client, col)
|
||
self.drop_index(client, col, vec_field)
|
||
self.alter_collection_properties(client, col,
|
||
properties={"query_mode": "large_topk"})
|
||
index_params = self.prepare_index_params(client)[0]
|
||
index_params.add_index(vec_field, index_type="FLAT", metric_type="L2")
|
||
self.create_index(client, col, index_params)
|
||
self.load_collection(client, col)
|
||
self.search(client, col, data=vectors, anns_field=vec_field,
|
||
limit=default_limit,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": default_nq,
|
||
"limit": default_limit,
|
||
"metric": "L2"})
|
||
|
||
# Phase 2: drop property
|
||
self.release_collection(client, col)
|
||
self.drop_index(client, col, vec_field)
|
||
self.drop_collection_properties(client, col, property_keys=["query_mode"])
|
||
self.create_index(client, col, index_params)
|
||
self.load_collection(client, col)
|
||
self.search(client, col, data=vectors, anns_field=vec_field,
|
||
limit=default_limit,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": default_nq,
|
||
"limit": default_limit,
|
||
"metric": "L2"})
|
||
|
||
desc = client.describe_collection(col)
|
||
assert "query_mode" not in desc.get("properties", {})
|
||
|
||
@pytest.mark.tags(CaseLabel.L1)
|
||
def test_property_persistence_after_reload(self):
|
||
"""
|
||
target: verify query_mode=large_topk persists after release + load
|
||
method:
|
||
1. create collection with property, insert total > large_topk_first rows, flush
|
||
2. release → load
|
||
3. describe_collection to verify property present
|
||
4. search with limit=large_topk_first, verify large_topk_first results returned
|
||
expected: property present; topk=large_topk_first returns exactly large_topk_first results
|
||
"""
|
||
client = self._client()
|
||
nb_total = large_topk_first + 1000
|
||
col = self._setup_col(client, enable_large_topk=True, nb=nb_total)
|
||
|
||
self.release_collection(client, col)
|
||
self.load_collection(client, col)
|
||
|
||
desc = client.describe_collection(col)
|
||
assert desc.get("properties", {}).get("query_mode") == "large_topk", \
|
||
f"property lost after reload: {desc.get('properties')}"
|
||
|
||
self.search(client, col, data=cf.gen_vectors(default_nq, default_dim),
|
||
anns_field=vec_field, limit=large_topk_first,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": default_nq,
|
||
"limit": large_topk_first,
|
||
"metric": "L2"})
|
||
|
||
@pytest.mark.tags(CaseLabel.L1)
|
||
def test_create_index_after_insert_with_large_topk(self):
|
||
"""
|
||
target: verify create_index on a populated collection with query_mode=large_topk,
|
||
and that large topk search (>16384) works correctly after index build
|
||
method:
|
||
1. create collection with query_mode=large_topk (no index yet)
|
||
2. insert large_topk_first + default_nb rows, flush
|
||
3. create_index on the populated collection
|
||
4. load
|
||
5. search with limit=large_topk_first (>16384) to verify large topk is functional
|
||
6. search with limit > nb rows (capped to actual nb) to verify normal search works
|
||
expected: create_index succeeds on populated collection;
|
||
large topk search returns large_topk_first results;
|
||
normal search returns min(limit, nb) results
|
||
note: this catches the timeout seen when rebuilding index after alter/drop property,
|
||
isolating whether the issue is data-at-index-build-time or the alter step itself
|
||
"""
|
||
client = self._client()
|
||
nb = large_topk_first + default_nb
|
||
col = cf.gen_collection_name_by_testcase_name()
|
||
schema = self.create_schema(client)[0]
|
||
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
||
schema.add_field(vec_field, DataType.FLOAT_VECTOR, dim=default_dim)
|
||
self.create_collection(client, col, schema=schema,
|
||
properties={"query_mode": "large_topk"}, force_teardown=True)
|
||
|
||
rows = [{vec_field: cf.gen_vectors(1, default_dim)[0]} for _ in range(nb)]
|
||
self.insert(client, col, rows)
|
||
self.flush(client, col)
|
||
|
||
index_params = self.prepare_index_params(client)[0]
|
||
index_params.add_index(vec_field, index_type="FLAT", metric_type="L2")
|
||
self.create_index(client, col, index_params)
|
||
self.load_collection(client, col)
|
||
|
||
# verify large topk (>16384) is functional
|
||
self.search(client, col, data=cf.gen_vectors(default_nq, default_dim),
|
||
anns_field=vec_field, limit=large_topk_first,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": default_nq, "limit": large_topk_first, "metric": "L2"})
|
||
|
||
# verify normal search also works
|
||
self.search(client, col, data=cf.gen_vectors(default_nq, default_dim),
|
||
anns_field=vec_field, limit=100,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": default_nq, "limit": 100, "metric": "L2"})
|
||
|
||
@pytest.mark.tags(CaseLabel.L1)
|
||
def test_large_topk_growing_segment(self):
|
||
"""
|
||
target: verify large_topk works on growing segments (before flush)
|
||
method:
|
||
1. create collection with property, build FLAT index, load
|
||
2. insert default_nb rows WITHOUT flush (growing segment)
|
||
3. search with limit=100 — should return results from growing segment
|
||
expected: search succeeds and returns hits; no error about topk limit
|
||
"""
|
||
client = self._client()
|
||
col = self._setup_col(client, enable_large_topk=True, nb=0)
|
||
|
||
# Insert without flush → growing segment
|
||
rows = [{vec_field: cf.gen_vectors(1, default_dim)[0]} for _ in range(default_nb)]
|
||
self.insert(client, col, rows)
|
||
# No flush — data stays in growing segment
|
||
|
||
self.search(client, col, data=cf.gen_vectors(default_nq, default_dim),
|
||
anns_field=vec_field, limit=100,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": default_nq,
|
||
"limit": 100,
|
||
"metric": "L2"})
|
||
|
||
@pytest.mark.tags(CaseLabel.L2)
|
||
def test_invalid_property_value(self):
|
||
"""
|
||
target: verify invalid query_mode value is rejected at collection creation
|
||
method: create_collection with properties={"query_mode": "invalid_mode"}
|
||
expected: MilvusException with message containing valid values hint
|
||
"""
|
||
client = self._client()
|
||
col = cf.gen_collection_name_by_testcase_name()
|
||
schema = self.create_schema(client)[0]
|
||
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
||
schema.add_field(vec_field, DataType.FLOAT_VECTOR, dim=default_dim)
|
||
|
||
error = {ct.err_code: 65535,
|
||
ct.err_msg: 'invalid query_mode value "invalid_mode", valid values: [large_topk]'}
|
||
self.create_collection(client, col, schema=schema,
|
||
properties={"query_mode": "invalid_mode"},
|
||
check_task=CheckTasks.err_res,
|
||
check_items=error)
|
||
|
||
@pytest.mark.tags(CaseLabel.L2)
|
||
@pytest.mark.parametrize("value", ["LARGE_TOPK", "Large_TopK", "large_TOPK"])
|
||
def test_query_mode_value_case_insensitive(self, value):
|
||
"""
|
||
target: verify query_mode value is case-sensitive
|
||
method: create collection with properties={"query_mode": value} (non-lowercase value)
|
||
expected: create collection fails with invalid query_mode value error
|
||
"""
|
||
client = self._client()
|
||
col = cf.gen_collection_name_by_testcase_name()
|
||
schema = self.create_schema(client)[0]
|
||
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
||
schema.add_field(vec_field, DataType.FLOAT_VECTOR, dim=default_dim)
|
||
error = {ct.err_code: 65535,
|
||
ct.err_msg: f'invalid query_mode value "{value}", valid values: [large_topk]'}
|
||
self.create_collection(client, col, schema=schema,
|
||
properties={"query_mode": value},
|
||
check_task=CheckTasks.err_res,
|
||
check_items=error)
|
||
|
||
@pytest.mark.tags(CaseLabel.L2)
|
||
@pytest.mark.parametrize("key", ["QUERY_MODE", "Query_Mode", "query_MODE"])
|
||
def test_query_mode_key_case_sensitive(self, key):
|
||
"""
|
||
target: verify query_mode key is case-sensitive
|
||
method: create collection with properties={key: "large_topk"} (wrong-cased key)
|
||
expected: create collection fails with invalid property key error
|
||
"""
|
||
client = self._client()
|
||
col = cf.gen_collection_name_by_testcase_name()
|
||
schema = self.create_schema(client)[0]
|
||
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
||
schema.add_field(vec_field, DataType.FLOAT_VECTOR, dim=default_dim)
|
||
error = {ct.err_code: 65535,
|
||
ct.err_msg: f'invalid property key "{key}", did you mean "query_mode"?'}
|
||
self.create_collection(client, col, schema=schema,
|
||
properties={key: "large_topk"},
|
||
check_task=CheckTasks.err_res,
|
||
check_items=error)
|
||
|
||
# Note: search_iterator and query_iterator are NOT affected by query_mode=large_topk.
|
||
# The SDK enforces batch_size <= 16384 client-side (ParamError code=1, regardless of property).
|
||
# Iterator `limit` (total results) uses internal pagination with batch_size <= 16384 per request,
|
||
# so per-request topk never exceeds 16384. No large_topk-specific iterator tests are needed.
|
||
|
||
# -------------------------------------------------------------------------
|
||
# Large topk boundary tests (L3, large data volume)
|
||
# -------------------------------------------------------------------------
|
||
|
||
@pytest.mark.tags(CaseLabel.L3)
|
||
def test_large_topk_boundary_2m_rows(self):
|
||
"""
|
||
target: verify large_topk topk boundary values with 2M rows
|
||
method:
|
||
1. create collection with query_mode=large_topk
|
||
2. insert 2,000,000 rows (128-dim) in batches, flush
|
||
3. create_index, load
|
||
4. search with limit=large_topk_max-1 (999,999) → should succeed
|
||
5. search with limit=large_topk_max (1,000,000) → should succeed, return 1M results
|
||
6. search with limit=large_topk_max+1 (1,000,001) → should fail with error
|
||
expected: boundary limits enforced correctly; max valid topk returns 1M results
|
||
"""
|
||
client = self._client()
|
||
total_nb = 2_000_000
|
||
batch_size = 50_000
|
||
col = cf.gen_collection_name_by_testcase_name()
|
||
dim = 64
|
||
schema = self.create_schema(client)[0]
|
||
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
||
schema.add_field(vec_field, DataType.FLOAT_VECTOR, dim=dim)
|
||
self.create_collection(client, col, schema=schema,
|
||
properties={"query_mode": "large_topk"}, force_teardown=True)
|
||
|
||
for _ in range(total_nb // batch_size):
|
||
vecs = cf.gen_vectors(batch_size, dim)
|
||
rows = [{vec_field: v} for v in vecs]
|
||
self.insert(client, col, rows)
|
||
self.flush(client, col)
|
||
|
||
index_params = self.prepare_index_params(client)[0]
|
||
index_params.add_index(vec_field, index_type="FLAT", metric_type="L2")
|
||
self.create_index(client, col, index_params)
|
||
self.load_collection(client, col)
|
||
|
||
# just below max → should succeed
|
||
self.search(client, col, data=cf.gen_vectors(1, dim),
|
||
anns_field=vec_field, limit=large_topk_max - 1,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": 1, "limit": large_topk_max - 1, "metric": "L2"})
|
||
|
||
# max valid large topk → should succeed, return 1M results
|
||
self.search(client, col, data=cf.gen_vectors(1, dim),
|
||
anns_field=vec_field, limit=large_topk_max,
|
||
check_task=CheckTasks.check_search_results,
|
||
check_items={"enable_milvus_client_api": True,
|
||
"nq": 1, "limit": large_topk_max, "metric": "L2"})
|
||
|
||
# over max → should fail
|
||
error = {ct.err_code: 65535,
|
||
ct.err_msg: f"topk [{large_topk_max + 1}] is invalid, "
|
||
f"it should be in range [1, {large_topk_max}]"}
|
||
self.search(client, col, data=cf.gen_vectors(1, dim),
|
||
anns_field=vec_field, limit=large_topk_max + 1,
|
||
check_task=CheckTasks.err_res,
|
||
check_items=error)
|
||
|
||
# -------------------------------------------------------------------------
|
||
# Hybrid search interface tests
|
||
# -------------------------------------------------------------------------
|
||
|
||
def _setup_dual_vec_col(self, client, enable_large_topk=True, nb=default_nb):
|
||
"""Create collection with two float vector fields for hybrid search tests.
|
||
Uses IVF_FLAT index to keep index-build time reasonable for large nb."""
|
||
col = cf.gen_collection_name_by_testcase_name(module_index=2)
|
||
schema = self.create_schema(client)[0]
|
||
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
||
schema.add_field(vec_field, DataType.FLOAT_VECTOR, dim=default_dim)
|
||
schema.add_field("vec2", DataType.FLOAT_VECTOR, dim=default_dim)
|
||
query_mode_props = {"query_mode": "large_topk"} if enable_large_topk else None
|
||
self.create_collection(client, col, schema=schema,
|
||
properties=query_mode_props, force_teardown=True)
|
||
index_params = self.prepare_index_params(client)[0]
|
||
index_params.add_index(vec_field, index_type="IVF_FLAT", metric_type="L2",
|
||
params={"nlist": 64})
|
||
index_params.add_index("vec2", index_type="IVF_FLAT", metric_type="L2",
|
||
params={"nlist": 64})
|
||
self.create_index(client, col, index_params)
|
||
self.load_collection(client, col)
|
||
if nb > 0:
|
||
rows = [{vec_field: cf.gen_vectors(1, default_dim)[0],
|
||
"vec2": cf.gen_vectors(1, default_dim)[0]} for _ in range(nb)]
|
||
self.insert(client, col, rows)
|
||
self.flush(client, col)
|
||
return col
|
||
|
||
@pytest.mark.tags(CaseLabel.L1)
|
||
def test_hybrid_search_large_topk(self):
|
||
"""
|
||
target: verify hybrid_search with limit > 16384 works when query_mode=large_topk is set
|
||
method:
|
||
1. create collection with two float vector fields and query_mode=large_topk
|
||
2. insert large_topk_total rows with IVF_FLAT index
|
||
3. hybrid_search with limit=large_topk_first using RRFRanker
|
||
expected: hybrid_search completes without error; returns > 0 results; no error code
|
||
"""
|
||
client = self._client()
|
||
col = self._setup_dual_vec_col(client, enable_large_topk=True, nb=large_topk_total)
|
||
req_list = [
|
||
AnnSearchRequest(
|
||
data=cf.gen_vectors(default_nq, default_dim),
|
||
anns_field=vec_field,
|
||
param={"metric_type": "L2", "nprobe": 16},
|
||
limit=large_topk_first,
|
||
),
|
||
AnnSearchRequest(
|
||
data=cf.gen_vectors(default_nq, default_dim),
|
||
anns_field="vec2",
|
||
param={"metric_type": "L2", "nprobe": 16},
|
||
limit=large_topk_first,
|
||
),
|
||
]
|
||
res, _ = self.hybrid_search(client, col,
|
||
reqs=req_list, ranker=RRFRanker(),
|
||
limit=large_topk_first)
|
||
assert len(res) == default_nq, f"Expected {default_nq} query results, got {len(res)}"
|
||
for hits in res:
|
||
assert len(hits) > 0, "Expected non-empty hybrid search results"
|
||
|
||
@pytest.mark.tags(CaseLabel.L2)
|
||
def test_hybrid_search_without_property_fails(self):
|
||
"""
|
||
target: verify hybrid_search with limit > 16384 is rejected when query_mode=large_topk is NOT set
|
||
method:
|
||
1. create collection with two float vector fields and NO query_mode property
|
||
2. hybrid_search with limit=large_topk_first
|
||
expected: MilvusException with invalid topk message
|
||
"""
|
||
client = self._client()
|
||
col = self._setup_dual_vec_col(client, enable_large_topk=False)
|
||
req_list = [
|
||
AnnSearchRequest(
|
||
data=cf.gen_vectors(default_nq, default_dim),
|
||
anns_field=vec_field,
|
||
param={"metric_type": "L2"},
|
||
limit=large_topk_first,
|
||
),
|
||
AnnSearchRequest(
|
||
data=cf.gen_vectors(default_nq, default_dim),
|
||
anns_field="vec2",
|
||
param={"metric_type": "L2"},
|
||
limit=large_topk_first,
|
||
),
|
||
]
|
||
# hybrid_search uses "invalid max query result window" (not "topk [N] is invalid")
|
||
error = {ct.err_code: 65535,
|
||
ct.err_msg: f"invalid max query result window, (offset+limit) should be in range [1, 16384], but got {large_topk_first}"}
|
||
self.hybrid_search(client, col,
|
||
reqs=req_list, ranker=RRFRanker(),
|
||
limit=large_topk_first,
|
||
check_task=CheckTasks.err_res,
|
||
check_items=error)
|
||
|