From aa084bd40ff4c9d31158c2d27a607159679389d4 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Sat, 4 Jul 2026 18:57:44 +0800 Subject: [PATCH 1/2] perf(index): norm-addend cache and slim top-k heap for fts scoring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two hot-loop changes, both result-identical (scores are bit-equal; only tie ordering among equal scores can shift with the slimmer heap tuple): - Scorer::doc_norm exposes the BM25 doc-length denominator addend, and quantized (v3) searches bake a per-norm-code addend cache once per partition search (Lucene's norm cache). The OR streaming/drain loops and the bulk conjunction pass score with one byte-norm load plus the cached addend instead of recomputing the denominator per doc — the factored expressions are evaluated identically, so scores don't move. - Top-k heap entries shrink to a Copy tuple; the per-candidate (term, freq) pairs go to an append-only side log and only the final top-k materialize a Vec, removing the per-insert allocation. Warm mmlb, 8 concurrent: OR@200M k10 0.0343->0.0249s (318qps), k100 0.0628->0.0467s (170qps) — both now ahead of Lucene 10.4's sliced searcher; AND@200M k10 0.0456->0.0443s, k100 0.0916->0.0883s; phrase@50M 2w 0.0392->0.0370s. Co-Authored-By: Claude Fable 5 --- .../lance-index/src/scalar/inverted/scorer.rs | 38 ++- rust/lance-index/src/scalar/inverted/wand.rs | 299 +++++++++++++----- 2 files changed, 256 insertions(+), 81 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/scorer.rs b/rust/lance-index/src/scalar/inverted/scorer.rs index 35a3be60e0a..03eba7ccf4d 100644 --- a/rust/lance-index/src/scalar/inverted/scorer.rs +++ b/rust/lance-index/src/scalar/inverted/scorer.rs @@ -27,6 +27,16 @@ pub trait Scorer: Send + Sync { fn doc_weight_cache_key(&self) -> Option { None } + + /// The doc-length-dependent BM25 denominator addend, when `doc_weight` + /// factors as `(K1 + 1) * freq / (freq + addend)`; `None` for scorers + /// without that shape. Scoring hot loops use this to bake a per-norm-code + /// addend cache (Lucene's norm cache), which is bit-identical to calling + /// `doc_weight` because both paths evaluate the same expressions. + fn doc_norm(&self, doc_tokens: u32) -> Option { + let _ = doc_tokens; + None + } } impl Scorer for Arc { @@ -45,6 +55,18 @@ impl Scorer for Arc { fn doc_weight_cache_key(&self) -> Option { self.as_ref().doc_weight_cache_key() } + + fn doc_norm(&self, doc_tokens: u32) -> Option { + self.as_ref().doc_norm(doc_tokens) + } +} + +/// The frequency-dependent half of the BM25 doc weight; `doc_norm` is the +/// doc-length addend (from [`Scorer::doc_norm`] or a per-norm-code cache). +#[inline] +pub(super) fn bm25_doc_weight_with_norm(freq: u32, doc_norm: f32) -> f32 { + let freq = freq as f32; + (K1 + 1.0) * freq / (freq + doc_norm) } // BM25 parameters @@ -112,10 +134,14 @@ impl Scorer for MemBM25Scorer { } fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32 { - let freq = freq as f32; let doc_tokens = doc_tokens as f32; let doc_norm = K1 * (1.0 - B + B * doc_tokens / self.avg_doc_length()); - (K1 + 1.0) * freq / (freq + doc_norm) + bm25_doc_weight_with_norm(freq, doc_norm) + } + + fn doc_norm(&self, doc_tokens: u32) -> Option { + let doc_tokens = doc_tokens as f32; + Some(K1 * (1.0 - B + B * doc_tokens / self.avg_doc_length())) } fn doc_weight_upper_bound(&self) -> Option { @@ -184,10 +210,14 @@ impl Scorer for IndexBM25Scorer<'_> { } fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32 { - let freq = freq as f32; let doc_tokens = doc_tokens as f32; let doc_norm = K1 * (1.0 - B + B * doc_tokens / self.avg_doc_length); - (K1 + 1.0) * freq / (freq + doc_norm) + bm25_doc_weight_with_norm(freq, doc_norm) + } + + fn doc_norm(&self, doc_tokens: u32) -> Option { + let doc_tokens = doc_tokens as f32; + Some(K1 * (1.0 - B + B * doc_tokens / self.avg_doc_length)) } fn doc_weight_upper_bound(&self) -> Option { diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index c887c342c99..fdb6f71ef9f 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -25,7 +25,7 @@ use super::{ impact::{IMPACT_LEVEL1_BLOCKS, ImpactScoreCache, ImpactSkipData}, index::{PositionStreamCodec, dequantize_doc_length}, query::Operator, - scorer::{K1, idf}, + scorer::{K1, bm25_doc_weight_with_norm, idf}, }; use super::{ CompressedPostingList, DocSet, PostingList, RawDocInfo, @@ -41,8 +41,40 @@ use super::{DocInfo, builder::BLOCK_SIZE}; const TERMINATED_DOC_ID: u64 = u64::MAX; -/// Top-k heap entry: (scored doc, (term, freq) pairs, doc length, posting doc id). -type TopKHeap = BinaryHeap, u32, u64)>>; +/// Top-k heap entry: (scored doc, doc length, posting doc id, freq-log index). +/// The (term, freq) pairs live in a side [`FreqLog`] so heap churn moves a +/// small Copy tuple instead of allocating a `Vec` per insert; only the final +/// top-k get their pairs materialized. +type TopKHeap = BinaryHeap>; + +/// Append-only (term, freq) log for heap candidates: flat storage plus +/// per-entry offsets (entries vary in width across paths). Discarded entries +/// just leak into the log until the search ends — inserts are far rarer than +/// candidates, so the log stays tiny. +#[derive(Default)] +struct FreqLog { + offsets: Vec, + pairs: Vec<(u32, u32)>, +} + +impl FreqLog { + fn push(&mut self, pairs: impl Iterator) -> u32 { + let index = self.offsets.len() as u32; + self.offsets.push(self.pairs.len() as u32); + self.pairs.extend(pairs); + index + } + + fn get(&self, index: u32) -> &[(u32, u32)] { + let start = self.offsets[index as usize] as usize; + let end = self + .offsets + .get(index as usize + 1) + .map(|&offset| offset as usize) + .unwrap_or(self.pairs.len()); + &self.pairs[start..end] + } +} const LINEAR_BLOCK_SKIP_LIMIT: usize = 8; pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { std::env::var("LANCE_FLAT_SEARCH_PERCENT_THRESHOLD") @@ -1077,6 +1109,9 @@ impl PostingIterator { /// first doc beyond `up_to`. This is the Lucene `nextDocsAndScores` /// equivalent: it walks the decompressed block arrays directly, with no /// per-doc heap traffic. + // The norm cache arg tips this hot-path fn over the limit; bundling the + // scoring inputs isn't worth the churn here. + #[allow(clippy::too_many_arguments)] fn collect_window_scores( &mut self, window_min: u64, @@ -1084,6 +1119,7 @@ impl PostingIterator { clause_idx: usize, docs: &DocSet, scorer: &S, + norm_k: Option<(&[u8], &[f32; 256])>, acc: &mut WindowAccumulator, ) { if self.doc().is_some_and(|doc| doc.doc_id() < window_min) { @@ -1113,8 +1149,16 @@ impl PostingIterator { break 'blocks; } let freq = compressed.freqs[offset]; - let doc_length = docs.scoring_num_tokens(doc_id); - let score = self.query_weight * scorer.doc_weight(freq, doc_length); + // One byte-norm load plus a cached addend replaces + // recomputing the BM25 denominator per doc. + let doc_weight = match norm_k { + Some((norms, cache)) => bm25_doc_weight_with_norm( + freq, + cache[norms[doc_id as usize] as usize], + ), + None => scorer.doc_weight(freq, docs.scoring_num_tokens(doc_id)), + }; + let score = self.query_weight * doc_weight; let slot = (u64::from(doc_id) - window_min) as usize; acc.add(clause_idx, slot, score, freq); } @@ -1500,6 +1544,26 @@ impl<'a, S: Scorer> Wand<'a, S> { self } + /// Per-search norm→BM25-denominator cache (Lucene's norm cache): the doc + /// byte-norm slab plus the 256 possible denominator addends. Available + /// when the DocSet scores quantized (V3 partitions) and the scorer + /// factors `doc_weight` as `(K1+1)*freq/(freq + addend)`. Scoring through + /// the cache is bit-identical to `scorer.doc_weight`, because both + /// evaluate the same expressions on the same quantized lengths. + fn norm_k_cache(&self) -> Option<(&'a [u8], Box<[f32; 256]>)> { + let docs: &'a DocSet = self.docs; + let norms = docs.scoring_norms()?; + self.scorer.doc_norm(0)?; + let mut cache = Box::new([0f32; 256]); + for (code, slot) in cache.iter_mut().enumerate() { + *slot = self + .scorer + .doc_norm(dequantize_doc_length(code as u8)) + .expect("doc_norm support was just probed"); + } + 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. fn update_threshold(&mut self, local_kth: f32, wand_factor: f32) { @@ -1591,6 +1655,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let docs_has_row_ids = self.docs.has_row_ids(); let mut candidates = BinaryHeap::with_capacity(std::cmp::min(limit, BLOCK_SIZE * 10)); + let mut freq_log = FreqLog::default(); 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, @@ -1664,31 +1729,31 @@ impl<'a, S: Scorer> Wand<'a, S> { }; if candidates.len() < limit { - let freqs = self.iter_term_freqs().collect(); + let log_idx = freq_log.push(self.iter_term_freqs()); if let Some(and_stats) = and_search_stats.as_mut() { and_stats.freqs_collected += 1; } candidates.push(Reverse(( ScoredDoc::new(row_id, score), - freqs, doc_length, posting_doc_id, + log_idx, ))); if candidates.len() == limit { let kth = candidates.peek().unwrap().0.0.score.0; self.update_threshold(kth, params.wand_factor); } } else if score > candidates.peek().unwrap().0.0.score.0 { - let freqs = self.iter_term_freqs().collect(); + let log_idx = freq_log.push(self.iter_term_freqs()); if let Some(and_stats) = and_search_stats.as_mut() { and_stats.freqs_collected += 1; } candidates.pop(); candidates.push(Reverse(( ScoredDoc::new(row_id, score), - freqs, doc_length, posting_doc_id, + log_idx, ))); let kth = candidates.peek().unwrap().0.0.score.0; self.update_threshold(kth, params.wand_factor); @@ -1732,10 +1797,10 @@ impl<'a, S: Scorer> Wand<'a, S> { Ok(candidates .into_iter() .map( - |Reverse((doc, freqs, doc_length, posting_doc_id))| DocCandidate { + |Reverse((doc, doc_length, posting_doc_id, log_idx))| DocCandidate { addr: to_addr(doc.row_id), posting_doc_id, - freqs, + freqs: freq_log.get(log_idx).to_vec(), doc_length, }, ) @@ -1779,7 +1844,8 @@ impl<'a, S: Scorer> Wand<'a, S> { .unwrap_or(false); let mut num_comparisons = 0; - let mut candidates = BinaryHeap::new(); + let mut candidates: TopKHeap = BinaryHeap::new(); + let mut freq_log = FreqLog::default(); for (doc_id, row_id) in doc_ids { num_comparisons += 1; self.move_head_before_target_to_tail(doc_id); @@ -1827,25 +1893,25 @@ impl<'a, S: Scorer> Wand<'a, S> { let score = self.score(doc_length); if candidates.len() < limit { - let freqs = self.iter_term_freqs().collect(); + let log_idx = freq_log.push(self.iter_term_freqs()); candidates.push(Reverse(( ScoredDoc::new(row_id, score), - freqs, doc_length, doc_id, + log_idx, ))); if candidates.len() == limit { let kth = candidates.peek().unwrap().0.0.score.0; self.update_threshold(kth, params.wand_factor); } } else if score > candidates.peek().unwrap().0.0.score.0 { - let freqs = self.iter_term_freqs().collect(); + let log_idx = freq_log.push(self.iter_term_freqs()); candidates.pop(); candidates.push(Reverse(( ScoredDoc::new(row_id, score), - freqs, doc_length, doc_id, + log_idx, ))); let kth = candidates.peek().unwrap().0.0.score.0; self.update_threshold(kth, params.wand_factor); @@ -1860,10 +1926,10 @@ impl<'a, S: Scorer> Wand<'a, S> { Ok(candidates .into_iter() .map( - |Reverse((doc, freqs, doc_length, posting_doc_id))| DocCandidate { + |Reverse((doc, doc_length, posting_doc_id, log_idx))| DocCandidate { addr: CandidateAddr::RowId(doc.row_id), posting_doc_id, - freqs, + freqs: freq_log.get(log_idx).to_vec(), doc_length, }, ) @@ -1907,6 +1973,11 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut acc = WindowAccumulator::new(clauses.len()); let mut candidates: TopKHeap = BinaryHeap::with_capacity(std::cmp::min(limit, BLOCK_SIZE * 10)); + let mut freq_log = FreqLog::default(); + let norm_k = self.norm_k_cache(); + let norm_k_ref = norm_k + .as_ref() + .map(|(norms, cache)| (*norms, cache.as_ref())); let mut num_comparisons = 0usize; // Adaptive minimum window size (Lucene): grow windows when they yield // too few candidates to amortize the per-window bound computations. @@ -2055,8 +2126,21 @@ impl<'a, S: Scorer> Wand<'a, S> { let doc = $doc; let freq = $freq; num_comparisons += 1; - let doc_length = self.docs.scoring_num_tokens(doc as u32); - let score = essential_weight * self.scorer.doc_weight(freq, doc_length); + // One byte-norm load + cached addend when available; + // the exact doc length is only needed at insert time. + let norm_addend = + norm_k_ref.map(|(norms, cache)| cache[norms[doc as usize] as usize]); + let score = match norm_addend { + Some(addend) => { + essential_weight * bm25_doc_weight_with_norm(freq, addend) + } + None => { + essential_weight + * self + .scorer + .doc_weight(freq, self.docs.scoring_num_tokens(doc as u32)) + } + }; if !(self.threshold > 0.0 && score_sum_cannot_exceed( score, @@ -2094,8 +2178,20 @@ impl<'a, S: Scorer> Wand<'a, S> { if let Some(d) = probe.doc() && d.doc_id() == doc { - total += - probe.score(&self.scorer, d.frequency(), doc_length); + total += match norm_addend { + Some(addend) => { + probe.query_weight + * bm25_doc_weight_with_norm( + d.frequency(), + addend, + ) + } + None => probe.score( + &self.scorer, + d.frequency(), + self.docs.scoring_num_tokens(doc as u32), + ), + }; } } @@ -2108,26 +2204,29 @@ impl<'a, S: Scorer> Wand<'a, S> { let beats_kth = !full || total > candidates.peek().unwrap().0.0.score.0; if beats_kth { - let mut freqs = Vec::with_capacity(non_essential.len() + 1); - freqs.push((essential_term, freq)); - for clause in non_essential.iter() { - if let Some(d) = clause.posting.doc() - && d.doc_id() == doc - { - freqs.push(( - clause.posting.term_index(), - d.frequency(), - )); - } - } + let doc_length = self.docs.scoring_num_tokens(doc as u32); + let log_idx = freq_log.push( + std::iter::once((essential_term, freq)).chain( + non_essential.iter().filter_map(|clause| { + clause.posting.doc().and_then(|d| { + (d.doc_id() == doc).then(|| { + ( + clause.posting.term_index(), + d.frequency(), + ) + }) + }) + }), + ), + ); if full { candidates.pop(); } candidates.push(Reverse(( ScoredDoc::new(row_id, total), - freqs, doc_length, doc, + log_idx, ))); if candidates.len() == limit { let kth = candidates.peek().unwrap().0.0.score.0; @@ -2226,6 +2325,7 @@ impl<'a, S: Scorer> Wand<'a, S> { clause_idx, self.docs, &self.scorer, + norm_k_ref, &mut acc, ); } @@ -2270,7 +2370,12 @@ impl<'a, S: Scorer> Wand<'a, S> { continue; } - let doc_length = self.docs.scoring_num_tokens(doc as u32); + // Doc length is only needed at heap-insert time; the + // non-essential completion scores go through the norm + // cache when available. + let norm_addend = + norm_k_ref.map(|(norms, cache)| cache[norms[doc as usize] as usize]); + let mut doc_length_cell: Option = None; let mut rejected = false; for i in (0..first_essential).rev() { if self.threshold > 0.0 @@ -2291,7 +2396,19 @@ impl<'a, S: Scorer> Wand<'a, S> { if let Some(d) = posting.doc() && d.doc_id() == doc { - score += posting.score(&self.scorer, d.frequency(), doc_length); + score += match norm_addend { + Some(addend) => { + posting.query_weight + * bm25_doc_weight_with_norm(d.frequency(), addend) + } + None => { + let doc_length = + *doc_length_cell.get_or_insert_with(|| { + self.docs.scoring_num_tokens(doc as u32) + }); + posting.score(&self.scorer, d.frequency(), doc_length) + } + }; } } @@ -2299,10 +2416,10 @@ impl<'a, S: Scorer> Wand<'a, S> { let full = candidates.len() >= limit; let beats_kth = !full || score > candidates.peek().unwrap().0.0.score.0; if beats_kth { - let freqs = clauses - .iter() - .enumerate() - .filter_map(|(i, clause)| { + let doc_length = doc_length_cell + .unwrap_or_else(|| self.docs.scoring_num_tokens(doc as u32)); + let log_idx = freq_log.push(clauses.iter().enumerate().filter_map( + |(i, clause)| { if i >= first_essential { let freq = acc.clause_freq(i, slot); (freq > 0).then(|| (clause.posting.term_index(), freq)) @@ -2313,16 +2430,16 @@ impl<'a, S: Scorer> Wand<'a, S> { }) }) } - }) - .collect::>(); + }, + )); if full { candidates.pop(); } candidates.push(Reverse(( ScoredDoc::new(row_id, score), - freqs, doc_length, doc, + log_idx, ))); if candidates.len() == limit { let kth = candidates.peek().unwrap().0.0.score.0; @@ -2357,10 +2474,10 @@ impl<'a, S: Scorer> Wand<'a, S> { Ok(candidates .into_iter() .map( - |Reverse((doc, freqs, doc_length, posting_doc_id))| DocCandidate { + |Reverse((doc, doc_length, posting_doc_id, log_idx))| DocCandidate { addr: to_addr(doc.row_id), posting_doc_id, - freqs, + freqs: freq_log.get(log_idx).to_vec(), doc_length, }, ) @@ -2797,7 +2914,15 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut batch_docs: Vec = Vec::with_capacity(MAX_POSTING_BLOCK_SIZE); let mut batch_offs: Vec = Vec::with_capacity(MAX_POSTING_BLOCK_SIZE * num_lists); let mut batch_lens: Vec = Vec::with_capacity(MAX_POSTING_BLOCK_SIZE); + let mut batch_norms: Vec = Vec::with_capacity(MAX_POSTING_BLOCK_SIZE); let mut cursor_scratch: Vec = Vec::with_capacity(num_lists); + let mut freq_log = FreqLog::default(); + // Norm cache: one byte-norm load plus a cached addend replaces the + // per-clause BM25 denominator recompute in pass B. + let norm_k = self.norm_k_cache(); + let norm_k_ref = norm_k + .as_ref() + .map(|(norms, cache)| (*norms, cache.as_ref())); // Per-window prune LUT for the merge kernels: an upper bound of the // first (rarest) clause's score by clamped frequency. Lead docs whose @@ -2981,35 +3106,52 @@ impl<'a, S: Scorer> Wand<'a, S> { ), } - // Pass A: gather doc lengths for the whole batch up front so - // the loads issue back-to-back and their cache misses overlap. - // Quantized (V3) sets gather through the byte-norm slab: a - // quarter of the bytes through the cache versus the u32 vec. + // Pass A: gather doc norms/lengths for the whole batch up + // front so the loads issue back-to-back and their cache + // misses overlap. With the norm cache, only the byte code is + // gathered; the exact length decodes from a tiny LUT. batch_lens.clear(); - match self.docs.scoring_norms() { - Some(norms) => { + batch_norms.clear(); + match norm_k_ref { + Some((norms, _)) => { for &doc in batch_docs.iter() { - batch_lens.push(dequantize_doc_length(norms[doc as usize])); + batch_norms.push(norms[doc as usize]); } } - None => { - for &doc in batch_docs.iter() { - batch_lens.push(self.docs.scoring_num_tokens(doc)); + None => match self.docs.scoring_norms() { + Some(norms) => { + for &doc in batch_docs.iter() { + batch_lens.push(dequantize_doc_length(norms[doc as usize])); + } } - } + None => { + for &doc in batch_docs.iter() { + batch_lens.push(self.docs.scoring_num_tokens(doc)); + } + } + }, } // Pass B: prune / verify / score / insert, in doc order, with // exactly the classic loop's semantics. for (index, &doc) in batch_docs.iter().enumerate() { - let doc_length = batch_lens[index]; + let (norm_addend, doc_length) = match norm_k_ref { + Some((_, cache)) => { + let code = batch_norms[index]; + (Some(cache[code as usize]), dequantize_doc_length(code)) + } + None => (None, batch_lens[index]), + }; let offs = &batch_offs[index * num_lists..(index + 1) * num_lists]; if self.threshold > 0.0 && num_lists >= 2 { - let first_score = self.lead[0].score( - &self.scorer, - unsafe { *wins[0].freqs.add(offs[0] as usize) }, - doc_length, - ); + let first_freq = unsafe { *wins[0].freqs.add(offs[0] as usize) }; + let first_score = match norm_addend { + Some(addend) => { + self.lead[0].query_weight + * bm25_doc_weight_with_norm(first_freq, addend) + } + None => self.lead[0].score(&self.scorer, first_freq, doc_length), + }; if first_score + others_block_max <= self.threshold { self.and_candidates_pruned_before_return += 1; continue; @@ -3059,7 +3201,12 @@ impl<'a, S: Scorer> Wand<'a, S> { for ((win, posting), &off) in wins.iter().zip(self.lead.iter()).zip(offs.iter()) { let freq = unsafe { *win.freqs.add(off as usize) }; - score += posting.score(&self.scorer, freq, doc_length); + score += match norm_addend { + Some(addend) => { + posting.query_weight * bm25_doc_weight_with_norm(freq, addend) + } + None => posting.score(&self.scorer, freq, doc_length), + }; } let insert = if candidates.len() < limit { @@ -3069,24 +3216,22 @@ impl<'a, S: Scorer> Wand<'a, S> { }; if insert { stats.freqs_collected += 1; - let freqs = wins - .iter() - .zip(self.lead.iter()) - .zip(offs.iter()) - .map(|((win, posting), &off)| { - (posting.term_index(), unsafe { - *win.freqs.add(off as usize) - }) - }) - .collect(); + let log_idx = + freq_log.push(wins.iter().zip(self.lead.iter()).zip(offs.iter()).map( + |((win, posting), &off)| { + (posting.term_index(), unsafe { + *win.freqs.add(off as usize) + }) + }, + )); if candidates.len() >= limit { candidates.pop(); } candidates.push(Reverse(( ScoredDoc::new(row_id, score), - freqs, doc_length, u64::from(doc), + log_idx, ))); if candidates.len() == limit { let kth = candidates.peek().unwrap().0.0.score.0; @@ -3129,10 +3274,10 @@ impl<'a, S: Scorer> Wand<'a, S> { Ok(candidates .into_iter() .map( - |Reverse((doc, freqs, doc_length, posting_doc_id))| DocCandidate { + |Reverse((doc, doc_length, posting_doc_id, log_idx))| DocCandidate { addr: to_addr(doc.row_id), posting_doc_id, - freqs, + freqs: freq_log.get(log_idx).to_vec(), doc_length, }, ) From 5ce91ee5e1fd79115965f49e16007203e6c97e32 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Wed, 15 Jul 2026 23:53:13 +0800 Subject: [PATCH 2/2] fix(index): bound FTS top-k frequency storage --- .../lance-index/src/scalar/inverted/scorer.rs | 18 +- rust/lance-index/src/scalar/inverted/wand.rs | 499 ++++++++++-------- 2 files changed, 291 insertions(+), 226 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/scorer.rs b/rust/lance-index/src/scalar/inverted/scorer.rs index 03eba7ccf4d..3a33a67ff7a 100644 --- a/rust/lance-index/src/scalar/inverted/scorer.rs +++ b/rust/lance-index/src/scalar/inverted/scorer.rs @@ -73,6 +73,12 @@ pub(super) fn bm25_doc_weight_with_norm(freq: u32, doc_norm: f32) -> f32 { pub const K1: f32 = 1.2; pub const B: f32 = 0.75; +#[inline] +fn bm25_doc_norm(doc_tokens: u32, avg_doc_length: f32) -> f32 { + let doc_tokens = doc_tokens as f32; + K1 * (1.0 - B + B * doc_tokens / avg_doc_length) +} + #[derive(Debug, Clone)] pub struct MemBM25Scorer { pub total_tokens: u64, @@ -134,14 +140,12 @@ impl Scorer for MemBM25Scorer { } fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32 { - let doc_tokens = doc_tokens as f32; - let doc_norm = K1 * (1.0 - B + B * doc_tokens / self.avg_doc_length()); + let doc_norm = bm25_doc_norm(doc_tokens, self.avg_doc_length()); bm25_doc_weight_with_norm(freq, doc_norm) } fn doc_norm(&self, doc_tokens: u32) -> Option { - let doc_tokens = doc_tokens as f32; - Some(K1 * (1.0 - B + B * doc_tokens / self.avg_doc_length())) + Some(bm25_doc_norm(doc_tokens, self.avg_doc_length())) } fn doc_weight_upper_bound(&self) -> Option { @@ -210,14 +214,12 @@ impl Scorer for IndexBM25Scorer<'_> { } fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32 { - let doc_tokens = doc_tokens as f32; - let doc_norm = K1 * (1.0 - B + B * doc_tokens / self.avg_doc_length); + let doc_norm = bm25_doc_norm(doc_tokens, self.avg_doc_length); bm25_doc_weight_with_norm(freq, doc_norm) } fn doc_norm(&self, doc_tokens: u32) -> Option { - let doc_tokens = doc_tokens as f32; - Some(K1 * (1.0 - B + B * doc_tokens / self.avg_doc_length)) + Some(bm25_doc_norm(doc_tokens, self.avg_doc_length)) } fn doc_weight_upper_bound(&self) -> Option { diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index fdb6f71ef9f..465d0f8e2ab 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -41,38 +41,159 @@ use super::{DocInfo, builder::BLOCK_SIZE}; const TERMINATED_DOC_ID: u64 = u64::MAX; -/// Top-k heap entry: (scored doc, doc length, posting doc id, freq-log index). -/// The (term, freq) pairs live in a side [`FreqLog`] so heap churn moves a -/// small Copy tuple instead of allocating a `Vec` per insert; only the final -/// top-k get their pairs materialized. +/// Top-k heap entry: (scored doc, doc length, posting doc id, frequency slot). +/// The (term, freq) pairs live outside the heap so heap churn moves a compact +/// tuple instead of a `Vec`. type TopKHeap = BinaryHeap>; -/// Append-only (term, freq) log for heap candidates: flat storage plus -/// per-entry offsets (entries vary in width across paths). Discarded entries -/// just leak into the log until the search ends — inserts are far rarer than -/// candidates, so the log stays tiny. -#[derive(Default)] -struct FreqLog { - offsets: Vec, - pairs: Vec<(u32, u32)>, +/// Reusable (term, freq) storage for live heap candidates. Replacing the k-th +/// result reuses its slot and retained `Vec` capacity, so memory stays bounded +/// by the live top-k rather than every historical heap admission. +struct FrequencySlots { + slots: Vec>, } -impl FreqLog { - fn push(&mut self, pairs: impl Iterator) -> u32 { - let index = self.offsets.len() as u32; - self.offsets.push(self.pairs.len() as u32); - self.pairs.extend(pairs); - index +impl FrequencySlots { + fn with_capacity(capacity: usize) -> Self { + Self { + slots: Vec::with_capacity(capacity), + } + } + + fn push(&mut self, pairs: impl Iterator) -> Result { + let slot = u32::try_from(self.slots.len()).map_err(|_| { + Error::internal(format!( + "FTS top-k frequency slot count {} exceeds u32::MAX", + self.slots.len() + )) + })?; + self.slots.push(pairs.collect()); + Ok(slot) + } + + fn replace(&mut self, slot: u32, pairs: impl Iterator) -> Result<()> { + let num_slots = self.slots.len(); + let slot = self.slots.get_mut(slot as usize).ok_or_else(|| { + Error::internal(format!( + "FTS top-k frequency slot {slot} is out of bounds for {num_slots} slots" + )) + })?; + slot.clear(); + slot.extend(pairs); + Ok(()) } - fn get(&self, index: u32) -> &[(u32, u32)] { - let start = self.offsets[index as usize] as usize; - let end = self - .offsets - .get(index as usize + 1) - .map(|&offset| offset as usize) - .unwrap_or(self.pairs.len()); - &self.pairs[start..end] + fn take(&mut self, slot: u32) -> Result> { + let num_slots = self.slots.len(); + self.slots + .get_mut(slot as usize) + .map(std::mem::take) + .ok_or_else(|| { + Error::internal(format!( + "FTS top-k frequency slot {slot} is out of bounds for {num_slots} slots" + )) + }) + } +} + +/// Owns the top-k heap and the frequency slots referenced by its entries. +struct TopKCollector { + limit: usize, + heap: TopKHeap, + frequency_slots: FrequencySlots, +} + +impl TopKCollector { + fn new(limit: usize, initial_capacity: usize) -> Self { + let initial_capacity = initial_capacity.min(limit); + Self { + limit, + heap: BinaryHeap::with_capacity(initial_capacity), + frequency_slots: FrequencySlots::with_capacity(initial_capacity), + } + } + + /// 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( + &mut self, + doc: ScoredDoc, + doc_length: u32, + posting_doc_id: u64, + pairs: impl Iterator, + ) -> Result { + if self.limit == 0 { + 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 + ))); + } + + 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) { + 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)? + }; + + self.heap + .push(Reverse((doc, doc_length, posting_doc_id, frequency_slot))); + Ok(true) + } + + fn kth_score_if_full(&self) -> Option { + if self.heap.len() == self.limit { + self.heap.peek().map(|entry| entry.0.0.score.0) + } else { + None + } + } + + fn into_candidates( + self, + mut to_addr: impl FnMut(u64) -> CandidateAddr, + ) -> Result> { + let Self { + heap, + mut frequency_slots, + .. + } = self; + heap.into_iter() + .map( + |Reverse((doc, doc_length, posting_doc_id, frequency_slot))| { + Ok(DocCandidate { + addr: to_addr(doc.row_id), + posting_doc_id, + freqs: frequency_slots.take(frequency_slot)?, + doc_length, + }) + }, + ) + .collect() + } + + #[cfg(test)] + fn num_frequency_slots(&self) -> usize { + self.frequency_slots.slots.len() } } const LINEAR_BLOCK_SKIP_LIMIT: usize = 8; @@ -1553,13 +1674,9 @@ impl<'a, S: Scorer> Wand<'a, S> { fn norm_k_cache(&self) -> Option<(&'a [u8], Box<[f32; 256]>)> { let docs: &'a DocSet = self.docs; let norms = docs.scoring_norms()?; - self.scorer.doc_norm(0)?; let mut cache = Box::new([0f32; 256]); for (code, slot) in cache.iter_mut().enumerate() { - *slot = self - .scorer - .doc_norm(dequantize_doc_length(code as u8)) - .expect("doc_norm support was just probed"); + *slot = self.scorer.doc_norm(dequantize_doc_length(code as u8))?; } Some((norms, cache)) } @@ -1654,8 +1771,7 @@ impl<'a, S: Scorer> Wand<'a, S> { // row_ids post-wand. let docs_has_row_ids = self.docs.has_row_ids(); - let mut candidates = BinaryHeap::with_capacity(std::cmp::min(limit, BLOCK_SIZE * 10)); - let mut freq_log = FreqLog::default(); + let mut candidates = TopKCollector::new(limit, std::cmp::min(limit, BLOCK_SIZE * 10)); 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, @@ -1728,35 +1844,18 @@ impl<'a, S: Scorer> Wand<'a, S> { self.score(doc_length) }; - if candidates.len() < limit { - let log_idx = freq_log.push(self.iter_term_freqs()); + if candidates.insert( + ScoredDoc::new(row_id, score), + doc_length, + posting_doc_id, + self.iter_term_freqs(), + )? { if let Some(and_stats) = and_search_stats.as_mut() { and_stats.freqs_collected += 1; } - candidates.push(Reverse(( - ScoredDoc::new(row_id, score), - doc_length, - posting_doc_id, - log_idx, - ))); - if candidates.len() == limit { - let kth = candidates.peek().unwrap().0.0.score.0; + if let Some(kth) = candidates.kth_score_if_full() { self.update_threshold(kth, params.wand_factor); } - } else if score > candidates.peek().unwrap().0.0.score.0 { - let log_idx = freq_log.push(self.iter_term_freqs()); - if let Some(and_stats) = and_search_stats.as_mut() { - and_stats.freqs_collected += 1; - } - candidates.pop(); - candidates.push(Reverse(( - ScoredDoc::new(row_id, score), - doc_length, - posting_doc_id, - log_idx, - ))); - let kth = candidates.peek().unwrap().0.0.score.0; - self.update_threshold(kth, params.wand_factor); } if self.operator == Operator::Or { self.push_back_leads(doc.doc_id() + 1); @@ -1794,17 +1893,7 @@ impl<'a, S: Scorer> Wand<'a, S> { CandidateAddr::Pending(row_id_slot as u32) } }; - Ok(candidates - .into_iter() - .map( - |Reverse((doc, doc_length, posting_doc_id, log_idx))| DocCandidate { - addr: to_addr(doc.row_id), - posting_doc_id, - freqs: freq_log.get(log_idx).to_vec(), - doc_length, - }, - ) - .collect()) + candidates.into_candidates(to_addr) } fn flat_search( @@ -1844,8 +1933,7 @@ impl<'a, S: Scorer> Wand<'a, S> { .unwrap_or(false); let mut num_comparisons = 0; - let mut candidates: TopKHeap = BinaryHeap::new(); - let mut freq_log = FreqLog::default(); + let mut candidates = TopKCollector::new(limit, 0); for (doc_id, row_id) in doc_ids { num_comparisons += 1; self.move_head_before_target_to_tail(doc_id); @@ -1892,28 +1980,13 @@ impl<'a, S: Scorer> Wand<'a, S> { self.collect_tail_matches(doc_id); let score = self.score(doc_length); - if candidates.len() < limit { - let log_idx = freq_log.push(self.iter_term_freqs()); - candidates.push(Reverse(( - ScoredDoc::new(row_id, score), - doc_length, - doc_id, - log_idx, - ))); - if candidates.len() == limit { - let kth = candidates.peek().unwrap().0.0.score.0; - self.update_threshold(kth, params.wand_factor); - } - } else if score > candidates.peek().unwrap().0.0.score.0 { - let log_idx = freq_log.push(self.iter_term_freqs()); - candidates.pop(); - candidates.push(Reverse(( - ScoredDoc::new(row_id, score), - doc_length, - doc_id, - log_idx, - ))); - let kth = candidates.peek().unwrap().0.0.score.0; + if candidates.insert( + ScoredDoc::new(row_id, score), + doc_length, + doc_id, + self.iter_term_freqs(), + )? && let Some(kth) = candidates.kth_score_if_full() + { self.update_threshold(kth, params.wand_factor); } @@ -1923,17 +1996,7 @@ impl<'a, S: Scorer> Wand<'a, S> { // flat_search is driven by an explicit row_ids iterator, so // every candidate already has a real row_id. - Ok(candidates - .into_iter() - .map( - |Reverse((doc, doc_length, posting_doc_id, log_idx))| DocCandidate { - addr: CandidateAddr::RowId(doc.row_id), - posting_doc_id, - freqs: freq_log.get(log_idx).to_vec(), - doc_length, - }, - ) - .collect()) + candidates.into_candidates(CandidateAddr::RowId) } /// Bulk MAXSCORE top-k disjunction, mirroring Lucene's MaxScoreBulkScorer. @@ -1971,9 +2034,7 @@ 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: TopKHeap = - BinaryHeap::with_capacity(std::cmp::min(limit, BLOCK_SIZE * 10)); - let mut freq_log = FreqLog::default(); + let mut candidates = TopKCollector::new(limit, std::cmp::min(limit, BLOCK_SIZE * 10)); let norm_k = self.norm_k_cache(); let norm_k_ref = norm_k .as_ref() @@ -2200,38 +2261,23 @@ impl<'a, S: Scorer> Wand<'a, S> { // which drops zero-score matches (e.g. terms // with idf 0) exactly like Wand::next does. if !rejected && total > self.threshold { - let full = candidates.len() >= limit; - let beats_kth = - !full || total > candidates.peek().unwrap().0.0.score.0; - if beats_kth { - let doc_length = self.docs.scoring_num_tokens(doc as u32); - let log_idx = freq_log.push( - std::iter::once((essential_term, freq)).chain( - non_essential.iter().filter_map(|clause| { - clause.posting.doc().and_then(|d| { - (d.doc_id() == doc).then(|| { - ( - clause.posting.term_index(), - d.frequency(), - ) - }) + let doc_length = self.docs.scoring_num_tokens(doc as u32); + if candidates.insert( + ScoredDoc::new(row_id, total), + doc_length, + doc, + std::iter::once((essential_term, freq)).chain( + non_essential.iter().filter_map(|clause| { + clause.posting.doc().and_then(|d| { + (d.doc_id() == doc).then(|| { + (clause.posting.term_index(), d.frequency()) }) - }), - ), - ); - if full { - candidates.pop(); - } - candidates.push(Reverse(( - ScoredDoc::new(row_id, total), - doc_length, - doc, - log_idx, - ))); - if candidates.len() == limit { - let kth = candidates.peek().unwrap().0.0.score.0; - self.update_threshold(kth, params.wand_factor); - } + }) + }), + ), + )? && let Some(kth) = candidates.kth_score_if_full() + { + self.update_threshold(kth, params.wand_factor); } } } @@ -2413,38 +2459,27 @@ impl<'a, S: Scorer> Wand<'a, S> { } if !rejected && score > self.threshold { - let full = candidates.len() >= limit; - let beats_kth = !full || score > candidates.peek().unwrap().0.0.score.0; - if beats_kth { - let doc_length = doc_length_cell - .unwrap_or_else(|| self.docs.scoring_num_tokens(doc as u32)); - let log_idx = freq_log.push(clauses.iter().enumerate().filter_map( - |(i, clause)| { - if i >= first_essential { - let freq = acc.clause_freq(i, slot); - (freq > 0).then(|| (clause.posting.term_index(), freq)) - } else { - clause.posting.doc().and_then(|d| { - (d.doc_id() == doc).then(|| { - (clause.posting.term_index(), d.frequency()) - }) + let doc_length = doc_length_cell + .unwrap_or_else(|| self.docs.scoring_num_tokens(doc as u32)); + if candidates.insert( + ScoredDoc::new(row_id, score), + doc_length, + doc, + clauses.iter().enumerate().filter_map(|(i, clause)| { + if i >= first_essential { + let freq = acc.clause_freq(i, slot); + (freq > 0).then(|| (clause.posting.term_index(), freq)) + } else { + clause.posting.doc().and_then(|d| { + (d.doc_id() == doc).then(|| { + (clause.posting.term_index(), d.frequency()) }) - } - }, - )); - if full { - candidates.pop(); - } - candidates.push(Reverse(( - ScoredDoc::new(row_id, score), - doc_length, - doc, - log_idx, - ))); - if candidates.len() == limit { - let kth = candidates.peek().unwrap().0.0.score.0; - self.update_threshold(kth, params.wand_factor); - } + }) + } + }), + )? && let Some(kth) = candidates.kth_score_if_full() + { + self.update_threshold(kth, params.wand_factor); } } acc.clear_slot(slot); @@ -2471,17 +2506,7 @@ impl<'a, S: Scorer> Wand<'a, S> { CandidateAddr::Pending(row_id_slot as u32) } }; - Ok(candidates - .into_iter() - .map( - |Reverse((doc, doc_length, posting_doc_id, log_idx))| DocCandidate { - addr: to_addr(doc.row_id), - posting_doc_id, - freqs: freq_log.get(log_idx).to_vec(), - doc_length, - }, - ) - .collect()) + candidates.into_candidates(to_addr) } // calculate the score of the current document @@ -2898,8 +2923,7 @@ impl<'a, S: Scorer> Wand<'a, S> { } } - let mut candidates: TopKHeap = - BinaryHeap::with_capacity(std::cmp::min(limit, BLOCK_SIZE * 10)); + let mut candidates = TopKCollector::new(limit, std::cmp::min(limit, BLOCK_SIZE * 10)); let mut num_comparisons: usize = 0; let mut stats = AndSearchStats { pruned_before_return_start: self.and_candidates_pruned_before_return, @@ -2916,7 +2940,6 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut batch_lens: Vec = Vec::with_capacity(MAX_POSTING_BLOCK_SIZE); let mut batch_norms: Vec = Vec::with_capacity(MAX_POSTING_BLOCK_SIZE); let mut cursor_scratch: Vec = Vec::with_capacity(num_lists); - let mut freq_log = FreqLog::default(); // Norm cache: one byte-norm load plus a cached addend replaces the // per-clause BM25 denominator recompute in pass B. let norm_k = self.norm_k_cache(); @@ -3209,32 +3232,20 @@ impl<'a, S: Scorer> Wand<'a, S> { }; } - let insert = if candidates.len() < limit { - true - } else { - score > candidates.peek().unwrap().0.0.score.0 - }; - if insert { + if candidates.insert( + ScoredDoc::new(row_id, score), + doc_length, + u64::from(doc), + wins.iter().zip(self.lead.iter()).zip(offs.iter()).map( + |((win, posting), &off)| { + (posting.term_index(), unsafe { + *win.freqs.add(off as usize) + }) + }, + ), + )? { stats.freqs_collected += 1; - let log_idx = - freq_log.push(wins.iter().zip(self.lead.iter()).zip(offs.iter()).map( - |((win, posting), &off)| { - (posting.term_index(), unsafe { - *win.freqs.add(off as usize) - }) - }, - )); - if candidates.len() >= limit { - candidates.pop(); - } - candidates.push(Reverse(( - ScoredDoc::new(row_id, score), - doc_length, - u64::from(doc), - log_idx, - ))); - if candidates.len() == limit { - let kth = candidates.peek().unwrap().0.0.score.0; + if let Some(kth) = candidates.kth_score_if_full() { self.update_threshold(kth, params.wand_factor); } } @@ -3271,17 +3282,7 @@ impl<'a, S: Scorer> Wand<'a, S> { CandidateAddr::Pending(row_id_slot as u32) } }; - Ok(candidates - .into_iter() - .map( - |Reverse((doc, doc_length, posting_doc_id, log_idx))| DocCandidate { - addr: to_addr(doc.row_id), - posting_doc_id, - freqs: freq_log.get(log_idx).to_vec(), - doc_length, - }, - ) - .collect()) + candidates.into_candidates(to_addr) } fn and_move_to_next_block(&mut self, target: u64) { @@ -4130,6 +4131,68 @@ mod tests { } } + struct PartialNormScorer; + + impl Scorer for PartialNormScorer { + fn query_weight(&self, _token: &str) -> f32 { + 1.0 + } + + fn doc_weight(&self, freq: u32, _doc_tokens: u32) -> f32 { + freq as f32 + } + + fn doc_norm(&self, doc_tokens: u32) -> Option { + (doc_tokens == 0).then_some(0.0) + } + } + + #[test] + fn test_norm_cache_requires_all_norm_codes() { + let mut docs = DocSet::default(); + docs.append(0, 1); + docs.set_quantized_scoring(true); + let wand = Wand::new(Operator::Or, std::iter::empty(), &docs, PartialNormScorer); + + assert!(wand.norm_k_cache().is_none()); + } + + #[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); + + for doc in 0..NUM_DOCS { + let num_terms = doc % 4 + 1; + let inserted = collector.insert( + ScoredDoc::new(doc as u64, doc as f32), + num_terms as u32, + doc as u64, + (0..num_terms).map(|term| (term as u32, doc as u32)), + )?; + assert!(inserted); + assert!(collector.num_frequency_slots() <= LIMIT); + } + assert_eq!(collector.num_frequency_slots(), LIMIT); + + let mut candidates = collector.into_candidates(CandidateAddr::RowId)?; + candidates.sort_unstable_by_key(|candidate| candidate.posting_doc_id); + assert_eq!(candidates.len(), LIMIT); + for (candidate, expected_doc) in candidates.iter().zip(NUM_DOCS - LIMIT..NUM_DOCS) { + assert_eq!(candidate.posting_doc_id, expected_doc as u64); + assert!(matches!( + candidate.addr, + CandidateAddr::RowId(row_id) if row_id == expected_doc as u64 + )); + let expected_freqs = (0..expected_doc % 4 + 1) + .map(|term| (term as u32, expected_doc as u32)) + .collect::>(); + assert_eq!(candidate.freqs, expected_freqs); + } + Ok(()) + } + #[rstest] #[case::auto("auto", Some(BulkAndMode::Auto))] #[case::auto_case_and_whitespace(" AUTO ", Some(BulkAndMode::Auto))]