enhance: support group_size and hybridsearch+groupby on httpv2 side(m…

…ilvus-io#36386) Signed-off-by: MrPresent-Han <[email protected]>
MrPresent-Han · Sep 24, 2024 · 2a3081a · 2a3081a
1 parent ed95568
commit 2a3081a
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 10 deletions.
diff --git a/internal/distributed/proxy/httpserver/constant.go b/internal/distributed/proxy/httpserver/constant.go
@@ -129,13 +129,15 @@ const (
 )
 
 const (
-	ParamAnnsField    = "anns_field"
-	Params            = "params"
-	ParamRoundDecimal = "round_decimal"
-	ParamOffset       = "offset"
-	ParamLimit        = "limit"
-	ParamRadius       = "radius"
-	ParamRangeFilter  = "range_filter"
-	ParamGroupByField = "group_by_field"
-	BoundedTimestamp  = 2
+	ParamAnnsField       = "anns_field"
+	Params               = "params"
+	ParamRoundDecimal    = "round_decimal"
+	ParamOffset          = "offset"
+	ParamLimit           = "limit"
+	ParamRadius          = "radius"
+	ParamRangeFilter     = "range_filter"
+	ParamGroupByField    = "group_by_field"
+	ParamGroupSize       = "group_size"
+	ParamGroupStrictSize = "group_strict_size"
+	BoundedTimestamp     = 2
 )
diff --git a/internal/distributed/proxy/httpserver/handler_v2.go b/internal/distributed/proxy/httpserver/handler_v2.go
@@ -944,6 +944,8 @@ func (h *HandlersV2) search(ctx context.Context, c *gin.Context, anyReq any, dbN
 	searchParams = append(searchParams, &commonpb.KeyValuePair{Key: common.TopKKey, Value: strconv.FormatInt(int64(httpReq.Limit), 10)})
 	searchParams = append(searchParams, &commonpb.KeyValuePair{Key: ParamOffset, Value: strconv.FormatInt(int64(httpReq.Offset), 10)})
 	searchParams = append(searchParams, &commonpb.KeyValuePair{Key: ParamGroupByField, Value: httpReq.GroupByField})
+	searchParams = append(searchParams, &commonpb.KeyValuePair{Key: ParamGroupSize, Value: strconv.FormatInt(int64(httpReq.GroupSize), 10)})
+	searchParams = append(searchParams, &commonpb.KeyValuePair{Key: ParamGroupStrictSize, Value: strconv.FormatBool(httpReq.GroupStrictSize)})
 	searchParams = append(searchParams, &commonpb.KeyValuePair{Key: proxy.AnnsFieldKey, Value: httpReq.AnnsField})
 	body, _ := c.Get(gin.BodyBytesKey)
 	placeholderGroup, err := generatePlaceholderGroup(ctx, string(body.([]byte)), collSchema, httpReq.AnnsField)
@@ -1010,7 +1012,6 @@ func (h *HandlersV2) advancedSearch(ctx context.Context, c *gin.Context, anyReq
 		}
 		searchParams = append(searchParams, &commonpb.KeyValuePair{Key: common.TopKKey, Value: strconv.FormatInt(int64(subReq.Limit), 10)})
 		searchParams = append(searchParams, &commonpb.KeyValuePair{Key: ParamOffset, Value: strconv.FormatInt(int64(subReq.Offset), 10)})
-		searchParams = append(searchParams, &commonpb.KeyValuePair{Key: ParamGroupByField, Value: subReq.GroupByField})
 		searchParams = append(searchParams, &commonpb.KeyValuePair{Key: proxy.AnnsFieldKey, Value: subReq.AnnsField})
 		placeholderGroup, err := generatePlaceholderGroup(ctx, searchArray[i].Raw, collSchema, subReq.AnnsField)
 		if err != nil {
@@ -1039,6 +1040,9 @@ func (h *HandlersV2) advancedSearch(ctx context.Context, c *gin.Context, anyReq
 		{Key: proxy.RankParamsKey, Value: string(bs)},
 		{Key: ParamLimit, Value: strconv.FormatInt(int64(httpReq.Limit), 10)},
 		{Key: ParamRoundDecimal, Value: "-1"},
+		{Key: ParamGroupByField, Value: httpReq.GroupByField},
+		{Key: ParamGroupSize, Value: strconv.FormatInt(int64(httpReq.GroupSize), 10)},
+		{Key: ParamGroupStrictSize, Value: strconv.FormatBool(httpReq.GroupStrictSize)},
 	}
 	resp, err := wrapperProxyWithLimit(ctx, c, req, h.checkAuth, false, "/milvus.proto.milvus.MilvusService/HybridSearch", true, h.proxy, func(reqCtx context.Context, req any) (interface{}, error) {
 		return h.proxy.HybridSearch(reqCtx, req.(*milvuspb.HybridSearchRequest))

diff --git a/internal/distributed/proxy/httpserver/request_v2.go b/internal/distributed/proxy/httpserver/request_v2.go
@@ -156,6 +156,8 @@ type SearchReqV2 struct {
 	PartitionNames   []string      `json:"partitionNames"`
 	Filter           string        `json:"filter"`
 	GroupByField     string        `json:"groupingField"`
+	GroupSize        int32         `json:"groupSize"`
+	GroupStrictSize  bool          `json:"groupStrictSize"`
 	Limit            int32         `json:"limit"`
 	Offset           int32         `json:"offset"`
 	OutputFields     []string      `json:"outputFields"`
@@ -190,6 +192,9 @@ type HybridSearchReq struct {
 	Search           []SubSearchReq `json:"search"`
 	Rerank           Rand           `json:"rerank"`
 	Limit            int32          `json:"limit"`
+	GroupByField     string         `json:"groupingField"`
+	GroupSize        int32          `json:"groupSize"`
+	GroupStrictSize  bool           `json:"groupStrictSize"`
 	OutputFields     []string       `json:"outputFields"`
 	ConsistencyLevel string         `json:"consistencyLevel"`
 }