feat: support group_size parameter for search_group_by (milvus-io#2130)

related: milvus-io#2129 Signed-off-by: MrPresent-Han <[email protected]> Co-authored-by: MrPresent-Han <[email protected]>
zh19990906 · Jun 24, 2024 · 7c09c5c · 7c09c5c
1 parent 08eff03
commit 7c09c5c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/pymilvus/client/constants.py b/pymilvus/client/constants.py
@@ -10,6 +10,7 @@
 DEFAULT_RESOURCE_GROUP = "__default_resource_group"
 REDUCE_STOP_FOR_BEST = "reduce_stop_for_best"
 GROUP_BY_FIELD = "group_by_field"
+GROUP_SIZE = "group_size"
 ITERATOR_FIELD = "iterator"
 
 RANKER_TYPE_RRF = "rrf"

diff --git a/pymilvus/client/prepare.py b/pymilvus/client/prepare.py
@@ -15,6 +15,7 @@
 from .constants import (
     DEFAULT_CONSISTENCY_LEVEL,
     GROUP_BY_FIELD,
+    GROUP_SIZE,
     ITERATOR_FIELD,
     REDUCE_STOP_FOR_BEST,
 )
@@ -667,6 +668,10 @@ def search_requests_with_expr(
         if group_by_field is not None:
             search_params[GROUP_BY_FIELD] = group_by_field
 
+        group_size = kwargs.get(GROUP_SIZE)
+        if group_size is not None:
+            search_params[GROUP_SIZE] = group_size
+
         if param.get("metric_type") is not None:
             search_params["metric_type"] = param["metric_type"]