diff --git a/pymilvus/client/constants.py b/pymilvus/client/constants.py index 0d31c7af1..733ddf461 100644 --- a/pymilvus/client/constants.py +++ b/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" diff --git a/pymilvus/client/prepare.py b/pymilvus/client/prepare.py index f0210b69d..5b17bc726 100644 --- a/pymilvus/client/prepare.py +++ b/pymilvus/client/prepare.py @@ -28,6 +28,7 @@ PAGE_RETAIN_ORDER_FIELD, RANK_GROUP_SCORER, REDUCE_STOP_FOR_BEST, + COLLECTION_ID, STRICT_GROUP_SIZE, ) from .types import ( @@ -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: diff --git a/pymilvus/client/search_iterator.py b/pymilvus/client/search_iterator.py index cf87a43f8..25dcd8995 100644 --- a/pymilvus/client/search_iterator.py +++ b/pymilvus/client/search_iterator.py @@ -11,6 +11,7 @@ ITER_SEARCH_LAST_BOUND_KEY, ITER_SEARCH_V2_KEY, ITERATOR_FIELD, + COLLECTION_ID, ) from pymilvus.exceptions import ExceptionsMessage, ParamError, ServerVersionIncompatibleException from pymilvus.orm.connections import Connections @@ -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: diff --git a/pymilvus/orm/constants.py b/pymilvus/orm/constants.py index cf6e2ebd4..1adae5268 100644 --- a/pymilvus/orm/constants.py +++ b/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 diff --git a/pymilvus/orm/iterator.py b/pymilvus/orm/iterator.py index 6bd6115a5..1787bf5d1 100644 --- a/pymilvus/orm/iterator.py +++ b/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,6 +102,7 @@ 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 @@ -108,6 +110,7 @@ def __init__( 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()