enhance: restrict max group size(#33544) (#36223)

related: #33544

Signed-off-by: MrPresent-Han <chun.han@gmail.com>
Co-authored-by: MrPresent-Han <chun.han@gmail.com>

configs/milvus.yaml

@@ -831,6 +831,7 @@ quotaAndLimits:
     maxCollectionNumPerDB: 65536 # Maximum number of collections per database.
     maxInsertSize: -1 # maximum size of a single insert request, in bytes, -1 means no limit
     maxResourceGroupNumOfQueryNode: 1024 # maximum number of resource groups of query nodes
+    maxGroupSize: 10 # maximum size for one single group when doing search group by
   ddl:
     enabled: false # Whether DDL request throttling is enabled.
     # Maximum number of collection-related DDL requests per second.

internal/proxy/search_util.go

@@ -172,6 +172,10 @@ func parseSearchInfo(searchParamsPair []*commonpb.KeyValuePair, schema *schemapb
 			groupSize = 1
 		}
 	}
+	if groupSize > Params.QuotaConfig.MaxGroupSize.GetAsInt64() {
+		return nil, 0, merr.WrapErrParameterInvalidMsg(
+			fmt.Sprintf("input group size:%d exceeds configured max group size:%d", groupSize, Params.QuotaConfig.MaxGroupSize.GetAsInt64()))
+	}
 
 	var groupStrictSize bool
 	groupStrictSizeStr, err := funcutil.GetAttrByKeyFromRepeatedKV(GroupStrictSize, searchParamsPair)

internal/proxy/task_search_test.go

@@ -19,6 +19,7 @@ import (
 	"context"
 	"fmt"
 	"strconv"
+	"strings"
 	"testing"
 	"time"
 
@@ -2430,6 +2431,31 @@ func TestTaskSearch_parseSearchInfo(t *testing.T) {
 		assert.NoError(t, err)
 		assert.Equal(t, Params.QuotaConfig.TopKLimit.GetAsInt64(), info.Topk)
 	})
+
+	t.Run("check max group size", func(t *testing.T) {
+		normalParam := getValidSearchParams()
+		normalParam = append(normalParam, &commonpb.KeyValuePair{
+			Key:   GroupSizeKey,
+			Value: "128",
+		})
+		fields := make([]*schemapb.FieldSchema, 0)
+		fields = append(fields, &schemapb.FieldSchema{
+			FieldID: int64(101),
+			Name:    "string_field",
+		})
+		schema := &schemapb.CollectionSchema{
+			Fields: fields,
+		}
+		info, _, err := parseSearchInfo(normalParam, schema, false)
+		assert.Nil(t, info)
+		assert.Error(t, err)
+		assert.True(t, strings.Contains(err.Error(), "exceeds configured max group size"))
+
+		resetSearchParamsValue(normalParam, GroupSizeKey, `10`)
+		info, _, err = parseSearchInfo(normalParam, schema, false)
+		assert.NotNil(t, info)
+		assert.NoError(t, err)
+	})
 }
 
 func getSearchResultData(nq, topk int64) *schemapb.SearchResultData {

pkg/util/paramtable/quota_param.go

@@ -132,6 +132,7 @@ type quotaConfig struct {
 	MaxOutputSize                  ParamItem `refreshable:"true"`
 	MaxInsertSize                  ParamItem `refreshable:"true"`
 	MaxResourceGroupNumOfQueryNode ParamItem `refreshable:"true"`
+	MaxGroupSize                   ParamItem `refreshable:"true"`
 
 	// limit writing
 	ForceDenyWriting                      ParamItem `refreshable:"true"`
@@ -1588,6 +1589,15 @@ Check https://milvus.io/docs/limitations.md for more details.`,
 	}
 	p.MaxResourceGroupNumOfQueryNode.Init(base.mgr)
 
+	p.MaxGroupSize = ParamItem{
+		Key:          "quotaAndLimits.limits.maxGroupSize",
+		Version:      "2.5.0",
+		Doc:          `maximum size for one single group when doing search group by`,
+		DefaultValue: "10",
+		Export:       true,
+	}
+	p.MaxGroupSize.Init(base.mgr)
+
 	// limit writing
 	p.ForceDenyWriting = ParamItem{
 		Key:          "quotaAndLimits.limitWriting.forceDeny",