From f004f8e718890b404d954284deb35a6dc22c1588 Mon Sep 17 00:00:00 2001 From: MrPresent-Han Date: Thu, 19 Sep 2024 23:31:54 -0400 Subject: [PATCH] enhance: support group_size and hybridsearch+groupby on httpv2 side(#36386) Signed-off-by: MrPresent-Han --- internal/distributed/proxy/httpserver/constant.go | 1 + internal/distributed/proxy/httpserver/handler_v2.go | 4 +++- internal/distributed/proxy/httpserver/request_v2.go | 3 +++ 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/internal/distributed/proxy/httpserver/constant.go b/internal/distributed/proxy/httpserver/constant.go index f106e52787312..29d1b1d28ff0d 100644 --- 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 index 80a9b7e4829cd..31474acc20a52 100644 --- 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 index 650924f2ac337..9874e69c6da96 100644 --- 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"` }