fix: iterator alias mismatch(#2555) (#2557)

related: #2555

Signed-off-by: MrPresent-Han <[email protected]>
Co-authored-by: MrPresent-Han <[email protected]>

Author: MrPresent-Han

pymilvus/client/constants.py

@@ -10,6 +10,7 @@
 DEFAULT_RESOURCE_GROUP = "__default_resource_group"
 DYNAMIC_FIELD_NAME = "$meta"
 REDUCE_STOP_FOR_BEST = "reduce_stop_for_best"
+COLLECTION_ID = "collection_id"
 GROUP_BY_FIELD = "group_by_field"
 GROUP_SIZE = "group_size"
 RANK_GROUP_SCORER = "rank_group_scorer"

pymilvus/client/prepare.py

@@ -15,6 +15,7 @@
 from . import __version__, blob, check, entity_helper, ts_utils, utils
 from .check import check_pass_param, is_legal_collection_properties
 from .constants import (
+    COLLECTION_ID,
     DEFAULT_CONSISTENCY_LEVEL,
     DYNAMIC_FIELD_NAME,
     GROUP_BY_FIELD,
@@ -961,6 +962,10 @@ def search_requests_with_expr(
         if is_iterator is not None:
             search_params[ITERATOR_FIELD] = is_iterator
 
+        collection_id = kwargs.get(COLLECTION_ID)
+        if collection_id is not None:
+            search_params[COLLECTION_ID] = str(collection_id)
+
         is_search_iter_v2 = kwargs.get(ITER_SEARCH_V2_KEY)
         if is_search_iter_v2 is not None:
             search_params[ITER_SEARCH_V2_KEY] = is_search_iter_v2
@@ -1310,6 +1315,11 @@ def query_request(
             consistency_level=kwargs.get("consistency_level", 0),
             expr_template_values=cls.prepare_expression_template(kwargs.get("expr_params", {})),
         )
+        collection_id = kwargs.get(COLLECTION_ID)
+        if collection_id is not None:
+            req.query_params.append(
+                common_types.KeyValuePair(key=COLLECTION_ID, value=str(collection_id))
+            )
 
         limit = kwargs.get("limit")
         if limit is not None:

pymilvus/client/search_iterator.py

@@ -5,6 +5,7 @@
 from pymilvus.client import entity_helper, utils
 from pymilvus.client.abstract import Hits
 from pymilvus.client.constants import (
+    COLLECTION_ID,
     GUARANTEE_TIMESTAMP,
     ITER_SEARCH_BATCH_SIZE_KEY,
     ITER_SEARCH_ID_KEY,
@@ -50,6 +51,8 @@ def __init__(
             self._left_res_cnt = limit
 
         self._conn = connection
+        self.__set_up_collection_id(collection_name)
+        kwargs[COLLECTION_ID] = self._collection_id
         self._params = {
             "collection_name": collection_name,
             "data": data,
@@ -71,6 +74,10 @@ def __init__(
         self._saved_first_res = self.next()
         self._is_saved = True
 
+    def __set_up_collection_id(self, collection_name: str):
+        res = self._conn.describe_collection(collection_name)
+        self._collection_id = res[COLLECTION_ID]
+
     def next(self):
         # for compatibility
         if self._is_saved:

pymilvus/orm/constants.py

@@ -47,6 +47,7 @@
 EF = "ef"
 IS_PRIMARY = "is_primary"
 REDUCE_STOP_FOR_BEST = "reduce_stop_for_best"
+COLLECTION_ID = "collection_id"
 ITERATOR_FIELD = "iterator"
 ITERATOR_SESSION_TS_FIELD = "iterator_session_ts"
 DEFAULT_MAX_L2_DISTANCE = 99999999.0

pymilvus/orm/iterator.py

@@ -23,6 +23,7 @@
     CALC_DIST_JACCARD,
     CALC_DIST_L2,
     CALC_DIST_TANIMOTO,
+    COLLECTION_ID,
     DEFAULT_SEARCH_EXTENSION_RATE,
     EF,
     FIELDS,
@@ -101,13 +102,15 @@ def __init__(
     ) -> QueryIterator:
         self._conn = connection
         self._collection_name = collection_name
+        self.__set_up_collection_id()
         self._output_fields = output_fields
         self._partition_names = partition_names
         self._schema = schema
         self._timeout = timeout
         self._session_ts = 0
         self._kwargs = kwargs
         self._kwargs[ITERATOR_FIELD] = "True"
+        self._kwargs[COLLECTION_ID] = self._collection_id
         self.__check_set_batch_size(batch_size)
         self._limit = limit
         self.__check_set_reduce_stop_for_best()
@@ -120,6 +123,10 @@ def __init__(
         self.__set_up_ts_cp()
         self.__seek_to_offset()
 
+    def __set_up_collection_id(self):
+        res = self._conn.describe_collection(self._collection_name)
+        self._collection_id = res[COLLECTION_ID]
+
     def __seek_to_offset(self):
         # read pk cursor from cp file, no need to seek offset
         if self._next_id is not None:
@@ -502,6 +509,8 @@ def __init__(
         self.__check_for_special_index_param()
         self._kwargs = kwargs
         self._kwargs[ITERATOR_FIELD] = "True"
+        self.__set_up_collection_id()
+        self._kwargs[COLLECTION_ID] = self._collection_id
         self._filtered_ids = []
         self._filtered_distance = None
         self._schema = schema
@@ -513,6 +522,10 @@ def __init__(
         self.__setup__pk_prop()
         self.__init_search_iterator()
 
+    def __set_up_collection_id(self):
+        res = self._conn.describe_collection(self._collection_name)
+        self._collection_id = res[COLLECTION_ID]
+
     def __init_search_iterator(self):
         init_page = self.__execute_next_search(self._param, self._expr, False)
         self._session_ts = init_page.get_session_ts()