feat: introduce Arrow-based function chain pipeline to replace legacy rerank implementation - #47919
Conversation
|
@junjiejiangjjj This is a feature PR ( How to resolve: Design documents location: https://github.com/milvus-io/milvus-design-docs/tree/main/design_docs |
mergify
Bot
added
dco-passed
|
@junjiejiangjjj Please associate the related issue to the body of your Pull Request. (eg. "issue: #") |
25d5d6b to
89e4f1c
Compare
fb7592c to
57386de
Compare
sre-ci-robot
removed
the
do-not-merge/work-in-progress
|
[ci-v2-notice] To rerun ci-v2 checks, comment with:
If you have any questions or requests, please contact @zhikunyao. |
|
@junjiejiangjjj go-sdk check failed, comment |
sre-ci-robot
added
the
low-code-coverage
7b56659 to
6197c97
Compare
sre-ci-robot
added
the
low-code-coverage
liliu-z
left a comment
liliu-z
left a comment
There was a problem hiding this comment.
Code review: 7 issues identified (2 critical regressions, 5 important improvements)
| // Grouping parameters | ||
| GroupByField string // Field to group by (empty means no grouping) | ||
| GroupSize int64 // Maximum rows per group | ||
| GroupScorer GroupScorer // How to compute group score ("max", "sum", "avg") |
There was a problem hiding this comment.
🔴 strictGroupSize missing from SearchParams — API regression
strictGroupSize is parsed in proxy (search_util.go:33) and stored in rankParams, but rerankOperator never reads it, SearchParams doesn't define it, and GroupByOp doesn't enforce it. Users setting strictGroupSize=true will silently get non-strict results after this change.
There was a problem hiding this comment.
strictGroupSize is a QueryNode-level parameter that controls segment-level search behavior — whether each group must be fully filled to groupSize rows before the search considers it "enough". The proxy only parses it and passes it through to QueryNode via the QueryInfo proto. The old rerank code never used it on the proxy side either. The chain code operates at the proxy rerank/merge layer and does not need to handle this parameter.
| } | ||
| chunks = importChunkedBatch(data, offsets, getValidSlice, array.NewStringBuilder, alloc) | ||
|
|
||
| default: |
There was a problem hiding this comment.
🔴 Timestamptz type not handled — regression
The old rerank/util.go handled DataType_Timestamptz as int64. This type switch covers Bool through String/VarChar/Text but has no Timestamptz case. Collections using Timestamptz fields for decay rerank will hit this default error branch. Fix: add case schemapb.DataType_Timestamptz: handled like Int64.
There was a problem hiding this comment.
Timestamptz is a newly added data type. The scope of this PR is to cover all existing capabilities of the old rerank implementation. Support for new types like Timestamptz will be added in a follow-up.
| } | ||
|
|
||
| if hasNull { | ||
| builder.AppendNull() |
There was a problem hiding this comment.
🟡 NaN silently propagated to user output
Null inputs are correctly handled here (AppendNull), but NaN floats from upstream pass through all combine modes unchecked. Built-in decay expressions have parameter validation that prevents NaN production, but there's no defensive layer if a future expression produces NaN. Consider adding math.IsNaN(result) check after s.combine(values) and returning an error or substituting 0.
There was a problem hiding this comment.
The decay parameter validation in NewDecayExpr already prevents NaN production: decay is constrained to (0.001, 0.999) so log(decay) is always a finite negative, scale > 0 prevents division by zero, and offset >= 0. All three decay functions are mathematically bounded to (0, 1] under these constraints. The only theoretical path to NaN would be a user storing NaN in a float field — the old rerank code had no guard for that either. Acknowledged as a defensive improvement, but not a real risk in practice.
|
|
||
| // buildGroups builds groups from the chunk. | ||
| func (o *GroupByOp) buildGroups(groupChunk arrow.Array, scoreChunk *array.Float32, idChunk arrow.Array, chunkLen int) []*group { | ||
| groupMap := make(map[any]*group) |
There was a problem hiding this comment.
🟡 No group cardinality bound
This map grows to match the number of distinct group keys with no upper limit. While TopK (≤16384) caps it in practice, grouping by a near-unique field silently produces degenerate results (one row per group). A log.Warn when len(groupMap) > limit * 10 would help users catch misconfigured group-by fields early.
There was a problem hiding this comment.
The input size is already bounded by upstream TopK (≤16384), so the map cannot grow unbounded. The group-by field is user-specified — if they pick a near-unique field, the degenerate result is expected behavior, same as the old code. A warning log could be useful for observability but is not a correctness issue. Will consider adding it as a follow-up improvement.
| } | ||
|
|
||
| // Add grouping or sort+limit | ||
| if searchParams.HasGrouping() { |
There was a problem hiding this comment.
🟡 IsSupportGroup check removed — intentional?
The old code rejected group-by for incompatible rerankers via IsSupportGroup() in task_search.go. The new architecture decouples GroupByOp from rerankers, implicitly making all rerankers group-compatible. If intentional, a comment here explaining the design change would help — otherwise this is a missing validation.
There was a problem hiding this comment.
Yes, this is intentional. The new architecture decouples group-by into a standalone GroupByOp that works uniformly across all rerankers, so individual rerankers no longer need to declare group-by compatibility via IsSupportGroup(). This was one of the design goals of the refactor — eliminate per-reranker group-by logic in favor of a single composable operator.
|
|
||
| // collectRRFScores collects RRF scores for a single chunk. | ||
| func (op *MergeOp) collectRRFScores(inputs []*DataFrame, chunkIdx int) (map[any]float32, map[any]idLocation, error) { | ||
| idScores := make(map[any]float32) |
There was a problem hiding this comment.
🟡 map[any] loses compile-time type safety
map[any]float32 relies on runtime type consistency of ID keys. If different chunks are ever parsed into mismatched types (e.g., int64 vs string), lookups silently fail. A type-switched dispatch to map[int64] / map[string] would catch this at compile time.
There was a problem hiding this comment.
In Milvus, the ID (primary key) type is defined at the collection level — either Int64 or VarChar. All inputs to MergeOp come from sub-searches on the same collection, so the ID type is guaranteed to be consistent across all inputs. The type mismatch scenario (int64 vs string keys in the same map) cannot happen in practice.
| return compareTyped(a, i, j) | ||
| case *array.String: | ||
| return compareTyped(a, i, j) | ||
| default: |
There was a problem hiding this comment.
🟡 default: return 0 hides future bugs
If a new type is added to isComparableType() but not to compareArrayValues(), rows with that type silently compare as equal, producing wrong sort order with no error. Since reaching this branch indicates an internal logic inconsistency, this should be a panic or error rather than a silent fallback.
| // Decay Builder | ||
| // ============================================================================= | ||
|
|
||
| func buildDecayChain(fc *FuncChain, collSchema *schemapb.CollectionSchema, funcSchema *schemapb.FunctionSchema, searchMetrics []string) error { |
There was a problem hiding this comment.
buildDecayChain uses fc.Merge which loses sortDescending propagation. For L2 + norm_score=false, results are incorrectly sorted DESC instead of ASC.
There was a problem hiding this comment.
Confirmed and fixed. Root cause is slightly different from "lost sortDescending propagation": decay multiplies $score by a [0, 1] factor and assumes "higher = better", so the chain must always produce DESC scores. Legacy rerank/decay_function.go enforced this by passing toGreater=true to getNormalizeFunc regardless of norm_score; the chain refactor lost that.
Fix: added WithForceDescending option on MergeOp. When set, resolveMergeBehavior applies getDirectionConvertFunc (atan-based) for distance metrics regardless of normalize. buildDecayChain always passes WithForceDescending(true). Other rerankers unchanged.
Coverage: TestExecuteDecay_L2_NoNormScore_RanksByCombinedScore (asserts each score matches (1 − 2·atan(d)/π) × gauss_decay within 1e-5) + Python e2e test_milvus_client_search_decay_rerank_l2_metric_no_norm_score.
| if scorer == "" { | ||
| scorer = GroupScorerMax | ||
| } | ||
| fc.GroupByWithScorer(searchParams.GroupByField, searchParams.GroupSize, searchParams.Limit, searchParams.Offset, scorer) |
There was a problem hiding this comment.
GroupByWithScorer path does not consume sortDescending and hard-codes DESC, breaking weighted + group_by + L2 + normalize=false.
There was a problem hiding this comment.
Confirmed. This is also a long-standing latent bug in legacy rerank/util.go::newGroupingIDScores (no direction parameter, hard-coded DESC) — the chain refactor inherited it. Fix:
- Added
sortDescending boolfield onGroupByOp(defaultstrue, preserving legacy behavior for all existing callers andNewGroupByOpFromRepr) - Added builder-style
SetSortDescending(bool)setter - Direction-aware
sortAndLimitGroup(within-group trim) andprocessChunkstep 4 (cross-group sort) buildRerankChainInternalconstructs the op directly and propagatessortDescendingfrombuildWeightedChain
Max scorer code is unchanged: scores[0] is now correctly the best representative under either direction since sortAndLimitGroup orders the slice first. Public chain API (GroupBy/GroupByWithScorer) unchanged — no breaking change.
While validating this end-to-end I also discovered a related pre-existing proxy bug in reduceAdvanceGroupBy single-shard early return path (search_reduce_util.go:86) — without that fix the chain layer fix alone wasn't observable from Python e2e because the proxy was passing -L2 to the chain. Also fixed in this PR with its own table-driven test (TDD-validated by reverting the fix and watching the new test go red, then re-applying).
Coverage: TestExecuteWeightedGroupBy_L2_NoNormScore_PreservesAscOrder, TestExecuteWeightedGroupBy_L2_NormScore_PreservesDescOrder, TestGroupByOp_AscDirection, TestGroupByOp_DefaultSortDescending + Python e2e test_milvus_client_hybrid_search_weighted_groupby_l2_no_norm_score + new TestReduceAdvanceGroupBy_SingleShardScoreNegation and TestReduceAdvanceGroupBy_SingleShardMatchesMultiShard.
| } | ||
| } | ||
|
|
||
| // Assemble only the rows referenced by the reranked IDs. |
There was a problem hiding this comment.
Skipping nil computers without removing corresponding IDs/FieldsData entries causes row/column misalignment in assembled results.
There was a problem hiding this comment.
You're right that this is an invariant violation. Tracing the proxy code, the precondition (one sub-result with FieldsData while another has none) shouldn't be reachable today — task_search.go sets identical plan.OutputFieldIds across all sub-requests and the needRequery=false path used by hybridSearchPipe always includes the PK. So the defensive continue is dead code that would only fire on an upstream invariant violation, where silently dropping the row is the worst possible outcome.
Fix: continue → return merr.WrapErrServiceInternal(...) with the offending sub-result index, the PK that triggered it, and the collection ID. Surfacing the bug at its source is much better than corrupting Ids ↔ FieldsData row alignment.
Coverage: TestHybridAssembleOp_MixedFieldsDataLayoutErrors (synthesizes the offending state and asserts the operator returns an error mentioning FieldsData).
| return rerank.GetRerankName(funcScore.Functions[0]) | ||
| } | ||
|
|
||
| func validateInputField(collSchema *schemapb.CollectionSchema, fieldName string) error { |
There was a problem hiding this comment.
validateInputField accepts Timestamptz but GetNumericValue only handles Int/Float, causing runtime failures after passing schema validation.
There was a problem hiding this comment.
Investigated: legacy decay listed Timestamptz in its type-dispatch switch (added in c0d62268a) but the support never reached end-to-end. No production path or test ever exercised it — getField had a branch but neither converter nor GetNumericValue handled it, and there was no decay+Timestamptz test. The chain refactor reproduces the same incomplete state.
Since this PR's scope is to preserve legacy functionality rather than extend it, I removed Timestamptz from validateInputField so the chain rejects it explicitly. The user-visible error in proxy e2e actually comes from chain.FromSearchResultData (Arrow converter default branch) which runs before BuildRerankChain — both layers reject Timestamptz, just at different stages. Adding genuine Timestamptz support (converter + GetNumericValue + tests) belongs in a separate feature PR.
Coverage: TestBuildDecayChainTimestamptzInputField (chain unit) + Python e2e test_milvus_client_search_decay_rerank_timestamptz_field_rejected (asserts unsupported field type: Timestamptz).
|
internal/proxy/search_pipeline.go:1 strictGroupSize is silently dropped in the Arrow chain path, causing a functional regression between single-path and multi-path search. Must be restored before merge. |
|
internal/proxy/rerank/converter.go:280 When ids==nil && totalRows>0, a malformed DataFrame is silently constructed instead of returning an error. |
|
internal/proxy/rerank/converter.go:293 Same hole as ids branch: when totalRows>0 && len(scores)==0, the $score column is silently dropped. Return an explicit error. |
|
internal/proxy/rerank/converter.go:1 Legacy rerank/ directory is nearly emptied but not fully removed, and the modelProvider→ModelProvider rename affects other callers. PR cannot be cleanly reverted; finish cleanup or provide a rollback plan. |
|
internal/proxy/rerank/operator_filter.go:1 FilterOp.Execute assumes the FunctionExpr chunked output shares chunk boundaries with the input DataFrame. Arrow does not guarantee this. Add an assertion that input.chunkSizes[chunkIdx]==boolChunk.Len() to fail fast. |
|
internal/proxy/rerank/converter.go:1 Missing unit tests: (1) empty path + offset/limit on chain pipeline, (2) end-to-end hybrid + chain-path strictGroupSize. |
|
internal/proxy/search_pipeline.go:1 Single-shard rerank score negation inverts sort order vs prior versions. Must be loudly flagged in release notes as a breaking behavior change. |
|
Thanks for the additional pass. Quick rundown: 1. 2. 4. Legacy 5. 6. Missing tests — Will add (1) empty input + offset/limit on the chain pipeline. (2) follows from #1, so n/a. 7. Single-shard score-negation behavior change — Agreed. Will add a "Behavior change" entry to the PR description / release notes covering hybrid + group_by + distance metric (L2/HAMMING/JACCARD) + single-shard collections. Let me know if any of #1/#2/#3/#5 are blocking — happy to add the assertions as pure defense in depth if preferred. |
… rerank implementation Adds a new chain package (internal/util/function/chain/) that implements a composable, Apache Arrow-based pipeline for search result post-processing (reranking, scoring, merging, filtering, grouping, etc.). This replaces the legacy rerank utility functions with a DataFrame-oriented approach that provides better performance and extensibility. Key changes: - New chain package with DataFrame abstraction, operator registry, expression engine (decay, score combine, round decimal, rerank model), and operators (merge, sort, filter, limit, select, map, group-by) - Refactored search_pipeline.go to use chain-based reranking instead of the legacy rerank package, adding a hybrid_assemble operator for hybrid search - Removed legacy rerank utilities (decay_function, rrf_function, weighted_function, util.go, rerank_base) and simplified remaining rerank providers to thin wrappers - Added rerank_meta.go in proxy for structured rerank configuration parsing - Comprehensive unit tests and benchmarks for the chain package - Extended Python and Go integration tests for reranker functions Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
|
/ci-rerun-e2e-default |
|
/lgtm |
|
[APPROVALNOTIFIER] This PR is APPROVED This pull-request has been approved by: junjiejiangjjj, liliu-z The full list of commands accepted by this bot can be found here. The pull request process is described here DetailsNeeds approval from an approver in each of these files:
Approvers can indicate their approval by writing |
3 participants
#46565
design doc: https://github.com/milvus-io/milvus-design-docs/blob/main/design_docs/20260310-function-chain.md
Adds a new chain package (internal/util/function/chain/) that implements a composable, Apache Arrow-based pipeline for search result post-processing (reranking, scoring, merging, filtering, grouping, etc.). This replaces the legacy rerank utility functions with a DataFrame-oriented approach that provides better performance and extensibility.
Key changes: