fix(fts): deterministic top-k tiebreak (score DESC, row_id ASC) - #7846
fix(fts): deterministic top-k tiebreak (score DESC, row_id ASC)#7846sbrunk wants to merge 2 commits into
Conversation
📝 WalkthroughWalkthroughThe inverted-index search paths now preserve deterministic top-k results for equal BM25 scores. Boundary ties are retained during WAND and BM25 collection, deferred row IDs are resolved before final sorting, and results use score-descending then row_id-ascending ordering. ChangesDeterministic top-k ranking
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Query as WAND search
participant Collector as TopKCollector
participant Merge as BM25 merge
participant Results as Final candidates
Query->>Collector: insert scored documents
Collector->>Collector: retain k-th-score boundary ties
Collector->>Merge: materialize heap and deferred ties
Merge->>Results: resolve row_ids and sort deterministically
Results->>Query: return top-k candidates
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
BM25 top-k resolved equal scores by encounter and pruning order, which varies with k and with partition completion order. top-k1 was not a prefix of top-k2, so paginating a tied-score query duplicated and skipped rows. Give the top-k a total order (score DESC, row_id ASC), matching Lucene: - ScoredDoc::cmp breaks score ties by ascending row_id (the final sort). - The per-partition collector and the cross-partition merge evict on the full (score, row_id) key, so ties resolve by row_id regardless of visit order or which partition finishes first. - WAND pruning lowers the threshold one ULP below the k-th score, so the score-only prune tests keep, rather than drop, docs tied at the k-th score, including a slower partition's ties behind the shared floor. This needs no change to the pruning kernels. - On the no-filter path, where row_ids are resolved after the WAND walk, the collector retains the k-th-score tie band and the merge re-selects by row_id. This keeps the result correct when doc_id order diverges from row_id order after a compaction remap. Non-tied queries keep full pruning power. A large exactly-tied band is scanned instead of pruned and, on the no-filter path, buffered in memory for row_id resolution.
6badd5f to
a7f155e
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Note
Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.
🟡 Other comments (1)
rust/lance-index/src/scalar/inverted/index.rs-11628-11662 (1)
11628-11662: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winUse a masked non-flat predicate for this filtered FTS test.
bm25_searchpasses the prefilter’s OR result back intodocs_for_wand(...), and the relevant existing test shows an OR with an allow-list can still return!has_row_ids()before the outer deferredCandidateAddr::Pending(doc_id)resolver fixes candidates. Becausetest_fts_topk_tied_scores_stable_prefix_filteredalso usesOperator::Or, its “forces the real row_id / full-key eviction path” claim is not guaranteed by the current setup. Use a masked non-flat predicate such asOperator::Andwith an all-rows allow-list so this test matchestest_fts_topk_tied_scores_stable_prefix_unsorted_row_idswithout relying on path-coverage assumptions.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-index/src/scalar/inverted/index.rs` around lines 11628 - 11662, Update test_fts_topk_tied_scores_stable_prefix_filtered to use a masked non-flat prefilter, such as Operator::And combined with the all-rows AllowListFilter, instead of the current Operator::Or setup. Match the predicate construction used by test_fts_topk_tied_scores_stable_prefix_unsorted_row_ids, while preserving the existing top3/top5 assertions and stable row_id ordering.
🧹 Nitpick comments (4)
rust/lance-index/src/scalar/inverted/index.rs (3)
11539-11560: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd
#[cfg_attr(coverage, coverage(off))]to this test-utility impl.The other test doubles in this file (
DocsRowIdReadCounter,DocsRowIdCountingReader,DocsRowIdCountingStore) carry it. As per coding guidelines: "disable coverage for test utilities with#[cfg_attr(coverage, coverage(off))]".♻️ Proposed change
+ #[cfg_attr(coverage, coverage(off))] #[async_trait::async_trait] impl PreFilter for AllowListFilter {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-index/src/scalar/inverted/index.rs` around lines 11539 - 11560, Add #[cfg_attr(coverage, coverage(off))] to the AllowListFilter PreFilter implementation, matching the existing test utilities DocsRowIdReadCounter, DocsRowIdCountingReader, and DocsRowIdCountingStore. Do not alter the implementation methods or behavior.Source: Coding guidelines
1232-1232: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueRow-id resolution is IO-bound; consider
self.store.io_parallelism()for the fan-out.
resolve_row_idsreads the docs file's ROW_ID column, and every other IO fan-out in this file (aggregate_corpus_stats,df_for_term, partition loading) is bounded byself.store.io_parallelism(). Using the compute-intensive CPU count here is inconsistent and does not track store IO capacity.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-index/src/scalar/inverted/index.rs` at line 1232, Update the fan-out in resolve_row_ids to use self.store.io_parallelism() instead of get_num_compute_intensive_cpus(). Keep the existing buffer_unordered behavior and ensure the concurrency limit follows the store’s IO capacity, consistent with the other IO-bound paths.
938-962: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winReplace the bare
.unwrap()s on the heap peeks/pops with.expect("reason")(or return an error).Lines 942, 952, 954 panic if the invariant ever breaks. The analogous collector in
wand.rs(TopKCollector::take_worst) returnsError::internalfor exactly this case, so the two merge paths are inconsistent. As per coding guidelines: "Never use.unwrap(),.expect(),panic!(), orassert!()in library code for fallible operations" and "Avoid bare.unwrap(); ... If unavoidable, use.expect(\"reason\")."♻️ Minimal fix using `expect`
- let kth = candidates.peek().unwrap().0.score; + let kth = candidates + .peek() + .expect("heap is full, so it holds at least one candidate") + .0 + .score; match candidate.score.cmp(&kth) { // Below the k-th score: never competitive. std::cmp::Ordering::Less => {} // Tied at the k-th score: a potential winner once row_ids are // resolved, so keep it in the boundary band. std::cmp::Ordering::Equal => overflow.push(candidate), // Strictly better on score: it enters the heap. Whether the // displaced candidate stays a boundary tie depends on the new k-th. std::cmp::Ordering::Greater => { - let Reverse(displaced) = candidates.pop().unwrap(); + let Reverse(displaced) = + candidates.pop().expect("heap is full, so pop yields the k-th"); candidates.push(Reverse(candidate)); - let new_kth = candidates.peek().unwrap().0.score; + let new_kth = candidates + .peek() + .expect("a candidate was just pushed") + .0 + .score;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-index/src/scalar/inverted/index.rs` around lines 938 - 962, Replace the bare unwrap calls in the candidate heap handling around candidates.peek() and candidates.pop() with descriptive expect messages, or propagate an appropriate internal error consistent with wand.rs TopKCollector::take_worst. Preserve the existing heap-selection and overflow behavior while making invariant failures explicit rather than using bare unwraps.Source: Coding guidelines
rust/lance-index/src/scalar/inverted/wand.rs (1)
4413-4487: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGood collector coverage; add a direct unit test for
admit_ties_flooredge inputs.The tie-admission scheme rests on
admit_ties_floor, but its documented edge behavior (0.0, negative, and NaN inputs relative to thethreshold > 0.0guards) is only covered indirectly. A three-linerstestlocks the contract in place. As per coding guidelines: "Every bugfix and feature must have corresponding tests" and "Userstestfor Rust parameterized tests, use readable#[case::{name}(...)]names".💚 Proposed test
#[rstest] #[case::positive(1.0)] #[case::small(f32::MIN_POSITIVE)] fn test_admit_ties_floor_admits_kth_score_ties(#[case] kth: f32) { let floor = admit_ties_floor(kth); assert!(floor < kth, "floor {floor} must sit below the k-th score {kth}"); assert!(kth > floor, "a doc tied at the k-th score must pass `score > threshold`"); } #[rstest] #[case::zero(0.0)] #[case::negative(-1.0)] fn test_admit_ties_floor_keeps_non_positive_read_as_no_threshold(#[case] kth: f32) { assert!(!(admit_ties_floor(kth) > 0.0)); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-index/src/scalar/inverted/wand.rs` around lines 4413 - 4487, Add direct rstest coverage for the admit_ties_floor function, using readable named cases for positive values (including f32::MIN_POSITIVE), zero, and negative input. Assert positive inputs produce a floor below the k-th score, and zero or negative inputs do not produce a positive threshold; place the tests alongside the existing TopKCollector tests.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/lance-index/src/scalar/inverted/index.rs`:
- Around line 1226-1239: Update the deferred row-ID resolution flow around
resolve_row_ids and the subsequent entries.into_iter().zip(row_ids) loop to
validate that every requested doc_id produced a row ID before assigning results.
Detect any short resolution instead of allowing zip to truncate and preserve the
placeholder, then return an appropriate error consistent with the existing
resolve_deferred_candidates guard.
---
Other comments:
In `@rust/lance-index/src/scalar/inverted/index.rs`:
- Around line 11628-11662: Update
test_fts_topk_tied_scores_stable_prefix_filtered to use a masked non-flat
prefilter, such as Operator::And combined with the all-rows AllowListFilter,
instead of the current Operator::Or setup. Match the predicate construction used
by test_fts_topk_tied_scores_stable_prefix_unsorted_row_ids, while preserving
the existing top3/top5 assertions and stable row_id ordering.
---
Nitpick comments:
In `@rust/lance-index/src/scalar/inverted/index.rs`:
- Around line 11539-11560: Add #[cfg_attr(coverage, coverage(off))] to the
AllowListFilter PreFilter implementation, matching the existing test utilities
DocsRowIdReadCounter, DocsRowIdCountingReader, and DocsRowIdCountingStore. Do
not alter the implementation methods or behavior.
- Line 1232: Update the fan-out in resolve_row_ids to use
self.store.io_parallelism() instead of get_num_compute_intensive_cpus(). Keep
the existing buffer_unordered behavior and ensure the concurrency limit follows
the store’s IO capacity, consistent with the other IO-bound paths.
- Around line 938-962: Replace the bare unwrap calls in the candidate heap
handling around candidates.peek() and candidates.pop() with descriptive expect
messages, or propagate an appropriate internal error consistent with wand.rs
TopKCollector::take_worst. Preserve the existing heap-selection and overflow
behavior while making invariant failures explicit rather than using bare
unwraps.
In `@rust/lance-index/src/scalar/inverted/wand.rs`:
- Around line 4413-4487: Add direct rstest coverage for the admit_ties_floor
function, using readable named cases for positive values (including
f32::MIN_POSITIVE), zero, and negative input. Assert positive inputs produce a
floor below the k-th score, and zero or negative inputs do not produce a
positive threshold; place the tests alongside the existing TopKCollector tests.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 07a492b0-2bd6-43ee-9af1-781768a51473
📒 Files selected for processing (3)
rust/lance-index/src/scalar/inverted/builder.rsrust/lance-index/src/scalar/inverted/index.rsrust/lance-index/src/scalar/inverted/wand.rs
- Assert row_id resolution length in the deferred merge (debug_assert_eq!), documenting that resolve_row_ids returns one row_id per requested doc_id so the zip cannot silently truncate to placeholder row 0. - Replace bare unwrap on the merge heap peeks/pops with expect messages, matching the invariant style of TopKCollector::take_worst. - Mark the AllowListFilter test double with #[cfg_attr(coverage, coverage(off))], consistent with the other test doubles in this file. - Add direct rstest coverage for admit_ties_floor edge inputs (positive, zero, negative, NaN).
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rust/lance-index/src/scalar/inverted/wand.rs (1)
160-195: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winNaN-score guard is unreachable while the heap is still filling up.
The NaN check at lines 191-195 only executes once the heap is full (
self.heap.len() == self.limit). The earlier "heap not yet full" fast path (lines 178-184) pushesdocunconditionally, before the NaN check runs. So a NaN-scored candidate is accepted into the top-k while the collector is still filling, but rejected once it reaches capacity — inconsistent with the stated intent ("A NaN custom-scorer score is never competitive; drop it"). This can let an uncompetitive/invalid document occupy a top-k slot (and, in deferred mode, be treated as tied at whatever the eventual k-th score becomes).🐛 Proposed fix: move the NaN check before the not-yet-full fast path
) -> Result<bool> { if self.limit == 0 { return Ok(false); } + // A NaN custom-scorer score is never competitive; drop it before it can enter the + // heap while still filling (the full-heap branch below already excludes it). + if doc.score.0.is_nan() { + return Ok(false); + } if self.heap.len() > self.limit { return Err(Error::internal(format!( "FTS top-k heap length {} exceeds limit {}", self.heap.len(), self.limit ))); } // Heap not yet full: everything is competitive. if self.heap.len() < self.limit { let slot = self.frequency_slots.push(pairs)?; self.heap .push(Reverse((doc, doc_length, posting_doc_id, slot))); return Ok(true); } let Some(worst) = self.heap.peek().map(|entry| entry.0.0.clone()) else { return Err(Error::internal( "FTS top-k heap is empty while its nonzero limit is reached", )); }; - // A NaN custom-scorer score is never competitive; drop it (also keeps NaN out of the - // `total_cmp` key, matching the pre-tiebreak collector). - if doc.score.0.is_nan() { - return Ok(false); - }None of the new tests (
test_top_k_collector_deferred_tiebreak_retains_boundary_ties,test_top_k_collector_full_key_eviction_by_row_id,test_top_k_collector_reuses_frequency_slots) exercise a NaN insert while the heap is belowlimit, so this regressed silently. Worth adding a regression test alongside the fix.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-index/src/scalar/inverted/wand.rs` around lines 160 - 195, Move the NaN-score guard in insert before the heap-not-full fast path so doc.score.0.is_nan() always returns Ok(false), including while the heap is filling. Preserve the existing limit and capacity checks, and add a regression test covering insertion of a NaN-scored document when the heap length is below limit.
♻️ Duplicate comments (1)
rust/lance-index/src/scalar/inverted/index.rs (1)
1234-1257: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
debug_assert_eq!still lets a shortresolve_row_idsreturn silently corrupt results in release builds.This is the same risk flagged in a prior review (
resolve_row_idsreturning fewer row_ids than requested leaves the placeholderrow_id = 0set at line 1230, via theziptruncating silently). The fix replaces the earlier no-check code withdebug_assert_eq!, butdebug_assert!compiles out in release builds, so in production a length mismatch still produces a wrong-row (0) result with no error — exactly the scenario the prior review asked to guard against with an explicitErr.As per coding guidelines: "Do not silently guard against impossible conditions; use
debug_assert!, return an explicit error, or remove the check" and "Preferdebug_assert!overassert!for non-safety invariants; reserveassert!for conditions preventing data corruption" — a silently-wrong row_id returned to the caller is exactly a data-corruption outcome, so this invariant belongs behind an explicitResultcheck (orassert!), notdebug_assert!.🛡️ Proposed fix
for (entries, row_ids) in batches { - // `resolve_row_ids` maps one row_id per requested doc_id, so the lengths always match; - // the assert documents that invariant (a short return would otherwise leave placeholder - // row 0 via the `zip` below). - debug_assert_eq!(entries.len(), row_ids.len()); + if entries.len() != row_ids.len() { + return Err(Error::internal(format!( + "resolve_row_ids returned {} row_ids for {} deferred doc_ids", + row_ids.len(), + entries.len() + ))); + } for ((pos, _), row_id) in entries.into_iter().zip(row_ids) { resolved[pos].1 = row_id; } }As per coding guidelines: "Do not silently guard against impossible conditions; use
debug_assert!, return an explicit error, or remove the check." Based on a past review comment on this exact code (resolve_row_idstruncation viazip), which proposed the same explicit-error fix that was not fully applied here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-index/src/scalar/inverted/index.rs` around lines 1234 - 1257, Replace the `debug_assert_eq!` in the deferred row-resolution loop with an explicit error check that returns an error when `entries.len()` and `row_ids.len()` differ, before the `zip` iteration. Preserve normal resolution for matching lengths and ensure release builds cannot leave placeholder row IDs after `resolve_row_ids` truncates the mismatch.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@rust/lance-index/src/scalar/inverted/wand.rs`:
- Around line 160-195: Move the NaN-score guard in insert before the
heap-not-full fast path so doc.score.0.is_nan() always returns Ok(false),
including while the heap is filling. Preserve the existing limit and capacity
checks, and add a regression test covering insertion of a NaN-scored document
when the heap length is below limit.
---
Duplicate comments:
In `@rust/lance-index/src/scalar/inverted/index.rs`:
- Around line 1234-1257: Replace the `debug_assert_eq!` in the deferred
row-resolution loop with an explicit error check that returns an error when
`entries.len()` and `row_ids.len()` differ, before the `zip` iteration. Preserve
normal resolution for matching lengths and ensure release builds cannot leave
placeholder row IDs after `resolve_row_ids` truncates the mismatch.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 7f896b81-bab5-4b14-be10-8e898b1ce3d4
📒 Files selected for processing (2)
rust/lance-index/src/scalar/inverted/index.rsrust/lance-index/src/scalar/inverted/wand.rs
Problem
Lance's BM25 top-k resolved equal scores by encounter and pruning order, which varies with
kand with the order partitions finish. Sotop-k1was not an ordered prefix oftop-k2, and paginating a broad tied-score term (many docs sharing tf and doc length, so identical BM25) duplicated and skipped rows acrossfrom/sizepages.Lucene avoids this with a deterministic secondary sort
(score DESC, docId ASC). This PR brings the same guarantee to Lance, keyed onrow_id.Change
Give the FTS top-k a total order
(score DESC, row_id ASC):ScoredDoc::cmpbreaks score ties by ascendingrow_id(the final sort).(score, row_id)key, and the WAND threshold is lowered one ULP below the k-th score (admit_ties_floor) so the score-only prune tests keep, not drop, docs tied at the k-th score. No change to the pruning kernels.(score DESC, row_id ASC).Row id stability
Correctness within a dataset version does not depend on stable row ids. The tiebreak only needs
row_idto be unique and comparable, which holds for both legacy row addresses and stable row ids, so the dup/skip fix works in either mode. Stable row ids additionally keep the order stable across compaction (legacy addresses change on compaction, so the key moves). This mirrors Lucene, whosedocIdtiebreak is deterministic within a searcher version but not stable across segment merges.Interaction with deferred row_id resolution (#7897)
#7897 resolves row_ids only for the top-k survivors, so a broad query loads only those partitions' ROW_ID columns. Deterministic tie resolution needs the row_ids of the whole boundary tie band, so on the no-filter path this loads the ROW_ID column of every partition holding a doc at the k-th score:
This is the inherent cost of an exact row_id tiebreak: you must see the tied row_ids to pick the lowest k. It only applies to tied no-filter queries, which are exactly the queries that need stable pagination.
Tests
top-k1is an ordered prefix oftop-k2) for single-partition (sorted and post-compaction unsorted), multi-partition, and filtered (non-deferred) paths.Performance
The change is a no-op on the non-tied path by construction. A stabilized local benchmark on the full 1M corpus (both binaries pre-built, interleaved before/after with alternating round order, warm-up,
caffeinate, median over 6 rounds) found no meaningful regression:invert_searchwas neutral (median favored the change; min-to-min within ~1%) andinvert_phrase_searchwas +1.8%, both inside the run-to-run spread. This matches a 250k run (+0.4% / +0.9%).@bench-bot benchmarkon the CI ARM64 runner remains the authoritative check, since it touches the WAND pruning inequality.