feat: add page_retain_order param during search with offset

Signed-off-by: Patrick Weizhi Xu <[email protected]>

Author: PwzXxm

pymilvus/client/constants.py

@@ -12,6 +12,7 @@
 GROUP_BY_FIELD = "group_by_field"
 GROUP_SIZE = "group_size"
 ITERATOR_FIELD = "iterator"
+PAGE_RETAIN_ORDER_FIELD = "page_retain_order"
 
 RANKER_TYPE_RRF = "rrf"
 RANKER_TYPE_WEIGHTED = "weighted"

pymilvus/client/prepare.py

@@ -17,6 +17,7 @@
     GROUP_BY_FIELD,
     GROUP_SIZE,
     ITERATOR_FIELD,
+    PAGE_RETAIN_ORDER_FIELD,
     REDUCE_STOP_FOR_BEST,
 )
 from .types import (
@@ -690,6 +691,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

pymilvus/orm/collection.py

@@ -677,13 +677,16 @@ 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::
 
                     {
                         "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.
 

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,
                     }
 

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: