From dfd9873a4b97dfd66540e30699023deb54ff234c Mon Sep 17 00:00:00 2001 From: junjiejiangjjj Date: Fri, 17 Jul 2026 23:10:39 +0800 Subject: [PATCH] fix: [cp 3.0] cover boost score in L0 rerank latency metric (#51517) pr: https://github.com/milvus-io/milvus/pull/51347 issue: https://github.com/milvus-io/milvus/issues/51310 https://github.com/milvus-io/milvus/issues/51306 Signed-off-by: junjie.jiang --- internal/proxy/task_search.go | 7 ++ internal/proxy/task_search_test.go | 70 ++++++++--- .../querynodev2/tasks/l0_function_chain.go | 26 ++++- .../tasks/l0_function_chain_test.go | 109 ++++++++++++++++++ pkg/metrics/metrics.go | 3 + pkg/metrics/querynode_metrics.go | 14 +++ 6 files changed, 210 insertions(+), 19 deletions(-) diff --git a/internal/proxy/task_search.go b/internal/proxy/task_search.go index 11237faa91..3e11dae0eb 100644 --- a/internal/proxy/task_search.go +++ b/internal/proxy/task_search.go @@ -911,6 +911,13 @@ func (t *searchTask) initSearchRequest(ctx context.Context) error { t.rerankMeta = newRerankMeta(t.schema.CollectionSchema, t.request.FunctionScore) } + // Search iterator v2 uses the final result score as the ANN continuation + // bound. Function rerank rewrites that score, so the next iterator request + // would interpret a rerank score in the ANN metric domain. + if queryInfo.GetSearchIteratorV2Info() != nil && (t.rerankMeta != nil || len(querynodeFunctionChains) > 0) { + return merr.WrapErrParameterInvalidMsg("function rerank is not supported with search iterator v2") + } + // order_by and function rerank cannot be used together if len(t.orderByFields) > 0 && (t.rerankMeta != nil || len(querynodeFunctionChains) > 0) { return merr.WrapErrParameterInvalidMsg("order_by and function rerank cannot be used together: they specify conflicting sort criteria") diff --git a/internal/proxy/task_search_test.go b/internal/proxy/task_search_test.go index 0f99ecaca5..aa638a7ff5 100644 --- a/internal/proxy/task_search_test.go +++ b/internal/proxy/task_search_test.go @@ -5504,6 +5504,33 @@ func TestSearchTask_FunctionChainRerankMeta(t *testing.T) { request.FunctionChains = []*schemapb.FunctionChain{l0Chain, l2Chain} return request, l0Chain, l2Chain } + newFunctionScoreRequest := func() *milvuspb.SearchRequest { + request := newRequest() + request.FunctionScore = &schemapb.FunctionScore{ + Functions: []*schemapb.FunctionSchema{ + { + Name: "decay", + Type: schemapb.FunctionType_Rerank, + InputFieldNames: []string{"ts"}, + OutputFieldNames: []string{}, + Params: []*commonpb.KeyValuePair{ + {Key: "reranker", Value: "decay"}, + {Key: "origin", Value: "100"}, + {Key: "scale", Value: "10"}, + }, + }, + }, + } + return request + } + withSearchIteratorV2 := func(request *milvuspb.SearchRequest) *milvuspb.SearchRequest { + request.SearchParams = append(request.SearchParams, + &commonpb.KeyValuePair{Key: IteratorField, Value: "true"}, + &commonpb.KeyValuePair{Key: SearchIterV2Key, Value: "true"}, + &commonpb.KeyValuePair{Key: SearchIterBatchSizeKey, Value: "10"}, + ) + return request + } newTask := func(request *milvuspb.SearchRequest) *searchTask { translatedOutputFields, _, _, _, _, err := translateOutputFields(request.GetOutputFields(), schemaInfo, true) require.NoError(t, err) @@ -5581,23 +5608,7 @@ func TestSearchTask_FunctionChainRerankMeta(t *testing.T) { }) t.Run("ordinary search keeps function score rerank meta", func(t *testing.T) { - request := newRequest() - request.FunctionScore = &schemapb.FunctionScore{ - Functions: []*schemapb.FunctionSchema{ - { - Name: "decay", - Type: schemapb.FunctionType_Rerank, - InputFieldNames: []string{"ts"}, - OutputFieldNames: []string{}, - Params: []*commonpb.KeyValuePair{ - {Key: "reranker", Value: "decay"}, - {Key: "origin", Value: "100"}, - {Key: "scale", Value: "10"}, - }, - }, - }, - } - task := newTask(request) + task := newTask(newFunctionScoreRequest()) require.NoError(t, task.initSearchRequest(ctx)) meta, ok := task.rerankMeta.(*funcScoreRerankMeta) @@ -5606,6 +5617,31 @@ func TestSearchTask_FunctionChainRerankMeta(t *testing.T) { assert.Equal(t, []int64{101}, meta.GetInputFieldIDs()) }) + t.Run("search iterator v2 rejects function score", func(t *testing.T) { + task := newTask(withSearchIteratorV2(newFunctionScoreRequest())) + + err := task.initSearchRequest(ctx) + require.Error(t, err) + assert.Contains(t, err.Error(), "function rerank is not supported with search iterator v2") + }) + + t.Run("search iterator v2 rejects l2 function chain", func(t *testing.T) { + task := newTask(withSearchIteratorV2(newFunctionChainRequest())) + + err := task.initSearchRequest(ctx) + require.Error(t, err) + assert.Contains(t, err.Error(), "function rerank is not supported with search iterator v2") + }) + + t.Run("search iterator v2 rejects l0 function chain", func(t *testing.T) { + request, _ := newL0FunctionChainRequest() + task := newTask(withSearchIteratorV2(request)) + + err := task.initSearchRequest(ctx) + require.Error(t, err) + assert.Contains(t, err.Error(), "function rerank is not supported with search iterator v2") + }) + t.Run("ordinary search routes l0 chain to querynode plan", func(t *testing.T) { request, chainPB := newL0FunctionChainRequest() task := newTask(request) diff --git a/internal/querynodev2/tasks/l0_function_chain.go b/internal/querynodev2/tasks/l0_function_chain.go index f902dcdd6f..0e6f6a29ae 100644 --- a/internal/querynodev2/tasks/l0_function_chain.go +++ b/internal/querynodev2/tasks/l0_function_chain.go @@ -19,6 +19,7 @@ package tasks import ( "context" "fmt" + "time" "golang.org/x/sync/errgroup" @@ -27,6 +28,7 @@ import ( "github.com/milvus-io/milvus/internal/util/function/chain" chaintypes "github.com/milvus-io/milvus/internal/util/function/chain/types" "github.com/milvus-io/milvus/internal/util/segcore" + "github.com/milvus-io/milvus/pkg/v3/metrics" "github.com/milvus-io/milvus/pkg/v3/proto/planpb" "github.com/milvus-io/milvus/pkg/v3/util/merr" "github.com/milvus-io/milvus/pkg/v3/util/typeutil" @@ -160,11 +162,31 @@ func executeL0RerankChains(ctx context.Context, segDFs []*chain.DataFrame, build return nil } -func (t *SearchTask) applyL0Rerank(segDFs []*chain.DataFrame, prepared *preparedQueryNodeFunctionChains, searchedSegments []segments.Segment, searchReq *segcore.SearchRequest) error { +func (t *SearchTask) applyL0Rerank(segDFs []*chain.DataFrame, prepared *preparedQueryNodeFunctionChains, searchedSegments []segments.Segment, searchReq *segcore.SearchRequest) (retErr error) { if prepared == nil { return merr.WrapErrServiceInternalMsg("l0_rerank: prepared querynode function chains is nil") } - if len(prepared.l0Chains) > 0 { + + hasPublicL0 := len(prepared.l0Chains) > 0 + hasBoostScore := prepared.plan != nil && len(prepared.plan.GetScorers()) > 0 + if !hasPublicL0 && !hasBoostScore { + return t.applyBoostScoresWithPlan(segDFs, prepared.plan, searchedSegments, searchReq) + } + + start := time.Now() + defer func() { + status := metrics.SuccessLabel + if retErr != nil { + status = metrics.FailLabel + } + metrics.QueryNodeFunctionChainLatency.WithLabelValues( + fmt.Sprint(t.GetNodeID()), + metrics.FunctionChainLevelL0, + status, + ).Observe(float64(time.Since(start).Microseconds()) / 1000.0) + }() + + if hasPublicL0 { return t.applyPublicL0Rerank(segDFs, prepared) } return t.applyBoostScoresWithPlan(segDFs, prepared.plan, searchedSegments, searchReq) diff --git a/internal/querynodev2/tasks/l0_function_chain_test.go b/internal/querynodev2/tasks/l0_function_chain_test.go index f5e905cf25..ca1d8cac1b 100644 --- a/internal/querynodev2/tasks/l0_function_chain_test.go +++ b/internal/querynodev2/tasks/l0_function_chain_test.go @@ -17,21 +17,33 @@ package tasks import ( + "fmt" "testing" "github.com/apache/arrow/go/v17/arrow" "github.com/apache/arrow/go/v17/arrow/array" + "github.com/prometheus/client_golang/prometheus" + dto "github.com/prometheus/client_model/go" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/milvus-io/milvus-proto/go-api/v3/commonpb" "github.com/milvus-io/milvus-proto/go-api/v3/schemapb" + "github.com/milvus-io/milvus/internal/querynodev2/segments" "github.com/milvus-io/milvus/internal/util/function/chain" chainexpr "github.com/milvus-io/milvus/internal/util/function/chain/expr" "github.com/milvus-io/milvus/internal/util/function/chain/types" + "github.com/milvus-io/milvus/pkg/v3/metrics" "github.com/milvus-io/milvus/pkg/v3/proto/planpb" ) +func histogramSampleCount(t *testing.T, observer prometheus.Observer) uint64 { + t.Helper() + metric := &dto.Metric{} + require.NoError(t, observer.(prometheus.Metric).Write(metric)) + return metric.GetHistogram().GetSampleCount() +} + func TestPrepareQueryNodeFunctionChainsFromPlan(t *testing.T) { schema := &schemapb.CollectionSchema{ Fields: []*schemapb.FieldSchema{ @@ -264,6 +276,103 @@ func TestApplyPublicL0RerankPrunesInputsAndPreservesReduceSystemColumns(t *testi require.InDelta(t, 0.6, scores.Value(2), 1e-6) } +func TestApplyL0RerankMetrics(t *testing.T) { + withBoostScoreCheckedAllocator(t) + + repr, err := chain.ProtoChainToRepr(l0FunctionChainForTest(mapOpWithParamsForTest( + types.ScoreFieldName, + chainexpr.NumCombineFuncName, + map[string]*schemapb.FunctionParamValue{ + types.NumCombineParamMode: stringParamForTest(types.NumCombineModeSum), + }, + columnArgForTest(types.ScoreFieldName), + columnArgForTest(types.IDFieldName), + ))) + require.NoError(t, err) + publicPrepared := &preparedQueryNodeFunctionChains{l0Chains: []*chain.ChainRepr{repr}} + boostPlan := &planpb.PlanNode{ + Scorers: []*planpb.ScoreFunction{{Weight: 1}}, + ScoreOption: &planpb.ScoreOption{ + FunctionMode: planpb.FunctionMode_FunctionModeSum, + BoostMode: planpb.BoostMode_BoostModeMultiply, + }, + } + boostPrepared := &preparedQueryNodeFunctionChains{plan: boostPlan} + task := &SearchTask{ctx: t.Context()} + nodeID := fmt.Sprint(task.GetNodeID()) + + successObserver := metrics.QueryNodeFunctionChainLatency.WithLabelValues( + nodeID, + metrics.FunctionChainLevelL0, + metrics.SuccessLabel, + ) + failObserver := metrics.QueryNodeFunctionChainLatency.WithLabelValues( + nodeID, + metrics.FunctionChainLevelL0, + metrics.FailLabel, + ) + + t.Run("public L0 success", func(t *testing.T) { + before := histogramSampleCount(t, successObserver) + segDFs := []*chain.DataFrame{ + makeBoostScoreTestDF(t, []int64{1, 2}, []float32{0.5, 0.2}, []int64{10, 20}, []int64{2}), + makeBoostScoreTestDF(t, []int64{3, 4}, []float32{0.9, 0.1}, []int64{30, 40}, []int64{2}), + } + defer func() { + for _, df := range segDFs { + df.Release() + } + }() + + require.NoError(t, task.applyL0Rerank(segDFs, publicPrepared, nil, nil)) + require.Equal(t, before+1, histogramSampleCount(t, successObserver)) + }) + + t.Run("public L0 failure", func(t *testing.T) { + before := histogramSampleCount(t, failObserver) + err := task.applyL0Rerank([]*chain.DataFrame{nil}, publicPrepared, nil, nil) + require.Error(t, err) + require.Equal(t, before+1, histogramSampleCount(t, failObserver)) + }) + + t.Run("boost score success", func(t *testing.T) { + oldFactory := boostScoreRunnerFactory + boostScoreRunnerFactory = mockBoostScoreRunnerFactory(boostScoreOutput{ + scores: []float32{2.0}, + hasScore: []bool{true}, + }) + defer func() { boostScoreRunnerFactory = oldFactory }() + + before := histogramSampleCount(t, successObserver) + segDFs := []*chain.DataFrame{ + makeBoostScoreTestDF(t, []int64{1}, []float32{0.5}, []int64{10}, []int64{1}), + } + defer func() { + for _, df := range segDFs { + df.Release() + } + }() + + require.NoError(t, task.applyL0Rerank(segDFs, boostPrepared, []segments.Segment{nil}, nil)) + require.Equal(t, before+1, histogramSampleCount(t, successObserver)) + }) + + t.Run("boost score failure", func(t *testing.T) { + before := histogramSampleCount(t, failObserver) + err := task.applyL0Rerank(nil, boostPrepared, []segments.Segment{nil}, nil) + require.Error(t, err) + require.Equal(t, before+1, histogramSampleCount(t, failObserver)) + }) + + t.Run("no rerank does not record", func(t *testing.T) { + successBefore := histogramSampleCount(t, successObserver) + failBefore := histogramSampleCount(t, failObserver) + require.NoError(t, task.applyL0Rerank(nil, &preparedQueryNodeFunctionChains{}, nil, nil)) + require.Equal(t, successBefore, histogramSampleCount(t, successObserver)) + require.Equal(t, failBefore, histogramSampleCount(t, failObserver)) + }) +} + func l0FunctionChainForTest(ops ...*schemapb.FunctionChainOp) *schemapb.FunctionChain { return &schemapb.FunctionChain{ Stage: schemapb.FunctionChainStage_FunctionChainStageL0Rerank, diff --git a/pkg/metrics/metrics.go b/pkg/metrics/metrics.go index 71098ac81c..e5ee3aa165 100644 --- a/pkg/metrics/metrics.go +++ b/pkg/metrics/metrics.go @@ -103,6 +103,8 @@ const ( BatchReduce = "batch_reduce" + FunctionChainLevelL0 = "l0" + Pending = "pending" Executing = "executing" Done = "done" @@ -128,6 +130,7 @@ const ( channelNameLabelName = "channel_name" functionLabelName = "function_name" queryTypeLabelName = "query_type" + chainLevelLabelName = "chain_level" collectionName = "collection_name" databaseLabelName = "db_name" ResourceGroupLabelName = "rg" diff --git a/pkg/metrics/querynode_metrics.go b/pkg/metrics/querynode_metrics.go index 74e3f2b7e8..9dd711906a 100644 --- a/pkg/metrics/querynode_metrics.go +++ b/pkg/metrics/querynode_metrics.go @@ -289,6 +289,19 @@ var ( reduceType, }) + QueryNodeFunctionChainLatency = prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Namespace: milvusNamespace, + Subsystem: typeutil.QueryNodeRole, + Name: "function_chain_latency", + Help: "query-level function chain execution latency in milliseconds", + Buckets: subMsBuckets, + }, []string{ + nodeIDLabelName, + chainLevelLabelName, + statusLabelName, + }) + QueryNodeLoadSegmentLatency = prometheus.NewHistogramVec( prometheus.HistogramOpts{ Namespace: milvusNamespace, @@ -1005,6 +1018,7 @@ func RegisterQueryNode(registry *prometheus.Registry) { registry.MustRegister(QueryNodeSQSegmentLatency) registry.MustRegister(QueryNodeSQSegmentLatencyInCore) registry.MustRegister(QueryNodeReduceLatency) + registry.MustRegister(QueryNodeFunctionChainLatency) registry.MustRegister(QueryNodeLoadSegmentLatency) registry.MustRegister(QueryNodeReadTaskReadyLen) registry.MustRegister(QueryNodeReadTaskReadyNQ)