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 <junjie.jiang@zilliz.com>

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")

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)

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)

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,

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"

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)