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

…ilvus-io#36386) Signed-off-by: MrPresent-Han <[email protected]>
MrPresent-Han · Sep 20, 2024 · 42b6c3c · 42b6c3c
1 parent 3b10085
commit 42b6c3c
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 4 deletions.
diff --git a/internal/datacoord/import_util.go b/internal/datacoord/import_util.go
@@ -23,14 +23,13 @@ import (
 	"sort"
 	"time"
 
-	"github.com/milvus-io/milvus/internal/proto/indexpb"
-
 	"github.com/cockroachdb/errors"
 	"github.com/samber/lo"
 	"go.uber.org/zap"
 
 	"github.com/milvus-io/milvus/internal/datacoord/allocator"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
+	"github.com/milvus-io/milvus/internal/proto/indexpb"
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/storage"
 	"github.com/milvus-io/milvus/internal/util/importutilv2"

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
@@ -958,6 +958,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})
 	searchParams = append(searchParams, &commonpb.KeyValuePair{Key: ParamRoundDecimal, Value: "-1"})
 	body, _ := c.Get(gin.BodyBytesKey)
@@ -1020,7 +1021,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})
 		searchParams = append(searchParams, &commonpb.KeyValuePair{Key: ParamRoundDecimal, Value: "-1"})
 		placeholderGroup, err := generatePlaceholderGroup(ctx, searchArray[i].Raw, collSchema, subReq.AnnsField)
@@ -1051,6 +1051,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
@@ -149,6 +149,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"`
@@ -166,7 +167,6 @@ type SubSearchReq struct {
 	Data          []interface{}      `json:"data" binding:"required"`
 	AnnsField     string             `json:"annsField"`
 	Filter        string             `json:"filter"`
-	GroupByField  string             `json:"groupingField"`
 	MetricType    string             `json:"metricType"`
 	Limit         int32              `json:"limit"`
 	Offset        int32              `json:"offset"`
@@ -180,6 +180,8 @@ type HybridSearchReq struct {
 	PartitionNames []string       `json:"partitionNames"`
 	Search         []SubSearchReq `json:"search"`
 	Rerank         Rand           `json:"rerank"`
+	GroupByField   string         `json:"groupingField"`
+	GroupSize      int32          `json:"groupSize"`
 	Limit          int32          `json:"limit"`
 	OutputFields   []string       `json:"outputFields"`
 }