fix: iterator mismatch when alter alias and database(#2555) (#2568)

related: #2555 Signed-off-by: MrPresent-Han <[email protected]> Co-authored-by: MrPresent-Han <[email protected]>
milvus-io · Jan 15, 2025 · 5460c02 · 5460c02
1 parent 3a2abe0
commit 5460c02
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 0 deletions.
diff --git a/pymilvus/client/constants.py b/pymilvus/client/constants.py
@@ -9,6 +9,7 @@
 DEFAULT_CONSISTENCY_LEVEL = ConsistencyLevel.Bounded
 DEFAULT_RESOURCE_GROUP = "__default_resource_group"
 REDUCE_STOP_FOR_BEST = "reduce_stop_for_best"
+COLLECTION_ID = "collection_id"
 GROUP_BY_FIELD = "group_by_field"
 ITERATOR_FIELD = "iterator"
 PAGE_RETAIN_ORDER_FIELD = "page_retain_order"

diff --git a/pymilvus/client/prepare.py b/pymilvus/client/prepare.py
@@ -14,6 +14,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,
     GROUP_BY_FIELD,
     ITERATOR_FIELD,
@@ -861,6 +862,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)
+
         group_by_field = kwargs.get(GROUP_BY_FIELD)
         if group_by_field is not None:
             search_params[GROUP_BY_FIELD] = group_by_field
@@ -1145,6 +1150,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:

diff --git a/pymilvus/orm/constants.py b/pymilvus/orm/constants.py
@@ -39,6 +39,7 @@
 EF = "ef"
 IS_PRIMARY = "is_primary"
 REDUCE_STOP_FOR_BEST = "reduce_stop_for_best"
+COLLECTION_ID = "collection_id"
 ITERATOR_FIELD = "iterator"
 DEFAULT_MAX_L2_DISTANCE = 99999999.0
 DEFAULT_MIN_IP_DISTANCE = -99999999.0

diff --git a/pymilvus/orm/iterator.py b/pymilvus/orm/iterator.py
@@ -19,6 +19,7 @@
     CALC_DIST_JACCARD,
     CALC_DIST_L2,
     CALC_DIST_TANIMOTO,
+    COLLECTION_ID,
     DEFAULT_SEARCH_EXTENSION_RATE,
     EF,
     FIELDS,
@@ -72,11 +73,14 @@ 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._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()
@@ -87,6 +91,10 @@ def __init__(
         self.__seek()
         self._cache_id_in_use = NO_CACHE_ID
 
+    def __set_up_collection_id(self):
+        res = self._conn.describe_collection(self._collection_name)
+        self._collection_id = res[COLLECTION_ID]
+
     def __check_set_reduce_stop_for_best(self):
         if self._kwargs.get(REDUCE_STOP_FOR_BEST, True):
             self._kwargs[REDUCE_STOP_FOR_BEST] = "True"
@@ -336,11 +344,14 @@ def __init__(
             "timeout": timeout,
             "round_decimal": round_decimal,
         }
+        self._collection_name = collection_name
         self._expr = expr
         self.__check_set_params(param)
         self.__check_for_special_index_param()
         self._kwargs = kwargs
         self.__set_up_iteration_states()
+        self.__set_up_collection_id()
+        self._kwargs[COLLECTION_ID] = self._collection_id
         self._filtered_ids = []
         self._filtered_distance = None
         self._schema = schema
@@ -352,6 +363,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)
         if len(init_page) == 0: