From a7f155e5817a32456fcb9ac3007578ee9879674b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=B6ren=20Brunk?= Date: Sun, 19 Jul 2026 15:49:49 +0200 Subject: [PATCH 1/2] fix(fts): deterministic top-k tiebreak (score DESC, row_id ASC) 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. --- .../src/scalar/inverted/builder.rs | 9 +- rust/lance-index/src/scalar/inverted/index.rs | 411 ++++++++++++++++-- rust/lance-index/src/scalar/inverted/wand.rs | 303 +++++++++++-- 3 files changed, 657 insertions(+), 66 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 588c8822942..2bcaf2afaf8 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -1756,8 +1756,15 @@ impl PartialOrd for ScoredDoc { } impl Ord for ScoredDoc { + // Total order for the FTS top-k: score first, ties broken by ascending row_id (a lower row_id ranks + // higher, so it compares as "greater" here). A total tiebreak makes top-k a stable prefix across `k` + // (top-k1 is an ordered prefix of top-k2), which tied-score pagination needs. This orientation makes a + // min-heap's `peek` the worst candidate (lowest score, highest row_id) to evict first, and + // `into_sorted_vec` over `Reverse` yield (score DESC, row_id ASC). fn cmp(&self, other: &Self) -> std::cmp::Ordering { - self.score.cmp(&other.score) + self.score + .cmp(&other.score) + .then_with(|| other.row_id.cmp(&self.row_id)) } } diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index a6db0e355c9..2742105f486 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -914,8 +914,17 @@ impl InvertedIndex { } } + // Global top-k merge that stays deterministic under the deferred-row_id + // scheme: row_ids are resolved only after selection, so the merge cannot + // tiebreak by row_id here. Keep the top-k by score, plus every candidate + // tied at the current k-th score (the boundary band) in `overflow`. After + // the merge, row_ids are resolved for the heap and the band and the exact + // top-k by (score DESC, row_id ASC) is chosen. The band is cleared when the + // k-th score rises, so it holds at most the current tie band (empty for + // non-tied queries). fn push_scored_candidate( candidates: &mut BinaryHeap>, + overflow: &mut Vec, limit: usize, slot: u32, addr: CandidateAddr, @@ -928,15 +937,37 @@ impl InvertedIndex { }; if candidates.len() < limit { candidates.push(Reverse(candidate)); - } else if candidates.peek().unwrap().0.score.0 < score { - candidates.pop(); - candidates.push(Reverse(candidate)); + return; + } + let kth = candidates.peek().unwrap().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(); + candidates.push(Reverse(candidate)); + let new_kth = candidates.peek().unwrap().0.score; + if new_kth > displaced.score { + overflow.clear(); + } else { + overflow.push(displaced); + } + } } } let mask = prefilter.mask(); let mut candidates = BinaryHeap::new(); + // Boundary ties (tied at the current k-th score) held for the + // deterministic tiebreak; resolved and re-selected with the heap after the + // merge (see `push_scored_candidate`). + let mut overflow: Vec = Vec::new(); // Shared top-k floor across this query's partitions. Seeded to -inf so // the first real score wins; each partition publishes its local k-th // and prunes against the running global k-th (a lower bound on the true @@ -1122,7 +1153,7 @@ impl InvertedIndex { score += idf_by_position[term_index as usize] * scorer.doc_weight(freq, doc_length); } - push_scored_candidate(&mut candidates, limit, slot, addr, score); + push_scored_candidate(&mut candidates, &mut overflow, limit, slot, addr, score); } } else { let grouped_positions = grouped_expansions @@ -1153,33 +1184,37 @@ impl InvertedIndex { score += term.query_weight() * scorer.doc_weight(freq, doc_length); } } - push_scored_candidate(&mut candidates, limit, slot, addr, score); + push_scored_candidate(&mut candidates, &mut overflow, limit, slot, addr, score); } } } - // Resolve row_ids only for the candidates that survived the global - // top-k: group deferred survivors per partition and batch-resolve — - // at most `limit` lookups regardless of how many partitions - // contributed candidates. - /// One partition's surviving deferred candidates: positions in the - /// merged result list paired with the doc_ids to resolve. + // Merge the top-k heap with the boundary-tie band, resolve row_ids for all + // of them (at most `limit` + the tie band, and just `limit` when there are + // no ties), then choose and order the exact top-k by + // (score DESC, row_id ASC). Resolution stays deferred and batched per + // partition, so a non-tied query still resolves only ~`limit` row_ids. + /// One partition's deferred candidates: positions in the merged result + /// list paired with the doc_ids to resolve. type DeferredGroup = Vec<(usize, u32)>; - let sorted = candidates.into_sorted_vec(); - let mut row_ids = Vec::with_capacity(sorted.len()); - let mut scores = Vec::with_capacity(sorted.len()); + let mut merged: Vec = candidates + .into_vec() + .into_iter() + .map(|Reverse(candidate)| candidate) + .collect(); + merged.append(&mut overflow); + + let mut resolved: Vec<(f32, u64)> = Vec::with_capacity(merged.len()); let mut deferred: HashMap = HashMap::new(); - for (pos, Reverse(candidate)) in sorted.into_iter().enumerate() { - scores.push(candidate.score.0); + for (pos, candidate) in merged.into_iter().enumerate() { + let score = candidate.score.0; + let slot = candidate.slot; match candidate.addr { - CandidateAddr::RowId(row_id) => row_ids.push(row_id), + CandidateAddr::RowId(row_id) => resolved.push((score, row_id)), CandidateAddr::Pending(doc_id) => { - deferred - .entry(candidate.slot) - .or_default() - .push((pos, doc_id)); + deferred.entry(slot).or_default().push((pos, doc_id)); // Placeholder, overwritten by the batch resolution below. - row_ids.push(0); + resolved.push((score, 0)); } } } @@ -1188,21 +1223,27 @@ impl InvertedIndex { .into_iter() .map(|(slot, entries)| (resolving_parts[slot as usize].clone(), entries)) .collect::>(); - let resolved: Vec<(DeferredGroup, Vec)> = + let batches: Vec<(DeferredGroup, Vec)> = stream::iter(groups.into_iter().map(|(part, entries)| async move { let doc_ids: Vec = entries.iter().map(|&(_, doc_id)| doc_id).collect(); - let resolved = part.docs.resolve_row_ids(&doc_ids).await?; - Result::Ok((entries, resolved)) + let row_ids = part.docs.resolve_row_ids(&doc_ids).await?; + Result::Ok((entries, row_ids)) })) .buffer_unordered(get_num_compute_intensive_cpus()) .try_collect() .await?; - for (entries, resolved) in resolved { - for ((pos, _), row_id) in entries.into_iter().zip(resolved) { - row_ids[pos] = row_id; + for (entries, row_ids) in batches { + for ((pos, _), row_id) in entries.into_iter().zip(row_ids) { + resolved[pos].1 = row_id; } } } + + // Deterministic total order (score DESC, row_id ASC), then the exact top-k. + resolved.sort_unstable_by(|a, b| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1))); + resolved.truncate(limit); + let row_ids = resolved.iter().map(|&(_, row_id)| row_id).collect(); + let scores = resolved.iter().map(|&(score, _)| score).collect(); Ok((row_ids, scores)) } @@ -8981,7 +9022,12 @@ mod tests { for d in 0..ndocs { builder.posting_lists[0].add(d as u32, PositionRecorder::Count(1)); let row_id = pid * 1000 + d; - builder.docs.append(row_id, 1); + // Distinct doc lengths give distinct BM25 scores, so the top-k is a + // well-defined set spanning <= k partitions. That keeps the deferred + // "resolve only survivors" optimization (#7897) meaningful to test. + // The tied case (which must resolve the whole cross-partition tie + // band) is covered by test_bm25_search_tied_scores_* below. + builder.docs.append(row_id, (pid * 3 + d + 1) as u32); expected_row_ids.push(row_id); } builder.write(store).await.unwrap(); @@ -9179,6 +9225,76 @@ mod tests { ); } + #[tokio::test] + async fn test_bm25_search_tied_scores_deterministic_across_many_partitions() { + // A fully-tied query (every doc: same term, tf=1, length 1, so identical BM25) + // spread across many partitions. The deterministic (score, row_id) tiebreak + // must return the globally lowest row_ids, which forces resolving the whole + // cross-partition tie band: every matching partition's ROW_ID column loads. + // This is the documented cost of determinism vs #7897's resolve-only-survivors + // (which applies to non-tied queries, covered by the test above). + const PARTS: u64 = 12; + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + for pid in 0..PARTS { + let mut b = InnerBuilder::new(pid, false, TokenSetFormat::default()); + b.tokens.add("pipeline".to_owned()); + b.posting_lists.push(PostingListBuilder::new(false)); + for d in 0..3u64 { + b.posting_lists[0].add(d as u32, PositionRecorder::Count(1)); + // Identical length across every doc and partition => tied scores. + b.docs.append(pid * 1000 + d, 1); + } + b.write(store.as_ref()).await.unwrap(); + } + write_test_metadata(&store, (0..PARTS).collect(), InvertedIndexParams::default()).await; + + let counter = Arc::new(DocsRowIdReadCounter::default()); + let inner_store: Arc = store.clone(); + let counting_store: Arc = Arc::new(DocsRowIdCountingStore { + inner: inner_store, + counter: counter.clone(), + }); + let cache = Arc::new(LanceCache::with_capacity(64 * 1024 * 1024)); + let index = InvertedIndex::load(counting_store, None, cache.as_ref()) + .await + .unwrap(); + + async fn top_row_ids(index: &InvertedIndex, k: usize) -> Vec { + let (row_ids, _) = index + .bm25_search( + Arc::new(Tokens::new(vec!["pipeline".to_owned()], DocType::Text)), + Arc::new(FtsSearchParams::new().with_limit(Some(k))), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + row_ids + } + + // Each partition owns row_ids pid*1000 + {0,1,2}; deterministic top-k is the + // k lowest row_ids in order (partition 0 owns the lowest: 0, 1, 2). + let top3 = top_row_ids(&index, 3).await; + assert_eq!(top3, vec![0, 1, 2]); + // Resolving the tie band read every matching partition's ROW_ID column. + assert_eq!( + counter.full_column_reads(), + PARTS as usize, + "deterministic tie resolution reads all matching partitions' ROW_ID columns" + ); + + let top6 = top_row_ids(&index, 6).await; + assert_eq!(top6, vec![0, 1, 2, 1000, 1001, 1002]); + assert_eq!(top3, top6[..3], "stable prefix across k under ties"); + } + #[tokio::test] async fn test_prewarm_fills_row_ids_cache_entry_without_a_second_copy() { let tmpdir = TempObjDir::default(); @@ -11350,6 +11466,243 @@ mod tests { ); } + // Deterministic top-k tiebreak (score DESC, row_id ASC). + // A BM25 top-k that breaks ties arbitrarily is not a stable prefix across `k`: top-k1 is not a prefix + // of top-k2, so paginating a tied-score query (a broad term where many docs share `(tf, doc_len)`) + // duplicates and skips rows across pages. Four coupled changes fix this, all keyed on row_id: + // 1. `ScoredDoc::cmp` is the total order (score DESC, row_id ASC) used for the final sort. + // 2. The per-partition collector (`wand::TopKCollector::insert`) and the cross-partition merge + // (`push_scored_candidate`) both evict on the full (score, row_id) key, so a score-tie with a + // lower row_id wins no matter the order docs or partitions are visited (`buffer_unordered`). + // 3. The WAND threshold is lowered one ULP below the k-th score (`wand::admit_ties_floor`), so the + // score-only prune tests keep, not drop, docs tied at the k-th score, including a slower + // partition's boundary ties gated by the shared cross-partition floor. + // 4. On the deferred-row_id fast path (no filter) wand tiebreaks by a local doc_id proxy and + // resolves real row_ids only afterward. That proxy diverges from row_id order after a compaction + // remap, so the collector keeps the boundary tie band (`TopKCollector::tie_overflow`) and the + // merge re-selects by row_id once resolved. + // This covers a single partition (even post-compaction) and a multi-partition index. Non-tied queries + // keep full pruning power. A query whose k-th-score tie band is large (a low-diversity field where many + // docs share tf and doc length) scans that band instead of pruning it, and on the deferred path buffers + // it in memory for row_id resolution: the inherent cost of an exact row_id tiebreak there. + async fn build_tied_partition( + store: &Arc, + partition_id: u64, + row_ids: &[u64], + ) { + let mut builder = InnerBuilder::new(partition_id, false, TokenSetFormat::default()); + builder.tokens.add("alpha".to_owned()); + builder.posting_lists.push(PostingListBuilder::new(false)); + for (local_doc_id, &row_id) in row_ids.iter().enumerate() { + builder.posting_lists[0].add(local_doc_id as u32, PositionRecorder::Count(1)); + // Every doc: term "alpha" once, doc length 1, so identical BM25 and all scores tie. + builder.docs.append(row_id, 1); + } + builder.write(store.as_ref()).await.unwrap(); + } + + async fn search_alpha(index: &InvertedIndex, limit: usize) -> Vec { + search_alpha_with(index, limit, Arc::new(NoFilter)).await + } + + async fn search_alpha_with( + index: &InvertedIndex, + limit: usize, + prefilter: Arc, + ) -> Vec { + let tokens = Arc::new(Tokens::new(vec!["alpha".to_owned()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(limit))); + let (row_ids, scores) = index + .bm25_search( + tokens, + params, + Operator::Or, + prefilter, + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + // All docs tie, so every returned score must be identical (guards the "tied" premise). + for w in scores.windows(2) { + assert!( + (w[0] - w[1]).abs() < 1e-6, + "premise: scores must be tied, got {scores:?}" + ); + } + row_ids + } + + // A prefilter with a non-select-all allow-list. Any explicit allow-list makes `mask().is_select_all()` + // false, which routes the search off the deferred-row_id fast path and onto the collector that holds + // real row_ids: the path where full-key eviction (`doc <= worst`) decides the tie band directly. + struct AllowListFilter { + mask: Arc, + } + + #[async_trait::async_trait] + impl PreFilter for AllowListFilter { + async fn wait_for_ready(&self) -> Result<()> { + Ok(()) + } + fn is_empty(&self) -> bool { + false + } + fn mask(&self) -> Arc { + self.mask.clone() + } + fn filter_row_ids<'a>(&self, row_ids: Box + 'a>) -> Vec { + row_ids + .enumerate() + .filter_map(|(i, r)| self.mask.selected(*r).then_some(i as u64)) + .collect() + } + } + + #[tokio::test] + async fn test_fts_topk_tied_scores_stable_prefix_single_partition() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + // 8 docs, row_ids monotonic with local doc id, all tied. + build_tied_partition(&store, 0, &[100, 101, 102, 103, 104, 105, 106, 107]).await; + write_test_metadata(&store, vec![0], InvertedIndexParams::default()).await; + let cache = Arc::new(LanceCache::with_capacity(4096)); + let index = InvertedIndex::load(store.clone(), None, cache.as_ref()) + .await + .unwrap(); + + let top3 = search_alpha(&index, 3).await; + let top5 = search_alpha(&index, 5).await; + let top8 = search_alpha(&index, 8).await; + // Deterministic (score DESC, row_id ASC): the k lowest row_ids, ascending. + assert_eq!( + top3, + vec![100, 101, 102], + "top-3 must be the lowest 3 row_ids in order" + ); + assert_eq!(top5, vec![100, 101, 102, 103, 104]); + assert_eq!(top8, vec![100, 101, 102, 103, 104, 105, 106, 107]); + // Stable prefix: smaller-k is an exact ordered prefix of larger-k, so no cross-page dup/skip. + assert_eq!(top3, top5[..3], "top-3 must be a prefix of top-5"); + assert_eq!(top5, top8[..5], "top-5 must be a prefix of top-8"); + } + + // Single partition, row_ids NOT monotonic with doc_id (simulates a post-compaction stable-row-id remap: + // postings are still visited in doc_id order, but that order no longer ascends by row_id). The tie band + // must be decided by row_id, not encounter order: a strict-greater-score collector would keep whichever + // tied docs it saw first (the lowest doc_ids), not the lowest row_ids. + #[tokio::test] + async fn test_fts_topk_tied_scores_stable_prefix_unsorted_row_ids() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + // doc_id order: 105,100,107,102,104,101,106,103 ; correct top-k is still the lowest row_ids. + build_tied_partition(&store, 0, &[105, 100, 107, 102, 104, 101, 106, 103]).await; + write_test_metadata(&store, vec![0], InvertedIndexParams::default()).await; + let cache = Arc::new(LanceCache::with_capacity(4096)); + let index = InvertedIndex::load(store.clone(), None, cache.as_ref()) + .await + .unwrap(); + + let top3 = search_alpha(&index, 3).await; + let top5 = search_alpha(&index, 5).await; + assert_eq!( + top3, + vec![100, 101, 102], + "tie band must resolve by row_id, not doc_id/encounter order" + ); + assert_eq!(top5, vec![100, 101, 102, 103, 104]); + assert_eq!(top3, top5[..3], "top-3 must be a prefix of top-5"); + } + + // Same unsorted-row_id scenario, but with a prefilter so the search runs off the deferred fast path: + // the collector holds real row_ids and full-key eviction (`doc <= worst`), not the tie-overflow, must + // pick the tie band. The allow-list covers every doc, so results still match the unfiltered case. + #[tokio::test] + async fn test_fts_topk_tied_scores_stable_prefix_filtered() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + let row_ids = [105u64, 100, 107, 102, 104, 101, 106, 103]; + build_tied_partition(&store, 0, &row_ids).await; + write_test_metadata(&store, vec![0], InvertedIndexParams::default()).await; + let cache = Arc::new(LanceCache::with_capacity(4096)); + let index = InvertedIndex::load(store.clone(), None, cache.as_ref()) + .await + .unwrap(); + + // Allow every row: a non-select-all mask still forces the real-row_id (non-deferred) path. + let filter = || -> Arc { + Arc::new(AllowListFilter { + mask: Arc::new(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter( + row_ids.iter().copied(), + ))), + }) + }; + + let top3 = search_alpha_with(&index, 3, filter()).await; + let top5 = search_alpha_with(&index, 5, filter()).await; + assert_eq!( + top3, + vec![100, 101, 102], + "filtered (non-deferred) path must resolve ties by row_id via full-key eviction" + ); + assert_eq!(top5, vec![100, 101, 102, 103, 104]); + assert_eq!(top3, top5[..3], "top-3 must be a prefix of top-5"); + } + + // Multi-partition tied-score determinism. Tied docs are split across two partitions with interleaved + // row_ids, so the shared cross-partition WAND threshold and the merge must resolve ties by row_id + // independent of the order partitions complete in (`buffer_unordered`). Before the tiebreak-aware + // threshold (`admit_ties_floor`), whichever partition finished first published its k-th score as a + // global floor and its siblings pruned their equal-score docs, making the selection load-dependent. + #[tokio::test] + async fn test_fts_topk_tied_scores_stable_prefix_across_partitions() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + // Tied docs split across two partitions (interleaved row_ids), so the cross-partition merge must + // still resolve ties by row_id, independent of partition completion order (buffer_unordered). + build_tied_partition(&store, 0, &[100, 102, 104, 106]).await; + build_tied_partition(&store, 1, &[101, 103, 105, 107]).await; + write_test_metadata(&store, vec![0, 1], InvertedIndexParams::default()).await; + let cache = Arc::new(LanceCache::with_capacity(4096)); + let index = InvertedIndex::load(store.clone(), None, cache.as_ref()) + .await + .unwrap(); + + let top3 = search_alpha(&index, 3).await; + let top6 = search_alpha(&index, 6).await; + assert_eq!( + top3, + vec![100, 101, 102], + "cross-partition ties must resolve by row_id" + ); + assert_eq!( + top6, + vec![100, 101, 102, 103, 104, 105], + "cross-partition top-6 must be the lowest 6 row_ids in order" + ); + assert_eq!( + top3, + top6[..3], + "cross-partition top-3 must be a prefix of top-6" + ); + } + #[tokio::test] async fn test_and_query_accepts_same_position_alternatives() { let tmpdir = TempObjDir::default(); diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index f730d6f9d4b..7378e3a15bd 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -95,23 +95,66 @@ impl FrequencySlots { } } +/// A boundary tie kept outside the heap in deferred-row_id mode (see [`TopKCollector::defer_tiebreak`]): +/// a doc tied at the k-th score that lost the doc_id proxy tiebreak. `into_candidates` returns it +/// alongside the heap so the downstream merge can re-select by real row_id. Owns its freqs because it +/// never took a frequency slot. +struct TieCandidate { + doc: ScoredDoc, + doc_length: u32, + posting_doc_id: u64, + freqs: Vec<(u32, u32)>, +} + /// Owns the top-k heap and the frequency slots referenced by its entries. struct TopKCollector { limit: usize, heap: TopKHeap, frequency_slots: FrequencySlots, + /// Set when the walk scores against local doc_ids and resolves row_ids only post-wand. The heap + /// then tiebreaks by doc_id, which differs from real row_id order after a compaction remap, so + /// boundary ties go to `tie_overflow` for a real-row_id re-selection downstream. False when the heap + /// already holds real row_ids, in which case the overflow is unused. + defer_tiebreak: bool, + /// Boundary ties (score == current k-th) the heap cannot hold, populated only in deferred mode. + /// Cleared when the k-th score rises, so it holds the current k-th-score tie band. Usually tiny, but a + /// low-diversity field (many docs sharing tf and doc length) can make that band a large fraction of the + /// matches, so this grows with the band: the inherent cost of resolving exact row_ids for ties there. + tie_overflow: Vec, } impl TopKCollector { - fn new(limit: usize, initial_capacity: usize) -> Self { + fn new(limit: usize, initial_capacity: usize, defer_tiebreak: bool) -> Self { let initial_capacity = initial_capacity.min(limit); Self { limit, heap: BinaryHeap::with_capacity(initial_capacity), frequency_slots: FrequencySlots::with_capacity(initial_capacity), + defer_tiebreak, + tie_overflow: Vec::new(), } } + /// Pop the worst heap entry as an owned `TieCandidate` (freqs moved out of its slot), returning the + /// freed slot index so the caller can reuse it for the replacement. + fn take_worst(&mut self) -> Result<(TieCandidate, u32)> { + let Some(Reverse((doc, doc_length, posting_doc_id, slot))) = self.heap.pop() else { + return Err(Error::internal( + "FTS top-k heap entry disappeared during replacement", + )); + }; + let freqs = self.frequency_slots.take(slot)?; + Ok(( + TieCandidate { + doc, + doc_length, + posting_doc_id, + freqs, + }, + slot, + )) + } + /// Insert a competitive result. When the heap is full, its evicted entry's /// frequency slot is cleared and reused before the replacement is pushed. fn insert( @@ -132,31 +175,90 @@ impl TopKCollector { ))); } - let frequency_slot = if self.heap.len() == self.limit { - let Some(kth_score) = self.heap.peek().map(|entry| entry.0.0.score.0) else { - return Err(Error::internal( - "FTS top-k heap is empty while its nonzero limit is reached", - )); - }; - // Preserve the existing collector semantics for non-finite custom - // scorer output: only a strictly greater raw f32 replaces k-th. - if doc.score.0.partial_cmp(&kth_score) != Some(std::cmp::Ordering::Greater) { + // 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); + } + + if !self.defer_tiebreak { + // Heap holds real row_ids, so the full (score DESC, row_id ASC) key is exact: evict the worst + // only if the candidate beats it on the whole key. A score-tie with a lower row_id wins, a + // higher one loses. Together with the threshold sitting one ULP below the k-th score (see + // `update_threshold`, which lets every k-th-score tie reach here) this makes top-k a + // deterministic prefix. + if doc <= worst { return Ok(false); } - let Some(Reverse((_, _, _, frequency_slot))) = self.heap.pop() else { - return Err(Error::internal( - "FTS top-k heap entry disappeared during replacement", - )); - }; - self.frequency_slots.replace(frequency_slot, pairs)?; - frequency_slot - } else { - self.frequency_slots.push(pairs)? - }; + let (_, slot) = self.take_worst()?; + self.frequency_slots.replace(slot, pairs)?; + self.heap + .push(Reverse((doc, doc_length, posting_doc_id, slot))); + return Ok(true); + } - self.heap - .push(Reverse((doc, doc_length, posting_doc_id, frequency_slot))); - Ok(true) + // Deferred-row_id mode: the heap tiebreaks by the doc_id proxy, so any doc tied at the k-th score + // could still win once real row_ids are resolved. Keep the heap on the proxy but stash every + // boundary tie it cannot hold, so the post-wand re-selection sees them all. + match doc.score.cmp(&worst.score) { + // Below the k-th score: not competitive, not a boundary tie. + std::cmp::Ordering::Less => Ok(false), + std::cmp::Ordering::Equal => { + if doc > worst { + // Wins the proxy tiebreak: takes the heap slot; the displaced worst is a boundary tie + // (its real row_id may still win), so it goes to the overflow. + let (displaced, slot) = self.take_worst()?; + self.frequency_slots.replace(slot, pairs)?; + self.heap + .push(Reverse((doc, doc_length, posting_doc_id, slot))); + self.tie_overflow.push(displaced); + Ok(true) + } else { + // Ties on score but loses the proxy tiebreak: keep it for row_id re-selection. + self.tie_overflow.push(TieCandidate { + doc, + doc_length, + posting_doc_id, + freqs: pairs.collect(), + }); + Ok(false) + } + } + std::cmp::Ordering::Greater => { + // Better on score: it enters the heap. The displaced worst stays a boundary tie only if + // the k-th score did not rise. + let (displaced, slot) = self.take_worst()?; + self.frequency_slots.replace(slot, pairs)?; + self.heap + .push(Reverse((doc, doc_length, posting_doc_id, slot))); + let new_kth = self + .heap + .peek() + .map(|entry| entry.0.0.score) + .ok_or_else(|| Error::internal("FTS top-k heap is empty right after a push"))?; + if new_kth > displaced.doc.score { + // k-th score rose past the old ties: none of them can win now. + self.tie_overflow.clear(); + } else { + // Heap still holds docs at the old k-th score, so the displaced doc stays a tie. + self.tie_overflow.push(displaced); + } + Ok(true) + } + } } fn kth_score_if_full(&self) -> Option { @@ -174,9 +276,11 @@ impl TopKCollector { let Self { heap, mut frequency_slots, + tie_overflow, .. } = self; - heap.into_iter() + let mut candidates = heap + .into_iter() .map( |Reverse((doc, doc_length, posting_doc_id, frequency_slot))| { Ok(DocCandidate { @@ -187,7 +291,20 @@ impl TopKCollector { }) }, ) - .collect() + .collect::>>()?; + // Append the deferred boundary ties. Their real row_id may beat a heap entry; the downstream + // (score, row_id) merge keeps the winners and drops the rest, so returning more than k is safe. + // Empty unless the query had ties at the k-th score. + candidates.reserve(tie_overflow.len()); + for tie in tie_overflow { + candidates.push(DocCandidate { + addr: to_addr(tie.doc.row_id), + posting_doc_id: tie.posting_doc_id, + freqs: tie.freqs, + doc_length: tie.doc_length, + }); + } + Ok(candidates) } #[cfg(test)] @@ -1693,6 +1810,18 @@ fn atomic_store_max_f32(slot: &AtomicU32, val: f32) { } } +/// Lower a competitive-score threshold by one ULP so the score-only WAND prune tests +/// (`score <= threshold`) and their admit mirrors (`score > threshold`) keep, rather than drop, docs +/// whose score equals the k-th best. This is the admit-ties half of the deterministic tiebreak: the +/// kernels stay score-only and fast, and the equal-score band survives to be ordered by row_id in the +/// collector and merge. Docs strictly below the k-th score are still pruned, so non-tied queries keep +/// their pruning power: a block max rarely bit-equals the running k-th score. For a zero or negative +/// input `next_down` drops below zero, which the `threshold > 0.0` guards read as "no threshold yet". +#[inline] +fn admit_ties_floor(threshold: f32) -> f32 { + threshold.next_down() +} + // we were using row id as doc id in the past, which is u64, // but now we are using the index as doc id, which is u32. // so here WAND is a generic struct that can be used for both u32 and u64 doc ids. @@ -1776,8 +1905,12 @@ impl<'a, S: Scorer> Wand<'a, S> { Some((norms, cache)) } - /// Set the pruning threshold from this partition's k-th best, raised to the - /// shared cross-partition floor when one is attached. + /// Set the pruning threshold from this partition's k-th best, raised to the shared cross-partition + /// floor when one is attached. The stored value is one ULP below the k-th score (`admit_ties_floor`) + /// so k-th-score ties are admitted, not pruned, and resolved by row_id downstream. The raw k-th + /// score is what gets published to the shared floor, so every partition lowers from the same basis. + /// The tiebreak determinism holds for `wand_factor <= 1.0` (the default); a larger factor lifts the + /// floor above the k-th score and may prune ties. fn update_threshold(&mut self, local_kth: f32, wand_factor: f32) { let mut t = local_kth * wand_factor; if let Some(shared) = self.shared_threshold.as_ref() { @@ -1787,14 +1920,16 @@ impl<'a, S: Scorer> Wand<'a, S> { t = g; } } - self.threshold = t; + self.threshold = admit_ties_floor(t); } - /// Raise the local threshold to the shared cross-partition floor, picking up - /// updates published by sibling partitions. + /// Raise the local threshold to the shared cross-partition floor published by sibling partitions. + /// Lowered one ULP like `update_threshold` so a tie at the shared k-th score is not pruned by + /// whichever partition published first; the merge resolves it by row_id, making the result + /// independent of partition completion order. fn raise_to_shared_floor(&mut self, wand_factor: f32) { if let Some(shared) = self.shared_threshold.as_ref() { - let g = f32::from_bits(shared.load(Ordering::Relaxed)) * wand_factor; + let g = admit_ties_floor(f32::from_bits(shared.load(Ordering::Relaxed)) * wand_factor); if g > self.threshold { self.threshold = g; } @@ -1864,7 +1999,13 @@ impl<'a, S: Scorer> Wand<'a, S> { // row_ids post-wand. let docs_has_row_ids = self.docs.has_row_ids(); - let mut candidates = TopKCollector::new(limit, std::cmp::min(limit, BLOCK_SIZE * 10)); + // Deferred walks resolve real row_ids only post-wand, so the collector keeps boundary ties for + // row_id re-selection (see `TopKCollector::defer_tiebreak`). + let mut candidates = TopKCollector::new( + limit, + std::cmp::min(limit, BLOCK_SIZE * 10), + !docs_has_row_ids, + ); let mut num_comparisons = 0; let mut and_search_stats = (self.operator == Operator::And).then_some(AndSearchStats { pruned_before_return_start: self.and_candidates_pruned_before_return, @@ -2031,7 +2172,9 @@ impl<'a, S: Scorer> Wand<'a, S> { .unwrap_or(false); let mut num_comparisons = 0; - let mut candidates = TopKCollector::new(limit, 0); + // Flat search walks a resolved row_id allow-list, so the collector tiebreaks on real row_ids: + // no overflow needed. + let mut candidates = TopKCollector::new(limit, 0, false); for (doc_id, row_id) in doc_ids { num_comparisons += 1; self.move_head_before_target_to_tail(doc_id); @@ -2132,7 +2275,13 @@ impl<'a, S: Scorer> Wand<'a, S> { let total_sum_upper_bound_factor = score_sum_upper_bound_factor(clauses.len()); let mut acc = WindowAccumulator::new(clauses.len()); - let mut candidates = TopKCollector::new(limit, std::cmp::min(limit, BLOCK_SIZE * 10)); + // Deferred walks resolve real row_ids only post-wand, so the collector keeps boundary ties for + // row_id re-selection (see `TopKCollector::defer_tiebreak`). + let mut candidates = TopKCollector::new( + limit, + std::cmp::min(limit, BLOCK_SIZE * 10), + !docs_has_row_ids, + ); let norm_k = self.norm_k_cache(); let norm_k_ref = norm_k .as_ref() @@ -3021,7 +3170,13 @@ impl<'a, S: Scorer> Wand<'a, S> { } } - let mut candidates = TopKCollector::new(limit, std::cmp::min(limit, BLOCK_SIZE * 10)); + // Deferred walks resolve real row_ids only post-wand, so the collector keeps boundary ties for + // row_id re-selection (see `TopKCollector::defer_tiebreak`). + let mut candidates = TopKCollector::new( + limit, + std::cmp::min(limit, BLOCK_SIZE * 10), + !docs_has_row_ids, + ); let mut num_comparisons: usize = 0; let mut stats = AndSearchStats { pruned_before_return_start: self.and_candidates_pruned_before_return, @@ -4255,11 +4410,87 @@ mod tests { assert!(wand.norm_k_cache().is_none()); } + #[test] + fn test_top_k_collector_deferred_tiebreak_retains_boundary_ties() -> Result<()> { + // Deferred-row_id mode: the heap tiebreaks by the doc_id proxy, so every doc tied at the k-th + // score must survive (heap or overflow) for the post-wand real-row_id re-selection. `posting_doc_id` + // stands in for the doc_id here. + fn insert(collector: &mut TopKCollector, doc_id: u64, score: f32) -> Result { + collector.insert( + ScoredDoc::new(doc_id, score), + 1, + doc_id, + std::iter::once((0u32, 1u32)), + ) + } + fn candidate_doc_ids(collector: TopKCollector) -> Result> { + let mut ids = collector + .into_candidates(CandidateAddr::RowId)? + .into_iter() + .map(|candidate| candidate.posting_doc_id) + .collect::>(); + ids.sort_unstable(); + Ok(ids) + } + + // All tied at the k-th score: both the doc that wins the proxy tiebreak (lower doc_id 5 displaces + // 11) and the ones that lose it (20) must be retained, so the whole band reaches re-selection. + let mut collector = TopKCollector::new(2, 2, true); + assert!(insert(&mut collector, 10, 1.0)?); + assert!(insert(&mut collector, 11, 1.0)?); + assert!(insert(&mut collector, 5, 1.0)?); // Equal, wins proxy tiebreak -> displaces 11 to overflow + assert!(!insert(&mut collector, 20, 1.0)?); // Equal, loses proxy tiebreak -> overflow + assert_eq!(candidate_doc_ids(collector)?, vec![5, 10, 11, 20]); + + // When the k-th score rises, the old tie band can no longer win and must be dropped. + let mut collector = TopKCollector::new(2, 2, true); + assert!(insert(&mut collector, 0, 1.0)?); + assert!(insert(&mut collector, 1, 1.0)?); + assert!(!insert(&mut collector, 2, 1.0)?); // overflow at score 1.0 + assert!(insert(&mut collector, 3, 2.0)?); // heap still holds a 1.0 doc -> displaced 1.0 kept + assert!(insert(&mut collector, 4, 2.0)?); // k-th rises to 2.0 -> overflow cleared + assert_eq!(candidate_doc_ids(collector)?, vec![3, 4]); + + Ok(()) + } + + #[test] + fn test_top_k_collector_full_key_eviction_by_row_id() -> Result<()> { + // Non-deferred mode (heap holds real row_ids): a score-tie with a lower row_id evicts the current + // worst tied entry; a higher row_id at the tie, and any lower score, are rejected. No overflow. + fn insert(collector: &mut TopKCollector, row_id: u64, score: f32) -> Result { + collector.insert( + ScoredDoc::new(row_id, score), + 1, + row_id, + std::iter::once((0u32, 1u32)), + ) + } + let mut collector = TopKCollector::new(2, 2, false); + assert!(insert(&mut collector, 20, 1.0)?); + assert!(insert(&mut collector, 10, 1.0)?); // heap full {10, 20}, worst = 20 + assert!(insert(&mut collector, 5, 1.0)?); // ties, lower row_id -> evicts 20 + assert!(!insert(&mut collector, 30, 1.0)?); // ties, higher row_id -> rejected + assert!(!insert(&mut collector, 15, 0.5)?); // lower score -> rejected + + let mut ids = collector + .into_candidates(CandidateAddr::RowId)? + .into_iter() + .map(|candidate| match candidate.addr { + CandidateAddr::RowId(r) => r, + CandidateAddr::Pending(_) => unreachable!("non-deferred mode yields real row_ids"), + }) + .collect::>(); + ids.sort_unstable(); + assert_eq!(ids, vec![5, 10], "the two lowest row_ids among the tie"); + Ok(()) + } + #[test] fn test_top_k_collector_reuses_frequency_slots() -> Result<()> { const LIMIT: usize = 8; const NUM_DOCS: usize = 10_000; - let mut collector = TopKCollector::new(LIMIT, LIMIT); + let mut collector = TopKCollector::new(LIMIT, LIMIT, false); for doc in 0..NUM_DOCS { let num_terms = doc % 4 + 1; From 8fa1753384a72541aefc8d8854e3fbd697b1ba82 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=B6ren=20Brunk?= Date: Mon, 27 Jul 2026 16:26:22 +0200 Subject: [PATCH 2/2] chore(fts): address review feedback on deterministic top-k tiebreak - 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). --- rust/lance-index/src/scalar/inverted/index.rs | 24 ++++++++++++++++--- rust/lance-index/src/scalar/inverted/wand.rs | 22 +++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 2742105f486..f0ac15082cb 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -939,7 +939,14 @@ impl InvertedIndex { candidates.push(Reverse(candidate)); return; } - let kth = candidates.peek().unwrap().0.score; + // The `candidates.len() < limit` early return above guarantees the heap is full here, so + // `limit >= 1` (limit == 0 returns before any candidate is offered) makes these peeks/pops + // non-empty. `expect` documents that invariant rather than panicking silently. + let kth = candidates + .peek() + .expect("full top-k heap has a k-th element") + .0 + .score; match candidate.score.cmp(&kth) { // Below the k-th score: never competitive. std::cmp::Ordering::Less => {} @@ -949,9 +956,15 @@ impl InvertedIndex { // 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("full top-k heap has a worst element to evict"); candidates.push(Reverse(candidate)); - let new_kth = candidates.peek().unwrap().0.score; + let new_kth = candidates + .peek() + .expect("top-k heap is non-empty right after a push") + .0 + .score; if new_kth > displaced.score { overflow.clear(); } else { @@ -1233,6 +1246,10 @@ impl InvertedIndex { .try_collect() .await?; 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()); for ((pos, _), row_id) in entries.into_iter().zip(row_ids) { resolved[pos].1 = row_id; } @@ -11540,6 +11557,7 @@ mod tests { mask: Arc, } + #[cfg_attr(coverage, coverage(off))] #[async_trait::async_trait] impl PreFilter for AllowListFilter { async fn wait_for_ready(&self) -> Result<()> { diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 7378e3a15bd..924165c5348 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -4410,6 +4410,28 @@ mod tests { assert!(wand.norm_k_cache().is_none()); } + #[rstest] + #[case::one(1.0)] + #[case::small(f32::MIN_POSITIVE)] + #[case::large(1e30)] + fn test_admit_ties_floor_sits_just_below_kth(#[case] kth: f32) { + // A doc tied at the k-th score must pass the `score > threshold` admit test, so the floor sits + // strictly below the k-th score, exactly one ULP down. + let floor = admit_ties_floor(kth); + assert!(floor < kth, "floor {floor} must be < kth {kth}"); + assert_eq!(floor, kth.next_down()); + } + + #[test] + fn test_admit_ties_floor_zero_negative_and_nan() { + // Zero and negative inputs drop below zero, which the `threshold > 0.0` prune guards read as + // "no threshold yet" (pruning disabled). + assert!(admit_ties_floor(0.0) < 0.0); + assert!(admit_ties_floor(-1.0) < -1.0); + // NaN stays NaN; it is never a competitive score and is guarded at insert time. + assert!(admit_ties_floor(f32::NAN).is_nan()); + } + #[test] fn test_top_k_collector_deferred_tiebreak_retains_boundary_ties() -> Result<()> { // Deferred-row_id mode: the heap tiebreaks by the doc_id proxy, so every doc tied at the k-th