milvus/internal/querynodev2/segments/reducer.go

40 lines
1.4 KiB
Go

package segments
import (
"context"
"github.com/milvus-io/milvus/internal/proto/segcorepb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/proto/querypb"
)
type internalReducer interface {
Reduce(context.Context, []*internalpb.RetrieveResults) (*internalpb.RetrieveResults, error)
}
func CreateInternalReducer(req *querypb.QueryRequest, schema *schemapb.CollectionSchema) internalReducer {
if req.GetReq().GetIsCount() {
return &cntReducer{}
} else if req.GetReq().GetIterationExtensionReduceRate() > 0 {
extendedLimit := req.GetReq().GetIterationExtensionReduceRate() * req.GetReq().Limit
return newExtensionLimitReducer(req, schema, extendedLimit)
}
return newDefaultLimitReducer(req, schema)
}
type segCoreReducer interface {
Reduce(context.Context, []*segcorepb.RetrieveResults) (*segcorepb.RetrieveResults, error)
}
func CreateSegCoreReducer(req *querypb.QueryRequest, schema *schemapb.CollectionSchema) segCoreReducer {
if req.GetReq().GetIsCount() {
return &cntReducerSegCore{}
} else if req.GetReq().GetIterationExtensionReduceRate() > 0 {
extendedLimit := req.GetReq().GetIterationExtensionReduceRate() * req.GetReq().Limit
return newExtensionLimitSegcoreReducer(req, schema, extendedLimit)
}
return newDefaultLimitReducerSegcore(req, schema)
}