Skip to content

Commit c1200a4

Browse files
XuanYang-cnclaude
andauthored
fix: pass context to describe_collection in SearchIteratorV2 (#3271) (#3274)
- Forward `**kwargs` (which includes `context` carrying `db_name`) from `SearchIteratorV2.__init__` to `_set_up_collection_id` and then to `describe_collection`, matching the existing `QueryIterator` pattern. - Without this fix, `SearchIteratorV2` always queries the `default` database, causing "can't find collection" errors for non-default databases. - Added regression test verifying context is properly forwarded. See also: #3270, #3271 - [x] Existing `tests/test_search_iterator.py` tests pass - [x] New `test_context_passed_to_describe_collection` regression test passes - [x] Lint (`make lint`) passes --------- Signed-off-by: yangxuan <xuan.yang@zilliz.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent bc8c707 commit c1200a4

File tree

4 files changed

+25
-8
lines changed

4 files changed

+25
-8
lines changed

.gitignore

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ venv/
2525

2626
# Local Temp
2727
temp/
28+
.worktrees/
29+
.worktree/
2830
*.swp
2931
assets/
3032
TODO

pymilvus/bulk_writer/constants.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -56,13 +56,13 @@
5656
DataType.JSON.name: lambda x: isinstance(x, (str, list, dict)),
5757
DataType.TIMESTAMPTZ.name: lambda x: isinstance(x, str),
5858
DataType.GEOMETRY.name: lambda x: isinstance(x, str),
59-
DataType.STRUCT.name: lambda x, max_cap: struct_validator(x, max_cap),
60-
DataType.FLOAT_VECTOR.name: lambda x, dim: float_vector_validator(x, dim),
61-
DataType.BINARY_VECTOR.name: lambda x, dim: binary_vector_validator(x, dim),
59+
DataType.STRUCT.name: struct_validator,
60+
DataType.FLOAT_VECTOR.name: float_vector_validator,
61+
DataType.BINARY_VECTOR.name: binary_vector_validator,
6262
DataType.FLOAT16_VECTOR.name: lambda x, dim: float16_vector_validator(x, dim, False),
6363
DataType.BFLOAT16_VECTOR.name: lambda x, dim: float16_vector_validator(x, dim, True),
64-
DataType.SPARSE_FLOAT_VECTOR.name: lambda x: sparse_vector_validator(x),
65-
DataType.INT8_VECTOR.name: lambda x, dim: int8_vector_validator(x, dim),
64+
DataType.SPARSE_FLOAT_VECTOR.name: sparse_vector_validator,
65+
DataType.INT8_VECTOR.name: int8_vector_validator,
6666
DataType.ARRAY.name: lambda x, cap: (isinstance(x, (list, np.ndarray)) and len(x) <= cap),
6767
}
6868

pymilvus/client/search_iterator.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ def __init__(
4949
self._left_res_cnt = limit
5050

5151
self._conn = connection
52-
self._set_up_collection_id(collection_name)
52+
self._set_up_collection_id(collection_name, **kwargs)
5353
kwargs[COLLECTION_ID] = self._collection_id
5454
self._params = {
5555
"collection_name": collection_name,
@@ -73,8 +73,8 @@ def __init__(
7373
self._batch_size = batch_size
7474
self._probe_for_compability(self._params)
7575

76-
def _set_up_collection_id(self, collection_name: str):
77-
res = self._conn.describe_collection(collection_name)
76+
def _set_up_collection_id(self, collection_name: str, **kwargs):
77+
res = self._conn.describe_collection(collection_name, **kwargs)
7878
self._collection_id = res[COLLECTION_ID]
7979

8080
def _check_token_exists(self, token: Union[str, None]):

tests/test_search_iterator.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import numpy as np
44
import pytest
5+
from pymilvus.client.call_context import CallContext
56
from pymilvus.client.search_iterator import SearchIteratorV2
67
from pymilvus.client.search_result import SearchResult
78
from pymilvus.exceptions import ParamError, ServerVersionIncompatibleException
@@ -88,6 +89,20 @@ def test_multiple_vectors_error(self, mock_connection):
8889
batch_size=100,
8990
)
9091

92+
def test_context_passed_to_describe_collection(self, mock_connection, search_data):
93+
"""Regression test for #3270: context (db_name) must be forwarded to describe_collection"""
94+
ctx = CallContext(db_name="nondefault")
95+
96+
SearchIteratorV2(
97+
connection=mock_connection,
98+
collection_name="test_collection",
99+
data=search_data,
100+
batch_size=100,
101+
context=ctx,
102+
)
103+
104+
mock_connection.describe_collection.assert_called_once_with("test_collection", context=ctx)
105+
91106
@patch("pymilvus.client.search_iterator.SearchIteratorV2._probe_for_compability")
92107
def test_next_without_external_filter(self, mock_probe, mock_connection, search_data):
93108
mock_connection.search.return_value = self.create_mock_search_result()

0 commit comments

Comments
 (0)