From fda125ff595440653685660a8670c6503641abe1 Mon Sep 17 00:00:00 2001 From: Patrick Weizhi Xu Date: Fri, 16 Aug 2024 16:52:52 +0800 Subject: [PATCH] feat: [2.4] add page_retain_order param during search with offset (#2230) issue: https://github.com/milvus-io/milvus/issues/35464 --- pymilvus/client/constants.py | 1 + pymilvus/client/prepare.py | 15 +++++++++++++++ pymilvus/orm/collection.py | 6 ++++++ pymilvus/orm/partition.py | 4 ++++ tests/test_prepare.py | 6 ++++++ 5 files changed, 32 insertions(+) diff --git a/pymilvus/client/constants.py b/pymilvus/client/constants.py index fc1b31fa3..5bf4f9af1 100644 --- a/pymilvus/client/constants.py +++ b/pymilvus/client/constants.py @@ -11,6 +11,7 @@ REDUCE_STOP_FOR_BEST = "reduce_stop_for_best" GROUP_BY_FIELD = "group_by_field" ITERATOR_FIELD = "iterator" +PAGE_RETAIN_ORDER_FIELD = "page_retain_order" RANKER_TYPE_RRF = "rrf" RANKER_TYPE_WEIGHTED = "weighted" diff --git a/pymilvus/client/prepare.py b/pymilvus/client/prepare.py index e63b47a50..d4ba91a14 100644 --- a/pymilvus/client/prepare.py +++ b/pymilvus/client/prepare.py @@ -16,6 +16,7 @@ DEFAULT_CONSISTENCY_LEVEL, GROUP_BY_FIELD, ITERATOR_FIELD, + PAGE_RETAIN_ORDER_FIELD, REDUCE_STOP_FOR_BEST, ) from .types import ( @@ -689,6 +690,20 @@ def search_requests_with_expr( raise ParamError(message=f"wrong type for offset, expect int, got {type(offset)}") search_params["offset"] = offset + if PAGE_RETAIN_ORDER_FIELD in kwargs and PAGE_RETAIN_ORDER_FIELD in param: + raise ParamError( + message="Provide page_retain_order both in kwargs and param, expect just one" + ) + page_retain_order = kwargs.get(PAGE_RETAIN_ORDER_FIELD) or param.get( + PAGE_RETAIN_ORDER_FIELD + ) + if page_retain_order is not None: + if not isinstance(page_retain_order, bool): + raise ParamError( + message=f"wrong type for page_retain_order, expect bool, got {type(page_retain_order)}" + ) + search_params[PAGE_RETAIN_ORDER_FIELD] = page_retain_order + is_iterator = kwargs.get(ITERATOR_FIELD) if is_iterator is not None: search_params[ITERATOR_FIELD] = is_iterator diff --git a/pymilvus/orm/collection.py b/pymilvus/orm/collection.py index 31f93aa00..36e8f45f5 100644 --- a/pymilvus/orm/collection.py +++ b/pymilvus/orm/collection.py @@ -677,6 +677,8 @@ def search( similar metricy types, the value must be of type str. * *offset* (``int``, optional) offset for pagination. + * *page_retain_order* (``bool``, optional) + Whether to retain the order of the search result when offset is provided. * *params of index: *nprobe*, *ef*, *search_k*, etc Corresponding search params for a certain index. example for param:: @@ -684,6 +686,7 @@ def search( { "metric_type": "L2", "offset": 10, + "page_retain_order": True, "params": {"nprobe": 12}, } @@ -716,6 +719,9 @@ def search( * *offset* (``int``, optinal) offset for pagination. + * *page_retain_order* (``bool``, optional) + Whether to retain the order of the search result when offset is provided. + * *consistency_level* (``str/int``, optional) Which consistency level to use when searching in the collection. diff --git a/pymilvus/orm/partition.py b/pymilvus/orm/partition.py index e860f4cad..fca4cbd68 100644 --- a/pymilvus/orm/partition.py +++ b/pymilvus/orm/partition.py @@ -387,6 +387,9 @@ def search( * *offset* (``int``, optional) offset for pagination. + * *page_retain_order* (``bool``, optional) + Whether to retain the order of the search result when offset is provided. + * *limit* (``int``, optional) limit for the search results and pagination. @@ -396,6 +399,7 @@ def search( "nprobe": 128, "metric_type": "L2", "offset": 10, + "page_retain_order": True, "limit": 10, } diff --git a/tests/test_prepare.py b/tests/test_prepare.py index 363d316c0..d292409fb 100644 --- a/tests/test_prepare.py +++ b/tests/test_prepare.py @@ -24,17 +24,23 @@ def test_search_requests_with_expr_offset(self): search_params = { "metric_type": "L2", "offset": 10, + "page_retain_order": True, } ret = Prepare.search_requests_with_expr("name", data, "v", search_params, 100) offset_exists = False + page_retain_order_exists = False for p in ret.search_params: if p.key == "offset": offset_exists = True assert p.value == "10" + elif p.key == "page_retain_order": + page_retain_order_exists = True + assert p.value == "True" # it was dumped as string assert offset_exists is True + assert page_retain_order_exists is True class TestCreateCollectionRequest: