From baa697f8041d9121ee0cdfdc67298025e676be10 Mon Sep 17 00:00:00 2001 From: MrPresent-Han Date: Tue, 14 Jan 2025 02:53:13 -0500 Subject: [PATCH] fix: iterator mismatch when alter alias and database(#2555) Signed-off-by: MrPresent-Han --- pymilvus/client/constants.py | 1 + pymilvus/client/prepare.py | 10 ++++++++++ pymilvus/orm/constants.py | 1 + pymilvus/orm/iterator.py | 15 +++++++++++++++ 4 files changed, 27 insertions(+) diff --git a/pymilvus/client/constants.py b/pymilvus/client/constants.py index 5bf4f9af1..f45a42b46 100644 --- 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 index 1630f4cc6..ab5e12700 100644 --- 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 index b5a88ecb4..7c1c73442 100644 --- 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 index fef5dead2..67ffb85ae 100644 --- 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: