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 · f004f8e · f004f8e
1 parent 51cdee1
commit f004f8e
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 1 deletion.
diff --git a/internal/distributed/proxy/httpserver/constant.go b/internal/distributed/proxy/httpserver/constant.go
@@ -137,5 +137,6 @@ const (
 	ParamRadius       = "radius"
 	ParamRangeFilter  = "range_filter"
 	ParamGroupByField = "group_by_field"
+	ParamGroupSize    = "group_size"
 	BoundedTimestamp  = 2
 )
diff --git a/internal/distributed/proxy/httpserver/handler_v2.go b/internal/distributed/proxy/httpserver/handler_v2.go
@@ -944,6 +944,7 @@ 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: proxy.AnnsFieldKey, Value: httpReq.AnnsField})
 	body, _ := c.Get(gin.BodyBytesKey)
 	placeholderGroup, err := generatePlaceholderGroup(ctx, string(body.([]byte)), collSchema, httpReq.AnnsField)
@@ -1010,7 +1011,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 +1039,8 @@ 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)},
 	}
 	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,7 @@ type SearchReqV2 struct {
 	PartitionNames   []string      `json:"partitionNames"`
 	Filter           string        `json:"filter"`
 	GroupByField     string        `json:"groupingField"`
+	GroupSize        int32         `json:"groupSize"`
 	Limit            int32         `json:"limit"`
 	Offset           int32         `json:"offset"`
 	OutputFields     []string      `json:"outputFields"`
@@ -190,6 +191,8 @@ type HybridSearchReq struct {
 	Search           []SubSearchReq `json:"search"`
 	Rerank           Rand           `json:"rerank"`
 	Limit            int32          `json:"limit"`
+	GroupByField   string         `json:"groupingField"`
+	GroupSize      int32          `json:"groupSize"`
 	OutputFields     []string       `json:"outputFields"`
 	ConsistencyLevel string         `json:"consistencyLevel"`
 }