From c4b12f1f5d2579252df8f6cbd0f72d6f58277fab Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Sun, 5 Jul 2026 01:48:56 +0800 Subject: [PATCH 1/7] feat(index)!: quantized doc-length scoring and slimmer posting blocks for 256-doc blocks Folds the v3 breaking changes into this PR so the format/scoring break lands in one place; 128-block indexes are untouched on disk and keep exact scoring bit-for-bit. - BM25 doc lengths for 256-doc-block partitions are quantized to a Lucene SmallFloat-style byte code (4 mantissa bits, exact below 8, <= 6.25% relative error, decode = bucket floor). The byte-norm slab bakes lazily per loaded DocSet, quartering the doc-length bytes the scoring path pulls through the cache (200M docs: 800MB -> 200MB). - 256-doc posting blocks drop the leading block-max-score f32. The impact skip data introduced by the stacked follow-up supplies a tighter per-block bound; until it lands, block-level pruning for 256 falls back to the list-level max score, which is a valid (looser) bound. Block layout: [first_doc u32][doc num_bits u8][docs][pfor freqs]; posting_block_score_prefix_len(block_size) keys every reader/writer. BREAKING: 256-doc-block (v3) indexes must be rebuilt; v3 is unreleased so no migration is provided. BM25 scores on v3 differ from exact-length BM25 by the norm quantization, matching Lucene's norm semantics. Co-Authored-By: Claude Fable 5 --- .../src/scalar/inverted/encoding.rs | 73 +++++--- rust/lance-index/src/scalar/inverted/index.rs | 161 +++++++++++++++++- rust/lance-index/src/scalar/inverted/iter.rs | 1 + .../src/scalar/inverted/lazy_docset.rs | 13 +- rust/lance-index/src/scalar/inverted/wand.rs | 9 +- 5 files changed, 225 insertions(+), 32 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/encoding.rs b/rust/lance-index/src/scalar/inverted/encoding.rs index c8731dc1929..6296de077c0 100644 --- a/rust/lance-index/src/scalar/inverted/encoding.rs +++ b/rust/lance-index/src/scalar/inverted/encoding.rs @@ -75,9 +75,11 @@ pub fn compress_posting_list_with_tail_codec_and_block_size<'a>( if length < block_size { // directly do remainder compression to avoid overhead of creating buffer let mut builder = LargeBinaryBuilder::with_capacity(1, length * 4 * 2 + 1); - // write the max score of the block - let max_score = block_max_scores.next().unwrap(); - let _ = builder.write(max_score.to_le_bytes().as_ref())?; + // write the max score of the block (128-doc blocks only) + if posting_block_score_prefix_len(block_size) > 0 { + let max_score = block_max_scores.next().unwrap(); + let _ = builder.write(max_score.to_le_bytes().as_ref())?; + } compress_posting_remainder( doc_ids.copied().collect::>().as_slice(), frequencies.copied().collect::>().as_slice(), @@ -101,9 +103,11 @@ pub fn compress_posting_list_with_tail_codec_and_block_size<'a>( assert_eq!(doc_id_buffer.len(), block_size); - // write the max score of the block - let max_score = block_max_scores.next().unwrap(); - let _ = builder.write(max_score.to_le_bytes().as_ref())?; + // write the max score of the block (128-doc blocks only) + if posting_block_score_prefix_len(block_size) > 0 { + let max_score = block_max_scores.next().unwrap(); + let _ = builder.write(max_score.to_le_bytes().as_ref())?; + } encode_posting_block_payload(&doc_id_buffer, &freq_buffer, &mut builder)?; builder.append_value(""); doc_id_buffer.clear(); @@ -112,15 +116,30 @@ pub fn compress_posting_list_with_tail_codec_and_block_size<'a>( // we don't compress the last block if it is not full if !doc_id_buffer.is_empty() { - // write the max score of the block - let max_score = block_max_scores.next().unwrap(); - let _ = builder.write(max_score.to_le_bytes().as_ref())?; + // write the max score of the block (128-doc blocks only) + if posting_block_score_prefix_len(block_size) > 0 { + let max_score = block_max_scores.next().unwrap(); + let _ = builder.write(max_score.to_le_bytes().as_ref())?; + } compress_posting_remainder(&doc_id_buffer, &freq_buffer, tail_codec, &mut builder)?; builder.append_value(""); } Ok(builder.finish()) } +/// Byte length of the block-max-score prefix on posting blocks. 128-doc +/// blocks store a per-block max score, patched in at build time; 256-doc +/// (V3) blocks always carry impact skip data, which supersedes it, so they +/// store none. +#[inline] +pub fn posting_block_score_prefix_len(block_size: usize) -> usize { + if block_size == MAX_POSTING_BLOCK_SIZE { + 0 + } else { + 4 + } +} + pub fn encode_full_posting_block_into( doc_ids: &[u32], frequencies: &[u32], @@ -128,7 +147,9 @@ pub fn encode_full_posting_block_into( ) -> Result<()> { validate_block_size(doc_ids.len())?; debug_assert_eq!(doc_ids.len(), frequencies.len()); - block.extend_from_slice(&0f32.to_le_bytes()); + if posting_block_score_prefix_len(doc_ids.len()) > 0 { + block.extend_from_slice(&0f32.to_le_bytes()); + } encode_posting_block_payload(doc_ids, frequencies, block)?; Ok(()) } @@ -252,10 +273,13 @@ pub fn encode_remainder_posting_block_into( doc_ids: &[u32], frequencies: &[u32], codec: PostingTailCodec, + block_size: usize, block: &mut Vec, ) -> Result<()> { debug_assert_eq!(doc_ids.len(), frequencies.len()); - block.extend_from_slice(&0f32.to_le_bytes()); + if posting_block_score_prefix_len(block_size) > 0 { + block.extend_from_slice(&0f32.to_le_bytes()); + } compress_posting_remainder(doc_ids, frequencies, codec, block)?; Ok(()) } @@ -843,6 +867,7 @@ pub fn decompress_posting_list_with_tail_codec_and_block_size( compressed, remainder, tail_codec, + block_size, &mut doc_ids, &mut frequencies, ); @@ -884,8 +909,8 @@ pub fn decompress_posting_block( ) { debug_assert!(validate_block_size(block_size).is_ok()); debug_assert!(buffer.len() >= block_size); - // skip the first 4 bytes for the max block score - let mut block = &block[4..]; + // skip the block max score prefix (128-doc blocks only) + let mut block = &block[posting_block_score_prefix_len(block_size)..]; match block_size { BitPacker4x::BLOCK_LEN => { let num_bytes = decompress_sorted_block_with::(block, buffer, doc_ids); @@ -912,10 +937,11 @@ pub fn decompress_posting_remainder( block: &[u8], n: usize, codec: PostingTailCodec, + block_size: usize, doc_ids: &mut Vec, frequencies: &mut Vec, ) { - let block = &block[4..]; + let block = &block[posting_block_score_prefix_len(block_size)..]; match codec { PostingTailCodec::Fixed32 => { decompress_raw_remainder(block, n, doc_ids); @@ -1006,11 +1032,18 @@ pub fn decompress_raw_remainder(compressed: &[u8], n: usize, dest: &mut Vec } } -pub fn read_posting_tail_first_doc(block: &[u8], codec: PostingTailCodec) -> u32 { +pub fn read_posting_tail_first_doc( + block: &[u8], + codec: PostingTailCodec, + block_size: usize, +) -> u32 { + let prefix = posting_block_score_prefix_len(block_size); match codec { - PostingTailCodec::Fixed32 => u32::from_le_bytes(block[4..8].try_into().unwrap()), + PostingTailCodec::Fixed32 => { + u32::from_le_bytes(block[prefix..prefix + 4].try_into().unwrap()) + } PostingTailCodec::VarintDelta => { - let mut offset = 4usize; + let mut offset = prefix; decode_varint_u32(block, &mut offset) .expect("posting tail block should contain a valid first doc id") } @@ -1139,9 +1172,11 @@ mod tests { assert_eq!(posting_list.len(), 1); let block = posting_list.value(0); - let doc_num_bits = block[8]; + // 256-doc blocks carry no block-max-score prefix (impacts supply the + // per-block bound): [first_doc u32][doc num_bits u8][doc payload]... + let doc_num_bits = block[4]; let doc_bytes = BitPacker8x::compressed_block_size(doc_num_bits); - let freq_header_offset = 9 + doc_bytes; + let freq_header_offset = 5 + doc_bytes; // 256-doc blocks use patched FOR for frequencies: // [width u8][exception_count u8][body][exceptions...] let freq_num_bits = block[freq_header_offset]; diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index f1bfffa9fb4..b63f3ecaec7 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -52,7 +52,7 @@ use std::sync::LazyLock; use tokio::{sync::OnceCell, task::spawn_blocking}; use tracing::{info, instrument}; -use super::encoding::{PositionBlockBuilder, decode_group_starts}; +use super::encoding::{MAX_POSTING_BLOCK_SIZE, PositionBlockBuilder, decode_group_starts}; use super::iter::PostingListIterator; use super::lazy_docset::LazyDocSet; use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; @@ -1559,6 +1559,8 @@ impl InvertedPartition { num_docs, false, frag_reuse_index, + // V3 (256-doc block) partitions score with quantized doc lengths. + inverted_list.block_size() == MAX_POSTING_BLOCK_SIZE, )); Ok(Self { @@ -3990,6 +3992,13 @@ impl CompressedPostingList { } pub fn block_max_score(&self, block_idx: usize) -> f32 { + // 256-doc (V3) blocks store no per-block max score: their impact + // skip data supplies the tight per-block bound, so callers on that + // path never reach here. Fall back to the list-level max, which is + // still a valid (looser) bound for any block. + if super::encoding::posting_block_score_prefix_len(self.block_size) == 0 { + return self.max_score; + } let block = self.blocks.value(block_idx); block[0..4].try_into().map(f32::from_le_bytes).unwrap() } @@ -4012,9 +4021,14 @@ impl CompressedPostingList { return super::encoding::read_posting_tail_first_doc( block, self.posting_tail_codec, + self.block_size, ); } - block[4..8].try_into().map(u32::from_le_bytes).unwrap() + let prefix = super::encoding::posting_block_score_prefix_len(self.block_size); + block[prefix..prefix + 4] + .try_into() + .map(u32::from_le_bytes) + .unwrap() }) .collect() }) @@ -4068,12 +4082,14 @@ impl EncodedBlocks { doc_ids: &[u32], frequencies: &[u32], codec: PostingTailCodec, + block_size: usize, ) -> Result<()> { self.offsets.push(self.bytes.len() as u32); super::encoding::encode_remainder_posting_block_into( doc_ids, frequencies, codec, + block_size, &mut self.bytes, ) } @@ -4863,7 +4879,9 @@ impl PostingListBuilder { frequencies.iter().copied(), ); max_score = max_score.max(block_score); - encoded_blocks.set_block_score(index, block_score); + if super::encoding::posting_block_score_prefix_len(block_size) > 0 { + encoded_blocks.set_block_score(index, block_score); + } } if !tail_entries.is_empty() { @@ -4880,8 +4898,11 @@ impl PostingListBuilder { doc_ids.as_slice(), frequencies.as_slice(), posting_tail_codec, + block_size, )?; - encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); + if super::encoding::posting_block_score_prefix_len(block_size) > 0 { + encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); + } if with_positions { encoded_position_blocks.push_encoded_block( tail_position_block @@ -4898,15 +4919,18 @@ impl PostingListBuilder { )) } + #[allow(clippy::too_many_arguments)] fn build_compressed_with_block_scores_from_parts( with_positions: bool, posting_tail_codec: PostingTailCodec, + block_size: usize, mut encoded_blocks: EncodedBlocks, mut encoded_position_blocks: EncodedPositionBlocks, tail_entries: &[RawDocInfo], tail_position_block: Option>, mut block_max_scores: impl Iterator, ) -> Result<(LargeBinaryArray, Option, f32)> { + let has_score_prefix = super::encoding::posting_block_score_prefix_len(block_size) > 0; let mut max_score = f32::MIN; let mut doc_ids = Vec::with_capacity(BLOCK_SIZE); let mut frequencies = Vec::with_capacity(BLOCK_SIZE); @@ -4916,7 +4940,9 @@ impl PostingListBuilder { .next() .ok_or_else(|| Error::index("missing block max score".to_owned()))?; max_score = max_score.max(block_score); - encoded_blocks.set_block_score(index, block_score); + if has_score_prefix { + encoded_blocks.set_block_score(index, block_score); + } } if !tail_entries.is_empty() { @@ -4929,8 +4955,11 @@ impl PostingListBuilder { doc_ids.as_slice(), frequencies.as_slice(), posting_tail_codec, + block_size, )?; - encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); + if has_score_prefix { + encoded_blocks.set_block_score(encoded_blocks.len() - 1, block_score); + } if with_positions { encoded_position_blocks.push_encoded_block( tail_position_block @@ -4984,6 +5013,7 @@ impl PostingListBuilder { Self::build_compressed_with_block_scores_from_parts( with_positions, posting_tail_codec, + block_size, encoded_blocks .map(|encoded_blocks| *encoded_blocks) .unwrap_or_default(), @@ -5225,9 +5255,55 @@ impl Ord for RawDocInfo { } } +/// Lucene SmallFloat-style doc-length quantization for V3 (256-doc block) +/// scoring: a 4-mantissa-bit float-like byte code. Values 0-7 are exact; +/// larger values keep their top four significand bits (relative error +/// <= 6.25%) and decode to their bucket floor. The floor only ever shortens +/// a doc, and a shorter doc only raises its BM25 weight, so bounds baked +/// from quantized lengths stay valid upper bounds of quantized scores. +pub(super) fn quantize_doc_length(value: u32) -> u8 { + let num_bits = 32 - value.leading_zeros(); + if num_bits < 4 { + value as u8 + } else { + let shift = num_bits - 4; + (((value >> shift) as u8) & 0x07) | (((shift + 1) as u8) << 3) + } +} + +#[inline] +pub(super) fn dequantize_doc_length(code: u8) -> u32 { + DEQUANTIZED_DOC_LENGTHS[code as usize] +} + +pub(super) static DEQUANTIZED_DOC_LENGTHS: [u32; 256] = build_dequantized_doc_lengths(); + +const fn build_dequantized_doc_lengths() -> [u32; 256] { + let mut table = [0u32; 256]; + let mut code = 0usize; + while code < 256 { + let bits = (code & 0x07) as u64; + let shift = (code >> 3) as i64 - 1; + let decoded = if shift < 0 { + bits + } else { + (bits | 0x08) << shift + }; + // Codes past the largest u32 encoding are never produced; saturate so + // the table stays total. + table[code] = if decoded > u32::MAX as u64 { + u32::MAX + } else { + decoded as u32 + }; + code += 1; + } + table +} + // DocSet is a mapping from row ids to the number of tokens in the document // It's used to sort the documents by the bm25 score -#[derive(Debug, Clone, Default, DeepSizeOf)] +#[derive(Debug, Clone, Default)] pub struct DocSet { row_ids: Vec, num_tokens: Vec, @@ -5235,6 +5311,26 @@ pub struct DocSet { inv: Vec<(u64, u32)>, total_tokens: u64, + + // V3 (256-doc block) partitions score with quantized doc lengths: the + // flag is set at partition load and the byte-norm slab bakes lazily on + // first scoring use (shared by clones of the loaded set). 128-block + // partitions never set the flag and keep exact scoring. + scoring_quantized: bool, + norms: Arc>>, +} + +impl DeepSizeOf for DocSet { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.row_ids.deep_size_of_children(context) + + self.num_tokens.deep_size_of_children(context) + + self.inv.deep_size_of_children(context) + + self + .norms + .get() + .map(|slab| std::mem::size_of_val(slab.as_ref())) + .unwrap_or(0) + } } impl DocSet { @@ -5383,6 +5479,8 @@ impl DocSet { num_tokens, inv: Vec::new(), total_tokens, + scoring_quantized: false, + norms: Arc::new(std::sync::OnceLock::new()), } } @@ -5418,6 +5516,8 @@ impl DocSet { num_tokens, inv: Vec::new(), total_tokens, + scoring_quantized: false, + norms: Arc::new(std::sync::OnceLock::new()), }); } @@ -5459,6 +5559,8 @@ impl DocSet { num_tokens, inv, total_tokens, + scoring_quantized: false, + norms: Arc::new(std::sync::OnceLock::new()), }); } @@ -5478,6 +5580,8 @@ impl DocSet { num_tokens, inv, total_tokens, + scoring_quantized: false, + norms: Arc::new(std::sync::OnceLock::new()), }) } @@ -5488,6 +5592,7 @@ impl DocSet { let len = self.len(); let row_ids = std::mem::replace(&mut self.row_ids, Vec::with_capacity(len)); let num_tokens = std::mem::replace(&mut self.num_tokens, Vec::with_capacity(len)); + self.invalidate_norms(); self.total_tokens = 0; for (doc_id, (row_id, num_token)) in std::iter::zip(row_ids, num_tokens).enumerate() { match mapping.get(&row_id) { @@ -5514,6 +5619,39 @@ impl DocSet { self.num_tokens[doc_id as usize] } + /// Enable quantized doc-length scoring (V3 / 256-doc block partitions). + pub fn set_quantized_scoring(&mut self, quantized: bool) { + self.scoring_quantized = quantized; + } + + /// The quantized doc-length slab when this set scores quantized (V3 + /// partitions), baked on first use; `None` for exact-scoring sets. + pub fn scoring_norms(&self) -> Option<&[u8]> { + if !self.scoring_quantized { + return None; + } + Some( + self.norms + .get_or_init(|| { + self.num_tokens + .iter() + .map(|&n| quantize_doc_length(n)) + .collect() + }) + .as_ref(), + ) + } + + /// Doc length as scoring sees it: the quantized bucket floor for V3 + /// partitions, the exact value otherwise. + #[inline] + pub fn scoring_num_tokens(&self, doc_id: u32) -> u32 { + match self.scoring_norms() { + Some(norms) => dequantize_doc_length(norms[doc_id as usize]), + None => self.num_tokens[doc_id as usize], + } + } + // this can be used only if it's a legacy format, // which store the sorted row ids so that we can use binary search #[inline] @@ -5530,9 +5668,18 @@ impl DocSet { self.row_ids.push(row_id); self.num_tokens.push(num_tokens); self.total_tokens += num_tokens as u64; + self.invalidate_norms(); self.row_ids.len() as u32 - 1 } + // Drop the baked norm slab after a mutation; it re-bakes on the next + // scoring use. + fn invalidate_norms(&mut self) { + if self.norms.get().is_some() { + self.norms = Arc::new(std::sync::OnceLock::new()); + } + } + pub(crate) fn memory_size(&self) -> usize { self.row_ids.capacity() * std::mem::size_of::() + self.num_tokens.capacity() * std::mem::size_of::() diff --git a/rust/lance-index/src/scalar/inverted/iter.rs b/rust/lance-index/src/scalar/inverted/iter.rs index 51970c46c33..3755fa41d80 100644 --- a/rust/lance-index/src/scalar/inverted/iter.rs +++ b/rust/lance-index/src/scalar/inverted/iter.rs @@ -166,6 +166,7 @@ impl Iterator for CompressedPostingListIterator { compressed, self.remainder, self.posting_tail_codec, + self.block_size, &mut self.doc_ids, &mut self.frequencies, ); diff --git a/rust/lance-index/src/scalar/inverted/lazy_docset.rs b/rust/lance-index/src/scalar/inverted/lazy_docset.rs index f103c2cefa4..9d424e0eac0 100644 --- a/rust/lance-index/src/scalar/inverted/lazy_docset.rs +++ b/rust/lance-index/src/scalar/inverted/lazy_docset.rs @@ -64,6 +64,9 @@ pub struct DeferredDocSet { docs_path: String, is_legacy: bool, frag_reuse_index: Option>, + /// V3 (256-doc block) partitions score with quantized doc lengths; the + /// flag is applied to every `DocSet` this deferred set materializes. + quantized_scoring: bool, /// Doc count cached at construction so `len()` stays sync + IO-free. num_rows: usize, /// `sum(num_tokens)` cached on first compute. @@ -117,18 +120,21 @@ impl lance_core::deepsize::DeepSizeOf for LazyDocSet { } impl LazyDocSet { + #[allow(clippy::too_many_arguments)] pub fn new( store: Arc, docs_path: String, num_rows: usize, is_legacy: bool, frag_reuse_index: Option>, + quantized_scoring: bool, ) -> Self { Self::Deferred(Box::new(DeferredDocSet { store, docs_path, is_legacy, frag_reuse_index, + quantized_scoring, num_rows, total_tokens: OnceCell::new(), num_tokens_col: OnceCell::new(), @@ -290,7 +296,7 @@ impl DeferredDocSet { .get_or_try_init(|| async { // If the stats path already pulled NUM_TOKEN_COL, // read only ROW_ID and rebuild from the two columns. - let docs = if self.num_tokens_col.get().is_some() { + let mut docs = if self.num_tokens_col.get().is_some() { let num_tokens = self.num_tokens_column().await?; let row_ids = self.row_ids_column().await?; DocSet::from_columns( @@ -307,6 +313,7 @@ impl DeferredDocSet { ) .await? }; + docs.set_quantized_scoring(self.quantized_scoring); Result::Ok(Arc::new(docs)) }) .await? @@ -320,7 +327,9 @@ impl DeferredDocSet { return Ok(full.clone()); } let num_tokens = self.num_tokens_column().await?; - let docs = Arc::new(DocSet::from_num_tokens_only(num_tokens.as_ref())); + let mut docs = DocSet::from_num_tokens_only(num_tokens.as_ref()); + docs.set_quantized_scoring(self.quantized_scoring); + let docs = Arc::new(docs); let _ = self.total_tokens.set(docs.total_tokens_num()); Ok(docs) } diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index b94ac61aa0a..9a48a72dfbb 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -106,6 +106,7 @@ impl CompressedState { block, remainder, tail_codec, + block_size, &mut self.doc_ids, &mut self.freqs, ); @@ -915,7 +916,7 @@ impl<'a, S: Scorer> Wand<'a, S> { } let doc_length = match &doc { - DocInfo::Raw(doc) => self.docs.num_tokens(doc.doc_id), + DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), }; @@ -1089,7 +1090,7 @@ impl<'a, S: Scorer> Wand<'a, S> { // score the doc let doc_length = match is_compressed { - true => self.docs.num_tokens(doc_id as u32), + true => self.docs.scoring_num_tokens(doc_id as u32), false => self.docs.num_tokens_by_row_id(row_id), }; if self.operator == Operator::Or && !self.refine_or_candidate(doc_id, doc_length) { @@ -1239,7 +1240,7 @@ impl<'a, S: Scorer> Wand<'a, S> { continue; }; let doc_length = match &first_doc { - DocInfo::Raw(doc) => self.docs.num_tokens(doc.doc_id), + DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), }; let mut lead_score = 0.0; @@ -1321,7 +1322,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let lead_doc = self.lead.first().and_then(|posting| posting.doc())?; let doc_length = match &lead_doc { - DocInfo::Raw(doc) => self.docs.num_tokens(doc.doc_id), + DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), }; if self.and_candidate_cannot_beat_threshold(doc_length) { From 96b172d1edd1a15e9b3374e4d1cc688b93f0275c Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 3 Jul 2026 13:10:54 +0800 Subject: [PATCH 2/7] feat(fts): impact skip data for posting lists Store per-block (freq, doc_len) impact frontiers alongside 256-doc posting blocks (varint-encoded: 2-3 bytes per pair) plus one level1 entry per 32 blocks, and drive block-max WAND pruning from them instead of build-time scores that go stale as index stats drift: - Bounds bake once per cached list into an Arc-shared slab (max doc weight per entry plus the list-wide max); per-query clones reuse it, so query time pays one multiply per bound instead of frontier rescans. - Entry doc_up_tos decode once at construction. - Lagging iterators park in the WAND tail under the data-driven global bound (query_weight x baked list max) instead of INFINITY. 128-doc-block indexes keep fixed-width u32 impact entries. --- rust/lance-index/protos-cache/cache.proto | 2 + rust/lance-index/src/scalar/inverted.rs | 1 + .../src/scalar/inverted/builder.rs | 49 +- .../src/scalar/inverted/cache_codec.rs | 203 ++++- .../lance-index/src/scalar/inverted/impact.rs | 765 ++++++++++++++++++ rust/lance-index/src/scalar/inverted/index.rs | 568 +++++++++++-- .../lance-index/src/scalar/inverted/scorer.rs | 11 + rust/lance-index/src/scalar/inverted/wand.rs | 508 ++++++++++-- 8 files changed, 1977 insertions(+), 130 deletions(-) create mode 100644 rust/lance-index/src/scalar/inverted/impact.rs diff --git a/rust/lance-index/protos-cache/cache.proto b/rust/lance-index/protos-cache/cache.proto index e92da05815f..a861e2ef652 100644 --- a/rust/lance-index/protos-cache/cache.proto +++ b/rust/lance-index/protos-cache/cache.proto @@ -31,6 +31,8 @@ message CompressedPostingHeader { // Number of documents in each compressed posting block. Older cache entries // omit this field and decode as the legacy 128-doc block size. uint32 block_size = 6; + // Whether an impact IPC section follows the posting/position sections. + bool has_impacts = 7; } // Header for a serialized `PlainPostingList` cache entry. Followed by an Arrow diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 55e2b2abce5..0ba7e52e530 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors pub mod builder; +mod impact; mod cache_codec; mod encoding; mod index; diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 3301c1b1277..ddd2ed9ef82 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -1733,11 +1733,27 @@ pub fn inverted_list_schema_for_version_with_block_size( with_position: bool, format_version: InvertedListFormatVersion, block_size: usize, +) -> SchemaRef { + inverted_list_schema_for_version_with_block_size_and_impacts( + with_position, + format_version, + block_size, + true, + ) +} + +pub(crate) fn inverted_list_schema_for_version_with_block_size_and_impacts( + with_position: bool, + format_version: InvertedListFormatVersion, + block_size: usize, + with_impacts: bool, ) -> SchemaRef { validate_format_version_block_size(format_version, block_size) .expect("invalid FTS format version for posting block size"); match format_version { - InvertedListFormatVersion::V1 => inverted_list_schema_v1(with_position, block_size), + InvertedListFormatVersion::V1 => { + inverted_list_schema_v1(with_position, block_size, with_impacts) + } InvertedListFormatVersion::V2 | InvertedListFormatVersion::V3 => { inverted_list_schema_with_tail_codec_and_position_codec( with_position, @@ -1745,12 +1761,17 @@ pub fn inverted_list_schema_for_version_with_block_size( PostingTailCodec::VarintDelta, Some(PositionStreamCodec::PackedDelta), block_size, + with_impacts, ) } } } -fn inverted_list_schema_v1(with_position: bool, block_size: usize) -> SchemaRef { +fn inverted_list_schema_v1( + with_position: bool, + block_size: usize, + with_impacts: bool, +) -> SchemaRef { let mut fields = vec![ arrow_schema::Field::new( POSTING_COL, @@ -1764,6 +1785,17 @@ fn inverted_list_schema_v1(with_position: bool, block_size: usize) -> SchemaRef arrow_schema::Field::new(MAX_SCORE_COL, datatypes::DataType::Float32, false), arrow_schema::Field::new(LENGTH_COL, datatypes::DataType::UInt32, false), ]; + if with_impacts { + fields.push(arrow_schema::Field::new( + IMPACT_COL, + datatypes::DataType::List(Arc::new(Field::new( + "item", + datatypes::DataType::LargeBinary, + true, + ))), + false, + )); + } if with_position { fields.push(arrow_schema::Field::new( POSITION_COL, @@ -1806,6 +1838,7 @@ pub fn inverted_list_schema_with_tail_codec( posting_tail_codec, Some(PositionStreamCodec::PackedDelta), LEGACY_BLOCK_SIZE, + false, ) } @@ -1815,6 +1848,7 @@ fn inverted_list_schema_with_tail_codec_and_position_codec( posting_tail_codec: PostingTailCodec, position_codec: Option, block_size: usize, + with_impacts: bool, ) -> SchemaRef { let mut fields = vec![ // we compress the posting lists (including row ids and frequencies), @@ -1831,6 +1865,17 @@ fn inverted_list_schema_with_tail_codec_and_position_codec( arrow_schema::Field::new(MAX_SCORE_COL, datatypes::DataType::Float32, false), arrow_schema::Field::new(LENGTH_COL, datatypes::DataType::UInt32, false), ]; + if with_impacts { + fields.push(arrow_schema::Field::new( + IMPACT_COL, + datatypes::DataType::List(Arc::new(Field::new( + "item", + datatypes::DataType::LargeBinary, + true, + ))), + false, + )); + } if with_position { fields.push(arrow_schema::Field::new( COMPRESSED_POSITION_COL, diff --git a/rust/lance-index/src/scalar/inverted/cache_codec.rs b/rust/lance-index/src/scalar/inverted/cache_codec.rs index fffb080bda1..bf36209d449 100644 --- a/rust/lance-index/src/scalar/inverted/cache_codec.rs +++ b/rust/lance-index/src/scalar/inverted/cache_codec.rs @@ -39,6 +39,7 @@ use crate::cache_pb::{ PostingTailCodec as PbPostingTailCodec, }; +use super::impact::ImpactSkipData; use super::index::{ CompressedPositionStorage, CompressedPostingList, PlainPostingList, PositionStreamCodec, Positions, PostingList, PostingListGroup, PostingTailCodec, SharedPositionStream, @@ -96,6 +97,7 @@ const BLOCK_OFFSETS_COLUMN: &str = "block_offsets"; const ROW_IDS_COLUMN: &str = "row_ids"; const FREQUENCIES_COLUMN: &str = "frequencies"; const BLOCKS_COLUMN: &str = "blocks"; +const IMPACTS_COLUMN: &str = "impacts"; fn legacy_positions_batch(list: &ListArray) -> Result { let schema = Arc::new(Schema::new(vec![Field::new( @@ -179,7 +181,7 @@ fn read_position_sections( impl CacheCodecImpl for PostingList { const TYPE_ID: &'static str = "lance.fts.PostingList"; - const CURRENT_VERSION: u32 = 1; + const CURRENT_VERSION: u32 = 2; fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { match self { @@ -293,6 +295,7 @@ fn serialize_compressed( position_storage: position_storage as i32, position_stream_codec: position_stream_codec as i32, block_size: posting.block_size as u32, + has_impacts: posting.impacts.is_some(), }; w.write_header(&header)?; @@ -307,6 +310,15 @@ fn serialize_compressed( if let Some(storage) = &posting.positions { write_position_sections(w, storage)?; } + if let Some(impacts) = &posting.impacts { + let schema = Arc::new(Schema::new(vec![Field::new( + IMPACTS_COLUMN, + DataType::LargeBinary, + false, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(impacts.entries().clone())])?; + w.write_ipc(&batch)?; + } Ok(()) } @@ -329,6 +341,22 @@ fn deserialize_compressed(r: &mut CacheEntryReader<'_>) -> Result= 2 && header.has_impacts { + let batch = r.read_ipc()?; + let entries = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::io("impacts column is not a LargeBinaryArray".to_string()))? + .clone(); + Some(ImpactSkipData::new( + entries, + blocks.len(), + super::impact::ImpactFormat::for_block_size(block_size), + )?) + } else { + None + }; Ok(CompressedPostingList::new( blocks, @@ -337,6 +365,7 @@ fn deserialize_compressed(r: &mut CacheEntryReader<'_>) -> Result) -> Result) -> Result<()> { let count = u32::try_from(self.posting_lists.len()) @@ -418,12 +447,13 @@ impl CacheCodecImpl for Positions { #[cfg(test)] mod tests { use arrow::buffer::ScalarBuffer; - use arrow_array::LargeBinaryArray; - use arrow_array::builder::{Int32Builder, ListBuilder}; + use arrow_array::builder::{Int32Builder, LargeBinaryBuilder, ListBuilder}; + use arrow_array::{Array, LargeBinaryArray}; use bytes::Bytes; use lance_core::Result; use lance_core::cache::{CacheCodecImpl, CacheEntryReader, CacheEntryWriter}; + use super::super::impact::ImpactSkipData; use super::super::index::{ CompressedPositionStorage, CompressedPostingList, PlainPostingList, PositionStreamCodec, Positions, PostingList, PostingListGroup, PostingTailCodec, SharedPositionStream, @@ -469,6 +499,44 @@ mod tests { } } + fn impact_entry(doc_up_to: u32, pairs: &[(u32, u32)]) -> Vec { + let mut bytes = Vec::with_capacity(8 + pairs.len() * 8); + bytes.extend_from_slice(&doc_up_to.to_le_bytes()); + bytes.extend_from_slice(&(pairs.len() as u32).to_le_bytes()); + for (freq, doc_len) in pairs { + bytes.extend_from_slice(&freq.to_le_bytes()); + bytes.extend_from_slice(&doc_len.to_le_bytes()); + } + bytes + } + + fn impact_skip_data(level0_len: usize) -> ImpactSkipData { + let mut entries = Vec::with_capacity(level0_len + level0_len.div_ceil(32)); + for block_idx in 0..level0_len { + entries.push(impact_entry( + (block_idx as u32 + 1) * 10 - 1, + &[(block_idx as u32 + 1, 10)], + )); + } + let num_level1 = level0_len.div_ceil(32); + for group_idx in 0..num_level1 { + let group_end = ((group_idx + 1) * 32).min(level0_len); + let doc_up_to = group_end as u32 * 10 - 1; + entries.push(impact_entry(doc_up_to, &[(1, 10), (2, 8)])); + } + let mut builder = LargeBinaryBuilder::with_capacity(entries.len(), 0); + for entry in entries { + builder.append_value(entry); + } + let array = builder.finish(); + ImpactSkipData::new( + array, + level0_len, + super::super::impact::ImpactFormat::FixedU32, + ) + .unwrap() + } + /// Serialize a codec body (no envelope) into a standalone buffer. fn body_bytes(entry: &T) -> Bytes { let mut buf = Vec::new(); @@ -483,6 +551,11 @@ mod tests { T::deserialize(&mut r) } + fn from_body_version(data: &Bytes, version: u32) -> Result { + let mut r = CacheEntryReader::new(data, 0, version); + T::deserialize(&mut r) + } + fn roundtrip_posting_list(entry: &PostingList) -> PostingList { from_body::(&body_bytes(entry)).unwrap() } @@ -540,8 +613,15 @@ mod tests { Some(&[1u8, 2, 3, 4, 5][..]), Some(&[6, 7, 8, 9, 10][..]), ]); - let posting = - CompressedPostingList::new(blocks, 3.5, 42, PostingTailCodec::VarintDelta, 256, None); + let posting = CompressedPostingList::new( + blocks, + 3.5, + 42, + PostingTailCodec::VarintDelta, + 256, + None, + None, + ); let entry = PostingList::Compressed(posting.clone()); match roundtrip_posting_list(&entry) { PostingList::Compressed(restored) => { @@ -556,6 +636,53 @@ mod tests { } } + #[test] + fn compressed_posting_list_impacts_roundtrip() { + let blocks = LargeBinaryArray::from_opt_vec(vec![ + Some(&[1u8, 2, 3, 4, 5][..]), + Some(&[6, 7, 8, 9, 10][..]), + ]); + let impacts = impact_skip_data(blocks.len()); + let posting = CompressedPostingList::new( + blocks, + 3.5, + 42, + PostingTailCodec::VarintDelta, + 256, + None, + Some(impacts.clone()), + ); + let entry = PostingList::Compressed(posting); + match roundtrip_posting_list(&entry) { + PostingList::Compressed(restored) => { + let restored = restored.impacts.expect("impacts should roundtrip"); + assert_eq!(restored.level0_len(), impacts.level0_len()); + assert_eq!(restored.level1_len(), impacts.level1_len()); + assert_eq!(restored.entries(), impacts.entries()); + } + PostingList::Plain(_) => panic!("expected Compressed variant"), + } + } + + #[test] + fn compressed_posting_list_v1_cache_without_impacts_decodes() { + let posting = CompressedPostingList::new( + LargeBinaryArray::from_opt_vec(vec![Some(&[1u8, 2, 3][..])]), + 1.25, + 5, + PostingTailCodec::Fixed32, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, + None, + None, + ); + let data = body_bytes(&PostingList::Compressed(posting)); + let restored = from_body_version::(&data, 1).unwrap(); + let PostingList::Compressed(restored) = restored else { + panic!("expected Compressed variant"); + }; + assert!(restored.impacts.is_none()); + } + #[test] fn compressed_posting_list_legacy_positions_roundtrip() { let blocks = LargeBinaryArray::from_opt_vec(vec![Some(&[1u8, 2, 3][..])]); @@ -568,6 +695,7 @@ mod tests { Some(CompressedPositionStorage::LegacyPerDoc(legacy_positions( &[&[0, 4, 8]], ))), + None, ); let entry = PostingList::Compressed(posting.clone()); match roundtrip_posting_list(&entry) { @@ -601,6 +729,7 @@ mod tests { PostingTailCodec::VarintDelta, 256, Some(CompressedPositionStorage::SharedStream(stream)), + None, ); let entry = PostingList::Compressed(posting.clone()); match roundtrip_posting_list(&entry) { @@ -632,6 +761,7 @@ mod tests { Some(CompressedPositionStorage::SharedStream( expected_stream.clone(), )), + None, ); let serialized = body_bytes(&PostingList::Compressed(posting)); @@ -665,6 +795,7 @@ mod tests { PostingTailCodec::VarintDelta, 256, None, + None, )); for members in [ @@ -689,6 +820,64 @@ mod tests { } } + #[test] + fn posting_list_group_impacted_compressed_members_roundtrip() { + let first = CompressedPostingList::new( + LargeBinaryArray::from_opt_vec(vec![Some(&[1u8, 2, 3][..]), Some(&[4u8, 5, 6][..])]), + 3.0, + 256, + PostingTailCodec::VarintDelta, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, + None, + Some(impact_skip_data(2)), + ); + let second = CompressedPostingList::new( + LargeBinaryArray::from_opt_vec(vec![Some(&[7u8, 8, 9][..])]), + 5.0, + 128, + PostingTailCodec::Fixed32, + 256, + Some(CompressedPositionStorage::SharedStream( + SharedPositionStream::new( + PositionStreamCodec::PackedDelta, + vec![0u32, 12], + Bytes::from(vec![0xABu8; 32]), + ), + )), + Some(impact_skip_data(1)), + ); + let members = vec![ + PostingList::Compressed(first.clone()), + PostingList::Compressed(second.clone()), + ]; + let group = PostingListGroup::new(members); + let restored = from_body::(&body_bytes(&group)).unwrap(); + assert_eq!(restored.posting_lists.len(), 2); + + let expected = [&first, &second]; + for (expected, restored) in expected.iter().zip(restored.posting_lists.iter()) { + let PostingList::Compressed(restored) = restored else { + panic!("expected compressed member"); + }; + assert_eq!(restored.blocks, expected.blocks); + assert_eq!(restored.length, expected.length); + assert_eq!(restored.max_score, expected.max_score); + assert_eq!(restored.posting_tail_codec, expected.posting_tail_codec); + assert_eq!(restored.block_size, expected.block_size); + assert_eq!( + restored.impacts.as_ref().unwrap().entries(), + expected.impacts.as_ref().unwrap().entries() + ); + match (&expected.positions, &restored.positions) { + (Some(expected), Some(restored)) => { + assert_position_storage_eq(expected, restored); + } + (None, None) => {} + _ => panic!("position storage mismatch"), + } + } + } + #[test] fn positions_legacy_roundtrip() { let positions = Positions(CompressedPositionStorage::LegacyPerDoc(legacy_positions( @@ -775,6 +964,7 @@ mod tests { PostingTailCodec::VarintDelta, 256, Some(CompressedPositionStorage::SharedStream(stream)), + None, )) } @@ -826,6 +1016,7 @@ mod tests { PostingTailCodec::VarintDelta, 256, None, + None, )) }; let group = PostingListGroup::new(vec![make_member(9), make_member(1)]); diff --git a/rust/lance-index/src/scalar/inverted/impact.rs b/rust/lance-index/src/scalar/inverted/impact.rs new file mode 100644 index 00000000000..b465d0e7cfe --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/impact.rs @@ -0,0 +1,765 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::{Arc, OnceLock}; + +use arrow_array::builder::LargeBinaryBuilder; +use arrow_array::{Array, LargeBinaryArray}; +use lance_core::{Error, Result}; + +use super::scorer::Scorer; + +pub const IMPACT_LEVEL1_BLOCKS: usize = 32; +const SMALL_FRONTIER_FREQ_LIMIT: usize = 256; + +/// On-disk encoding of one impact entry. +/// +/// `FixedU32` (128-doc blocks): [doc_up_to u32][pair_count u32][(freq u32, +/// doc_len u32)...]. `Varint` (256-doc blocks, format V3): [doc_up_to varint] +/// [pair_count varint][(freq delta varint, doc_len varint)...] with the +/// frontier's freqs delta-encoded in ascending order — pairs shrink from 8 +/// bytes to ~2-3. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ImpactFormat { + FixedU32, + Varint, +} + +impl ImpactFormat { + pub fn for_block_size(block_size: usize) -> Self { + if block_size == 256 { + Self::Varint + } else { + Self::FixedU32 + } + } +} + +#[derive(Debug, Clone)] +pub struct ImpactSkipData { + entries: LargeBinaryArray, + level0_len: usize, + format: ImpactFormat, + // Last doc id covered by each entry (level0 entries then level1 entries), + // decoded once at construction; u32::MAX marks malformed entries. + entry_doc_up_tos: Arc<[u32]>, + // Per-entry max doc weight plus the list-wide max, baked on first use and + // shared across the per-query clones of a cached list. Doc weights depend + // only on (freq, doc_len) and index-wide stats (e.g. BM25 avgdl), not the + // query. Malformed entries bake to INFINITY so pruning stays safe. + doc_weight_bounds: Arc, f32)>>, +} + +impl PartialEq for ImpactSkipData { + fn eq(&self, other: &Self) -> bool { + self.entries == other.entries && self.level0_len == other.level0_len + } +} + +#[derive(Debug, Clone, Copy)] +pub struct ImpactScore { + pub score: f32, + pub entries_scanned: usize, +} + +#[derive(Debug, Default, Clone)] +pub struct ImpactScoreCache {} + +impl ImpactScoreCache { + fn entry_score( + &mut self, + impacts: &ImpactSkipData, + entry_idx: usize, + query_weight: f32, + scorer: &S, + ) -> f32 { + if query_weight <= 0.0 { + return 0.0; + } + query_weight * impacts.doc_weight_bounds(scorer)[entry_idx] + } +} + +#[derive(Debug, Clone, Copy)] +struct ImpactEntryHeader { + doc_up_to: u32, + pair_count: usize, +} + +impl ImpactSkipData { + pub fn new(entries: LargeBinaryArray, level0_len: usize, format: ImpactFormat) -> Result { + let expected_len = level0_len + level1_len(level0_len); + if entries.len() != expected_len { + return Err(Error::index(format!( + "impact entry count mismatch: got {}, expected {} for {} level0 blocks", + entries.len(), + expected_len, + level0_len + ))); + } + let entry_doc_up_tos = (0..entries.len()) + .map(|entry_idx| { + if entries.is_null(entry_idx) { + return u32::MAX; + } + decode_entry_doc_up_to(entries.value(entry_idx), format).unwrap_or(u32::MAX) + }) + .collect::>(); + Ok(Self { + entries, + level0_len, + format, + entry_doc_up_tos, + doc_weight_bounds: Arc::new(OnceLock::new()), + }) + } + + fn baked_bounds(&self, scorer: &S) -> &(Box<[f32]>, f32) { + self.doc_weight_bounds.get_or_init(|| { + // V3 (Varint impacts <=> 256-doc blocks) scores with quantized doc + // lengths, so bake the bounds against the same quantized lengths. + // Quantization is monotone, so pareto dominance among the stored + // (freq, doc_len) frontier pairs is preserved and the frontier max + // still bounds every doc in the range. + let quantized = self.format == ImpactFormat::Varint; + let per_entry = (0..self.entries.len()) + .map(|entry_idx| { + let bytes = self.entries.value(entry_idx); + let mut max_doc_weight = 0.0_f32; + match for_each_entry_pair(bytes, self.format, |freq, doc_len| { + let doc_len = if quantized { + super::index::dequantize_doc_length(super::index::quantize_doc_length( + doc_len, + )) + } else { + doc_len + }; + max_doc_weight = max_doc_weight.max(scorer.doc_weight(freq, doc_len)); + }) { + Ok(()) => max_doc_weight, + Err(_) => f32::INFINITY, + } + }) + .collect::>(); + // The level1 entries cover every block, so their max is the + // list-wide max doc weight; fall back to level0 entries for lists + // too short to have a level1 entry. + let global = if per_entry.len() > self.level0_len { + per_entry[self.level0_len..] + .iter() + .copied() + .fold(0.0_f32, f32::max) + } else { + per_entry.iter().copied().fold(0.0_f32, f32::max) + }; + (per_entry, global) + }) + } + + fn doc_weight_bounds(&self, scorer: &S) -> &[f32] { + &self.baked_bounds(scorer).0 + } + + /// List-wide max doc weight, from the baked bounds. The tightest valid + /// global score bound for this list is `query_weight * this`, matching + /// what the non-impact format stores as `max_score` at build time — but + /// computed against the current index stats. + pub fn global_max_doc_weight(&self, scorer: &S) -> f32 { + self.baked_bounds(scorer).1 + } + + pub fn entries(&self) -> &LargeBinaryArray { + &self.entries + } + + #[cfg(test)] + pub fn level0_len(&self) -> usize { + self.level0_len + } + + #[cfg(test)] + pub fn level1_len(&self) -> usize { + level1_len(self.level0_len) + } + + pub(crate) fn level1_doc_up_to(&self, group_idx: usize) -> Option { + if group_idx >= level1_len(self.level0_len) { + return None; + } + match self.entry_doc_up_tos[self.level0_len + group_idx] { + u32::MAX => None, + doc_up_to => Some(doc_up_to), + } + } + + /// Max score of the docs covered by the level0 entry of `block_idx`, + /// answered from the baked bounds slab. + // Only tests exercise the uncached form until the maxscore rework + // (stacked follow-up) anchors its block-max caches on it. + #[cfg_attr(not(test), allow(dead_code))] + pub fn level0_score( + &self, + block_idx: usize, + query_weight: f32, + scorer: &S, + ) -> f32 { + if block_idx >= self.level0_len || query_weight <= 0.0 { + return 0.0; + } + query_weight * self.doc_weight_bounds(scorer)[block_idx] + } + + pub fn level0_score_cached( + &self, + block_idx: usize, + query_weight: f32, + scorer: &S, + cache: &mut ImpactScoreCache, + ) -> f32 { + if block_idx >= self.level0_len { + return 0.0; + } + cache.entry_score(self, block_idx, query_weight, scorer) + } + + #[cfg(test)] + pub fn max_score_up_to( + &self, + start_block_idx: usize, + up_to: u64, + _block_least_doc_id: F, + query_weight: f32, + scorer: &S, + ) -> ImpactScore + where + S: Scorer + ?Sized, + F: FnMut(usize) -> u32, + { + self.max_score_up_to_with(start_block_idx, up_to, |impacts, entry_idx| { + impacts.entry_score(entry_idx, query_weight, scorer) + }) + } + + pub fn max_score_up_to_cached( + &self, + start_block_idx: usize, + up_to: u64, + query_weight: f32, + scorer: &S, + cache: &mut ImpactScoreCache, + ) -> ImpactScore + where + S: Scorer + ?Sized, + { + self.max_score_up_to_with(start_block_idx, up_to, |impacts, entry_idx| { + cache.entry_score(impacts, entry_idx, query_weight, scorer) + }) + } + + fn max_score_up_to_with( + &self, + start_block_idx: usize, + up_to: u64, + mut entry_score: E, + ) -> ImpactScore + where + E: FnMut(&Self, usize) -> f32, + { + let mut block_idx = start_block_idx; + let mut max_score = 0.0_f32; + let mut entries_scanned = 0usize; + + while block_idx < self.level0_len { + let group_idx = block_idx / IMPACT_LEVEL1_BLOCKS; + let group_start = group_idx * IMPACT_LEVEL1_BLOCKS; + let group_end = ((group_idx + 1) * IMPACT_LEVEL1_BLOCKS).min(self.level0_len); + if block_idx == group_start { + let level1_entry_idx = self.level0_len + group_idx; + match self.entry_doc_up_tos[level1_entry_idx] { + u32::MAX => { + return ImpactScore { + score: f32::INFINITY, + entries_scanned: entries_scanned + 1, + }; + } + doc_up_to if u64::from(doc_up_to) <= up_to => { + max_score = max_score.max(entry_score(self, level1_entry_idx)); + entries_scanned += 1; + block_idx = group_end; + continue; + } + _ => {} + } + } + + max_score = max_score.max(entry_score(self, block_idx)); + entries_scanned += 1; + match self.entry_doc_up_tos[block_idx] { + u32::MAX => { + return ImpactScore { + score: f32::INFINITY, + entries_scanned, + }; + } + doc_up_to if u64::from(doc_up_to) >= up_to => break, + _ => {} + } + block_idx += 1; + } + + ImpactScore { + score: max_score, + entries_scanned, + } + } + + #[cfg(test)] + fn entry_score( + &self, + entry_idx: usize, + query_weight: f32, + scorer: &S, + ) -> f32 { + if query_weight <= 0.0 { + return 0.0; + } + let bytes = self.entries.value(entry_idx); + let mut max_doc_weight = 0.0_f32; + if for_each_entry_pair(bytes, self.format, |freq, doc_len| { + max_doc_weight = max_doc_weight.max(scorer.doc_weight(freq, doc_len)); + }) + .is_err() + { + return f32::INFINITY; + } + query_weight * max_doc_weight + } +} + +pub struct ImpactSkipDataBuilder { + entries: LargeBinaryBuilder, + level0_len: usize, + level1_entries: Vec>, + level1_docs: Vec<(u32, u32, u32)>, + format: ImpactFormat, +} + +impl ImpactSkipDataBuilder { + pub fn with_capacity(level0_blocks: usize, block_size: usize) -> Self { + Self { + entries: LargeBinaryBuilder::with_capacity( + level0_blocks + level1_len(level0_blocks), + 0, + ), + level0_len: 0, + level1_entries: Vec::with_capacity(level1_len(level0_blocks)), + level1_docs: Vec::with_capacity(IMPACT_LEVEL1_BLOCKS * block_size), + format: ImpactFormat::for_block_size(block_size), + } + } + + pub fn append_block(&mut self, docs: &[(u32, u32, u32)]) -> Result<()> { + let bytes = encode_impact_entry(docs, self.format)?; + self.entries.append_value(bytes.as_slice()); + self.level0_len += 1; + self.level1_docs.extend_from_slice(docs); + if self.level0_len.is_multiple_of(IMPACT_LEVEL1_BLOCKS) { + self.flush_level1()?; + } + Ok(()) + } + + pub fn finish(mut self) -> Result { + if !self.level1_docs.is_empty() { + self.flush_level1()?; + } + for entry in self.level1_entries { + self.entries.append_value(entry.as_slice()); + } + ImpactSkipData::new(self.entries.finish(), self.level0_len, self.format) + } + + fn flush_level1(&mut self) -> Result<()> { + let bytes = encode_impact_entry(self.level1_docs.as_slice(), self.format)?; + self.level1_entries.push(bytes); + self.level1_docs.clear(); + Ok(()) + } +} + +#[cfg(test)] +pub fn build_impact_skip_data(blocks: &[Vec<(u32, u32, u32)>]) -> Result { + let block_size = blocks.iter().map(Vec::len).max().unwrap_or(0).max(1); + let mut builder = ImpactSkipDataBuilder::with_capacity(blocks.len(), block_size); + for block in blocks { + builder.append_block(block)?; + } + builder.finish() +} + +fn encode_impact_entry(docs: &[(u32, u32, u32)], format: ImpactFormat) -> Result> { + let doc_up_to = docs + .last() + .map(|(doc_id, _, _)| *doc_id) + .unwrap_or_default(); + let frontier = impact_frontier(docs); + let pair_count = u32::try_from(frontier.len()).map_err(|_| { + Error::index("impact frontier too large to encode as u32 pair count".to_string()) + })?; + let mut bytes = Vec::with_capacity(8 + frontier.len() * 8); + match format { + ImpactFormat::FixedU32 => { + bytes.extend_from_slice(&doc_up_to.to_le_bytes()); + bytes.extend_from_slice(&pair_count.to_le_bytes()); + for (freq, doc_len) in frontier { + bytes.extend_from_slice(&freq.to_le_bytes()); + bytes.extend_from_slice(&doc_len.to_le_bytes()); + } + } + ImpactFormat::Varint => { + super::encoding::encode_varint_u32(&mut bytes, doc_up_to); + super::encoding::encode_varint_u32(&mut bytes, pair_count); + let mut previous_freq = 0u32; + for (freq, doc_len) in frontier { + super::encoding::encode_varint_u32(&mut bytes, freq - previous_freq); + super::encoding::encode_varint_u32(&mut bytes, doc_len); + previous_freq = freq; + } + } + } + Ok(bytes) +} + +fn decode_entry_doc_up_to(bytes: &[u8], format: ImpactFormat) -> Result { + match format { + ImpactFormat::FixedU32 => decode_header(bytes).map(|header| header.doc_up_to), + ImpactFormat::Varint => { + let mut offset = 0usize; + super::encoding::decode_varint_u32(bytes, &mut offset) + } + } +} + +/// Walk an entry's (freq, doc_len) frontier pairs, validating the layout. +fn for_each_entry_pair( + bytes: &[u8], + format: ImpactFormat, + mut visit: impl FnMut(u32, u32), +) -> Result<()> { + match format { + ImpactFormat::FixedU32 => { + let header = decode_header(bytes)?; + let mut offset = 8usize; + for _ in 0..header.pair_count { + visit(read_u32_le(bytes, offset), read_u32_le(bytes, offset + 4)); + offset += 8; + } + } + ImpactFormat::Varint => { + let mut offset = 0usize; + let _doc_up_to = super::encoding::decode_varint_u32(bytes, &mut offset)?; + let pair_count = super::encoding::decode_varint_u32(bytes, &mut offset)?; + let mut freq = 0u32; + for _ in 0..pair_count { + freq = freq + .checked_add(super::encoding::decode_varint_u32(bytes, &mut offset)?) + .ok_or_else(|| Error::index("impact freq delta overflow".to_owned()))?; + let doc_len = super::encoding::decode_varint_u32(bytes, &mut offset)?; + visit(freq, doc_len); + } + if offset != bytes.len() { + return Err(Error::index(format!( + "impact varint entry has {} trailing bytes", + bytes.len() - offset + ))); + } + } + } + Ok(()) +} + +fn impact_frontier(docs: &[(u32, u32, u32)]) -> Vec<(u32, u32)> { + let max_freq = docs.iter().map(|(_, freq, _)| *freq).max().unwrap_or(0) as usize; + if max_freq <= SMALL_FRONTIER_FREQ_LIMIT { + return impact_frontier_small_freq(docs, max_freq); + } + + impact_frontier_sparse_freq(docs) +} + +fn impact_frontier_small_freq(docs: &[(u32, u32, u32)], max_freq: usize) -> Vec<(u32, u32)> { + let mut min_doc_len_by_freq = [u32::MAX; SMALL_FRONTIER_FREQ_LIMIT + 1]; + for (_, freq, doc_len) in docs { + min_doc_len_by_freq[*freq as usize] = min_doc_len_by_freq[*freq as usize].min(*doc_len); + } + + let min_doc_lens = min_doc_len_by_freq[..=max_freq] + .iter() + .enumerate() + .filter_map(|(freq, doc_len)| (*doc_len != u32::MAX).then_some((freq as u32, *doc_len))) + .collect::>(); + frontier_from_min_doc_lens(min_doc_lens) +} + +fn impact_frontier_sparse_freq(docs: &[(u32, u32, u32)]) -> Vec<(u32, u32)> { + let mut pairs = docs + .iter() + .map(|(_, freq, doc_len)| (*freq, *doc_len)) + .collect::>(); + pairs.sort_unstable_by_key(|(freq, _)| *freq); + + let mut min_doc_lens: Vec<(u32, u32)> = Vec::with_capacity(pairs.len()); + for (freq, doc_len) in pairs { + match min_doc_lens.last_mut() { + Some((last_freq, last_doc_len)) if *last_freq == freq => { + *last_doc_len = (*last_doc_len).min(doc_len); + } + _ => min_doc_lens.push((freq, doc_len)), + } + } + + frontier_from_min_doc_lens(min_doc_lens) +} + +fn frontier_from_min_doc_lens(min_doc_lens: Vec<(u32, u32)>) -> Vec<(u32, u32)> { + let mut best_doc_len = u32::MAX; + let mut frontier = Vec::with_capacity(min_doc_lens.len()); + for (freq, doc_len) in min_doc_lens.into_iter().rev() { + if doc_len < best_doc_len { + frontier.push((freq, doc_len)); + best_doc_len = doc_len; + } + } + frontier.reverse(); + frontier +} + +fn decode_header(bytes: &[u8]) -> Result { + if bytes.len() < 8 { + return Err(Error::index(format!( + "impact entry too short: {} bytes", + bytes.len() + ))); + } + let pair_count = read_u32_le(bytes, 4) as usize; + let expected_len = 8 + pair_count * 8; + if bytes.len() != expected_len { + return Err(Error::index(format!( + "impact entry length mismatch: got {} bytes, expected {} for {} pairs", + bytes.len(), + expected_len, + pair_count + ))); + } + Ok(ImpactEntryHeader { + doc_up_to: read_u32_le(bytes, 0), + pair_count, + }) +} + +#[inline] +fn read_u32_le(bytes: &[u8], offset: usize) -> u32 { + let mut value = [0u8; 4]; + value.copy_from_slice(&bytes[offset..offset + 4]); + u32::from_le_bytes(value) +} + +fn level1_len(level0_len: usize) -> usize { + level0_len.div_ceil(IMPACT_LEVEL1_BLOCKS) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::scalar::inverted::scorer::{MemBM25Scorer, Scorer}; + + #[test] + fn impact_entry_frontier_drops_dominated_pairs() { + let docs = vec![(0, 1, 10), (1, 1, 8), (2, 2, 9), (3, 3, 20)]; + assert_eq!(impact_frontier(&docs), vec![(1, 8), (2, 9), (3, 20)]); + } + + #[test] + fn impact_entry_frontier_handles_sparse_large_frequencies() { + let docs = vec![ + (0, 1, 100), + (1, 1, 80), + (2, 512, 90), + (3, 1_000, 120), + (4, 1_000, 110), + ]; + assert_eq!( + impact_frontier(&docs), + vec![(1, 80), (512, 90), (1_000, 110)] + ); + } + + #[test] + fn impact_max_score_can_use_level1_entry() { + let blocks = (0..40) + .map(|block| vec![(block as u32, 1 + block as u32 % 3, 10)]) + .collect::>(); + let impacts = build_impact_skip_data(&blocks).unwrap(); + assert_eq!(impacts.level0_len(), 40); + assert_eq!(impacts.level1_len(), 2); + let scorer = MemBM25Scorer::new(400, 40, HashMap::from([(String::from("token"), 40usize)])); + let score = impacts.max_score_up_to(0, 31, |idx| idx as u32, 1.0, &scorer); + assert!(score.entries_scanned < IMPACT_LEVEL1_BLOCKS); + assert!(score.score > 0.0); + } + + #[test] + fn impact_level1_doc_up_to_reports_full_and_partial_groups() { + let blocks = (0..40) + .map(|block| vec![(block as u32, 1, 10)]) + .collect::>(); + let impacts = build_impact_skip_data(&blocks).unwrap(); + + assert_eq!( + impacts.level1_doc_up_to(0), + Some((IMPACT_LEVEL1_BLOCKS - 1) as u32) + ); + assert_eq!(impacts.level1_doc_up_to(1), Some(39)); + assert_eq!(impacts.level1_doc_up_to(2), None); + } + + #[test] + fn impact_level1_doc_up_to_returns_none_for_malformed_entry() { + let level0 = encode_impact_entry(&[(0, 1, 10)], ImpactFormat::FixedU32).unwrap(); + let malformed_level1 = vec![1, 2, 3]; + let entries = LargeBinaryArray::from_opt_vec(vec![ + Some(level0.as_slice()), + Some(malformed_level1.as_slice()), + ]); + let impacts = ImpactSkipData::new(entries, 1, ImpactFormat::FixedU32).unwrap(); + + assert_eq!(impacts.level1_doc_up_to(0), None); + } + + #[test] + fn impact_score_cache_matches_uncached_scores() { + let blocks = (0..40) + .map(|block| vec![(block as u32, 1 + block as u32 % 3, 10)]) + .collect::>(); + let impacts = build_impact_skip_data(&blocks).unwrap(); + let scorer = MemBM25Scorer::new(400, 40, HashMap::from([(String::from("token"), 40usize)])); + let mut cache = ImpactScoreCache::default(); + + let uncached_level0 = impacts.level0_score(3, 1.0, &scorer); + let cached_level0 = impacts.level0_score_cached(3, 1.0, &scorer, &mut cache); + assert_eq!(cached_level0, uncached_level0); + + let uncached = impacts.max_score_up_to(0, 31, |idx| idx as u32, 1.0, &scorer); + let cached = impacts.max_score_up_to_cached(0, 31, 1.0, &scorer, &mut cache); + assert_eq!(cached.score, uncached.score); + assert_eq!(cached.entries_scanned, uncached.entries_scanned); + } + + #[test] + fn impact_entries_are_decoded_lazily() { + let level0_0 = encode_impact_entry(&[(0, 1, 10)], ImpactFormat::FixedU32).unwrap(); + let malformed_level0_1 = vec![1, 2, 3]; + let level1 = + encode_impact_entry(&[(0, 1, 10), (1, 1, 10)], ImpactFormat::FixedU32).unwrap(); + let entries = LargeBinaryArray::from_opt_vec(vec![ + Some(level0_0.as_slice()), + Some(malformed_level0_1.as_slice()), + Some(level1.as_slice()), + ]); + let impacts = ImpactSkipData::new(entries, 2, ImpactFormat::FixedU32).unwrap(); + let scorer = MemBM25Scorer::new(10, 10, HashMap::from([(String::from("token"), 2usize)])); + + let score = impacts.max_score_up_to(0, 0, |_| 0, 1.0, &scorer); + assert!(score.score.is_finite()); + assert_eq!(score.entries_scanned, 1); + + let mut cache = ImpactScoreCache::default(); + assert_eq!( + impacts.level0_score_cached(1, 1.0, &scorer, &mut cache), + f32::INFINITY + ); + } + + #[test] + fn impact_varint_entries_roundtrip_and_match_fixed() { + let docs = vec![(3, 1, 100), (9, 2, 40), (200, 7, 80), (4095, 130, 900)]; + let fixed = encode_impact_entry(&docs, ImpactFormat::FixedU32).unwrap(); + let varint = encode_impact_entry(&docs, ImpactFormat::Varint).unwrap(); + assert!(varint.len() < fixed.len()); + assert_eq!( + decode_entry_doc_up_to(&fixed, ImpactFormat::FixedU32).unwrap(), + decode_entry_doc_up_to(&varint, ImpactFormat::Varint).unwrap() + ); + let mut fixed_pairs = Vec::new(); + for_each_entry_pair(&fixed, ImpactFormat::FixedU32, |f, l| { + fixed_pairs.push((f, l)) + }) + .unwrap(); + let mut varint_pairs = Vec::new(); + for_each_entry_pair(&varint, ImpactFormat::Varint, |f, l| { + varint_pairs.push((f, l)) + }) + .unwrap(); + assert_eq!(fixed_pairs, varint_pairs); + assert!(!fixed_pairs.is_empty()); + + // a 256-doc-block skip data goes through the varint path end to end + let blocks: Vec> = (0..3) + .map(|b| (0..256).map(|i| (b * 256 + i, 1 + i % 5, 10)).collect()) + .collect(); + let impacts = build_impact_skip_data(&blocks).unwrap(); + assert_eq!(impacts.level1_doc_up_to(0), Some(767)); + let scorer = MemBM25Scorer::new(400, 768, HashMap::from([(String::from("t"), 768usize)])); + assert!(impacts.level0_score(0, 1.0, &scorer).is_finite()); + let level1 = impacts.max_score_up_to(0, 767, |idx| (idx * 256) as u32, 1.0, &scorer); + assert!(level1.score.is_finite() && level1.score > 0.0); + } + + #[test] + fn impact_upper_bound_covers_real_scores() { + let blocks = vec![ + vec![(0, 1, 100), (3, 2, 40), (7, 4, 80)], + vec![(9, 3, 15), (10, 1, 5), (12, 5, 30)], + vec![(16, 2, 10), (18, 6, 70), (21, 3, 12)], + vec![(24, 1, 4), (28, 7, 100), (30, 2, 8)], + ]; + let impacts = build_impact_skip_data(&blocks).unwrap(); + let scorer = MemBM25Scorer::new(474, 31, HashMap::from([(String::from("token"), 4usize)])); + let query_weight = scorer.query_weight("token"); + + for start_block_idx in 0..blocks.len() { + let up_to = blocks + .iter() + .skip(start_block_idx) + .take(2) + .flatten() + .map(|(doc_id, _, _)| *doc_id) + .max() + .unwrap(); + let upper_bound = impacts.max_score_up_to( + start_block_idx, + u64::from(up_to), + |idx| blocks[idx][0].0, + query_weight, + &scorer, + ); + let exact_max = blocks + .iter() + .skip(start_block_idx) + .flatten() + .take_while(|(doc_id, _, _)| *doc_id <= up_to) + .map(|(_, freq, doc_len)| query_weight * scorer.doc_weight(*freq, *doc_len)) + .fold(0.0_f32, f32::max); + assert!( + upper_bound.score + 1e-6 >= exact_max, + "upper bound {} should cover exact max {} from block {} up to doc {}", + upper_bound.score, + exact_max, + start_block_idx, + up_to + ); + } + } +} diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index b63f3ecaec7..7b2d3357312 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -53,6 +53,7 @@ use tokio::{sync::OnceCell, task::spawn_blocking}; use tracing::{info, instrument}; use super::encoding::{MAX_POSTING_BLOCK_SIZE, PositionBlockBuilder, decode_group_starts}; +use super::impact::{ImpactSkipData, ImpactSkipDataBuilder}; use super::iter::PostingListIterator; use super::lazy_docset::LazyDocSet; use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; @@ -60,7 +61,8 @@ use super::{InvertedIndexBuilder, InvertedIndexParams, wand::*}; use super::{ builder::{ BLOCK_SIZE, PostingGroupAccumulator, PostingGroupConfig, ScoredDoc, doc_file_path, - inverted_list_schema_for_version_with_block_size, posting_file_path, token_file_path, + inverted_list_schema_for_version_with_block_size_and_impacts, posting_file_path, + token_file_path, }, iter::PlainPostingListIterator, query::*, @@ -105,6 +107,7 @@ pub const POSITION_COL: &str = "_position"; pub const COMPRESSED_POSITION_COL: &str = "_compressed_position"; pub const POSITION_BLOCK_OFFSET_COL: &str = "_position_block_offset"; pub const POSTING_COL: &str = "_posting"; +pub const IMPACT_COL: &str = "_impacts"; pub const MAX_SCORE_COL: &str = "_max_score"; pub const LENGTH_COL: &str = "_length"; pub const BLOCK_MAX_SCORE_COL: &str = "_block_max_score"; @@ -291,6 +294,7 @@ impl PartitionCandidates { struct LoadedPostings { postings: Vec, grouped_expansions: Vec, + impact_safe: bool, } impl LoadedPostings { @@ -298,6 +302,7 @@ impl LoadedPostings { Self { postings: Vec::new(), grouped_expansions: Vec::new(), + impact_safe: false, } } } @@ -854,7 +859,7 @@ impl InvertedIndex { // hits per-token `posting_len`; building a `MemBM25Scorer` with // precomputed per-term IDFs avoids the v2 bulk metadata pull. let local_scorer; - let scorer: &dyn Scorer = if let Some(base_scorer) = base_scorer { + let scorer: &MemBM25Scorer = if let Some(base_scorer) = base_scorer { base_scorer } else { local_scorer = self @@ -862,6 +867,7 @@ impl InvertedIndex { .await?; &local_scorer }; + let impact_scorer = Arc::new(scorer.clone()); let limit = params.limit.unwrap_or(usize::MAX); if limit == 0 { @@ -900,8 +906,9 @@ impl InvertedIndex { // 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 - // global k-th — see `Wand::shared_threshold`). - let shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + // global k-th - see `Wand::shared_threshold`). + let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + let legacy_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); let parts = self .partitions .iter() @@ -911,19 +918,23 @@ impl InvertedIndex { let params = params.clone(); let mask = mask.clone(); let metrics = metrics.clone(); - let shared_threshold = shared_threshold.clone(); + let impact_scorer = impact_scorer.clone(); + let impact_shared_threshold = impact_shared_threshold.clone(); + let legacy_shared_threshold = legacy_shared_threshold.clone(); async move { let loaded_postings = part .load_posting_lists( tokens.as_ref(), params.as_ref(), operator, + Some(impact_scorer.as_ref()), metrics.as_ref(), ) .await?; let LoadedPostings { postings, grouped_expansions, + impact_safe, } = loaded_postings; if postings.is_empty() { // No hits in this partition; its DocSet stays @@ -947,6 +958,7 @@ impl InvertedIndex { let metrics = metrics.clone(); let part_for_wand = part.clone(); let has_grouped_expansions = !grouped_expansions.is_empty(); + let use_impact_path = impact_safe && !has_grouped_expansions; let wand_params = if has_grouped_expansions { let mut rescoring_params = params.as_ref().clone(); rescoring_params.limit = @@ -957,9 +969,12 @@ impl InvertedIndex { }; let partition_threshold = if has_grouped_expansions { Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())) + } else if use_impact_path { + impact_shared_threshold } else { - shared_threshold + legacy_shared_threshold }; + let wand_scorer = use_impact_path.then(|| impact_scorer.clone()); let candidates = spawn_cpu(move || { let candidates = part_for_wand.bm25_search( docs_for_wand.as_ref(), @@ -967,6 +982,7 @@ impl InvertedIndex { operator, mask, postings, + wand_scorer, metrics.as_ref(), partition_threshold, )?; @@ -1741,6 +1757,7 @@ impl InvertedPartition { tokens: &Tokens, params: &FtsSearchParams, operator: Operator, + impact_scorer: Option<&MemBM25Scorer>, metrics: &dyn MetricsCollector, ) -> Result { let is_fuzzy = matches!(params.fuzziness, Some(n) if n != 0); @@ -1818,11 +1835,18 @@ impl InvertedPartition { } if !is_fuzzy_and_query { + let impact_safe = impact_scorer.is_some() + && loaded_postings + .iter() + .all(|(_, _, _, posting)| posting.has_impacts()); return Ok(LoadedPostings { postings: loaded_postings .into_iter() .map(|(token_id, token, position, posting)| { - let query_weight = idf(posting.len(), num_docs); + let query_weight = impact_scorer + .filter(|_| impact_safe) + .map(|scorer| scorer.query_weight(&token)) + .unwrap_or_else(|| idf(posting.len(), num_docs)); PostingIterator::with_query_weight( token, token_id, @@ -1834,6 +1858,7 @@ impl InvertedPartition { }) .collect(), grouped_expansions: Vec::new(), + impact_safe, }); } @@ -1901,6 +1926,7 @@ impl InvertedPartition { Ok(LoadedPostings { postings: grouped_postings, grouped_expansions, + impact_safe: false, }) } @@ -1916,6 +1942,7 @@ impl InvertedPartition { operator: Operator, mask: Arc, postings: Vec, + impact_scorer: Option>, metrics: &dyn MetricsCollector, shared_threshold: Arc, ) -> Result> { @@ -1926,10 +1953,16 @@ impl InvertedPartition { // Caller selects the DocSet shape via `LazyDocSet::docs_for_wand` // and passes it in here; wand uses `docs.has_row_ids()` to // handle the num_tokens-only case. - let scorer = IndexBM25Scorer::new(std::iter::once(self)); - let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) - .with_shared_threshold(shared_threshold); - let hits = wand.search(params, mask, metrics)?; + let hits = if let Some(scorer) = impact_scorer { + let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) + .with_shared_threshold(shared_threshold); + wand.search(params, mask, metrics)? + } else { + let scorer = IndexBM25Scorer::new(std::iter::once(self)); + let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) + .with_shared_threshold(shared_threshold); + wand.search(params, mask, metrics)? + }; Ok(hits) } @@ -2348,6 +2381,7 @@ pub struct PostingListReader { metadata: PostingMetadata, has_position: bool, + has_impacts: bool, posting_tail_codec: PostingTailCodec, block_size: usize, positions_layout: PositionsLayout, @@ -2450,6 +2484,7 @@ impl PostingListReader { let posting_tail_codec = parse_posting_tail_codec(&reader.schema().metadata)?; let block_size = parse_posting_block_size(&reader.schema().metadata)?; let has_position = positions_layout != PositionsLayout::None; + let has_impacts = reader.schema().field(IMPACT_COL).is_some(); let metadata = if reader.schema().field(POSTING_COL).is_none() { let (offsets, max_scores) = Self::load_metadata(reader.schema())?; PostingMetadata::LegacyV1 { @@ -2468,6 +2503,7 @@ impl PostingListReader { reader, metadata, has_position, + has_impacts, posting_tail_codec, block_size, positions_layout, @@ -2664,7 +2700,7 @@ impl PostingListReader { self.posting_batch_legacy(token_id, with_position).await } else { let token_id = token_id as usize; - let columns = if with_position { + let mut columns = if with_position { match self.positions_layout { PositionsLayout::SharedStream(_) => { vec![ @@ -2679,6 +2715,9 @@ impl PostingListReader { } else { vec![POSTING_COL] }; + if self.has_impacts { + columns.push(IMPACT_COL); + } let batch = self .reader .read_range(token_id..token_id + 1, Some(&columns)) @@ -2723,11 +2762,18 @@ impl PostingListReader { Some((start, end)) => { let group = self .index_cache - .get_or_insert_with_key(PostingListGroupKey { start, end }, || async move { - metrics.record_part_load(); - info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="inverted", part_id=start); - self.load_posting_list_group(start, end).await - }) + .get_or_insert_with_key( + PostingListGroupKey { + start, + end, + has_impacts: self.has_impacts, + }, + || async move { + metrics.record_part_load(); + info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="inverted", part_id=start); + self.load_posting_list_group(start, end).await + }, + ) .await?; let slot = (token_id - start) as usize; group @@ -2743,19 +2789,25 @@ impl PostingListReader { // per token. None => self .index_cache - .get_or_insert_with_key(PostingListKey { token_id }, || async move { - metrics.record_part_load(); - info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="inverted", part_id=token_id); - // Fetch the posting batch and this token's (max_score, - // length) in parallel; for cold v2 partitions this is one - // single-row metadata read plus one posting-row read, - // instead of pulling the full per-token metadata table. - let (batch, (max_score, length)) = futures::try_join!( - self.posting_batch(token_id, false), - self.posting_metadata_for_token(token_id), - )?; - self.posting_list_from_batch(&batch, max_score, length) - }) + .get_or_insert_with_key( + PostingListKey { + token_id, + has_impacts: self.has_impacts, + }, + || async move { + metrics.record_part_load(); + info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="inverted", part_id=token_id); + // Fetch the posting batch and this token's (max_score, + // length) in parallel; for cold v2 partitions this is one + // single-row metadata read plus one posting-row read, + // instead of pulling the full per-token metadata table. + let (batch, (max_score, length)) = futures::try_join!( + self.posting_batch(token_id, false), + self.posting_metadata_for_token(token_id), + )?; + self.posting_list_from_batch(&batch, max_score, length) + }, + ) .await? .as_ref() .clone(), @@ -2800,12 +2852,13 @@ impl PostingListReader { /// [`PostingListGroup`] cache value (issue #7040). Positions are excluded; /// phrase queries load them on demand via [`Self::read_positions`]. async fn load_posting_list_group(&self, start: u32, end: u32) -> Result { + let mut columns = vec![POSTING_COL, MAX_SCORE_COL, LENGTH_COL]; + if self.has_impacts { + columns.push(IMPACT_COL); + } let batch = self .reader - .read_range( - start as usize..end as usize, - Some(&[POSTING_COL, MAX_SCORE_COL, LENGTH_COL]), - ) + .read_range(start as usize..end as usize, Some(&columns)) .await?; let max_scores = batch[MAX_SCORE_COL].as_primitive::(); let lengths = batch[LENGTH_COL].as_primitive::(); @@ -3128,7 +3181,14 @@ impl PostingListReader { let hi = end as usize - tok_start; let group = PostingListGroup::new(chunk_postings[lo..hi].to_vec()); self.index_cache - .insert_with_key(&PostingListGroupKey { start, end }, Arc::new(group)) + .insert_with_key( + &PostingListGroupKey { + start, + end, + has_impacts: self.has_impacts, + }, + Arc::new(group), + ) .await; } } @@ -3137,7 +3197,13 @@ impl PostingListReader { self.cache_positions(&mut posting_list, token_id, with_position) .await; self.index_cache - .insert_with_key(&PostingListKey { token_id }, Arc::new(posting_list)) + .insert_with_key( + &PostingListKey { + token_id, + has_impacts: self.has_impacts, + }, + Arc::new(posting_list), + ) .await; } } @@ -3307,6 +3373,9 @@ impl PostingListReader { } } } + if self.has_impacts { + base_columns.push(IMPACT_COL); + } base_columns } } @@ -3403,13 +3472,18 @@ impl DeepSizeOf for Positions { #[derive(Debug, Clone)] pub struct PostingListKey { pub token_id: u32, + pub has_impacts: bool, } impl CacheKey for PostingListKey { type ValueType = PostingList; fn key(&self) -> std::borrow::Cow<'_, str> { - format!("postings-{}", self.token_id).into() + if self.has_impacts { + format!("postings-{}-impacts", self.token_id).into() + } else { + format!("postings-{}", self.token_id).into() + } } fn type_name() -> &'static str { @@ -3429,13 +3503,18 @@ impl CacheKey for PostingListKey { pub struct PostingListGroupKey { pub start: u32, pub end: u32, + pub has_impacts: bool, } impl CacheKey for PostingListGroupKey { type ValueType = PostingListGroup; fn key(&self) -> std::borrow::Cow<'_, str> { - format!("postings-{}-{}", self.start, self.end).into() + if self.has_impacts { + format!("postings-{}-{}-impacts", self.start, self.end).into() + } else { + format!("postings-{}-{}", self.start, self.end).into() + } } fn type_name() -> &'static str { @@ -3581,6 +3660,7 @@ impl PostingListGroup { } #[derive(Debug, Clone, DeepSizeOf)] +#[allow(clippy::large_enum_variant)] pub enum PostingList { Plain(PlainPostingList), Compressed(CompressedPostingList), @@ -3645,7 +3725,7 @@ impl PostingList { posting_tail_codec, block_size, shared_position_codec, - ); + )?; Ok(Self::Compressed(posting)) } None => { @@ -3666,6 +3746,13 @@ impl PostingList { } } + pub fn has_impacts(&self) -> bool { + match self { + Self::Plain(_) => false, + Self::Compressed(posting) => posting.impacts.is_some(), + } + } + pub fn set_positions(&mut self, positions: CompressedPositionStorage) { match self { Self::Plain(posting) => match positions { @@ -3886,6 +3973,7 @@ pub struct CompressedPostingList { pub posting_tail_codec: PostingTailCodec, pub block_size: usize, pub positions: Option, + pub(crate) impacts: Option, // First doc id per block, baked lazily and shared across per-query clones // of the cached list. See `block_first_docs`. first_docs: Arc>>, @@ -3899,6 +3987,7 @@ impl PartialEq for CompressedPostingList { && self.posting_tail_codec == other.posting_tail_codec && self.block_size == other.block_size && self.positions == other.positions + && self.impacts == other.impacts } } @@ -3910,17 +3999,23 @@ impl DeepSizeOf for CompressedPostingList { .as_ref() .map(|positions| positions.deep_size_of_children(context)) .unwrap_or(0) + + self + .impacts + .as_ref() + .map(|impacts| sliced_cache_bytes(impacts.entries())) + .unwrap_or(0) } } impl CompressedPostingList { - pub fn new( + pub(crate) fn new( blocks: LargeBinaryArray, max_score: f32, length: u32, posting_tail_codec: PostingTailCodec, block_size: usize, positions: Option, + impacts: Option, ) -> Self { Self { max_score, @@ -3929,6 +4024,7 @@ impl CompressedPostingList { posting_tail_codec, block_size, positions, + impacts, first_docs: Arc::new(OnceLock::new()), } } @@ -3940,7 +4036,7 @@ impl CompressedPostingList { posting_tail_codec: PostingTailCodec, block_size: usize, shared_position_codec: Option, - ) -> Self { + ) -> Result { debug_assert_eq!(batch.num_rows(), 1); let blocks = batch[POSTING_COL] .as_list::() @@ -3969,16 +4065,28 @@ impl CompressedPostingList { ) }) }; + let impacts = batch + .column_by_name(IMPACT_COL) + .map(|col| { + let entries = col.as_list::().value(0).as_binary::().clone(); + ImpactSkipData::new( + entries, + blocks.len(), + super::impact::ImpactFormat::for_block_size(block_size), + ) + }) + .transpose()?; - Self { + Ok(Self { max_score, length, blocks, posting_tail_codec, block_size, positions, + impacts, first_docs: Arc::new(OnceLock::new()), - } + }) } pub fn iter(&self) -> CompressedPostingListIterator { @@ -4166,6 +4274,7 @@ pub struct PostingListBuilder { pub(super) struct PostingListBatchBuilder { schema: SchemaRef, postings: ListBuilder, + impacts: Option>, max_scores: Float32Builder, lengths: UInt32Builder, positions: BatchPositionsBuilder, @@ -4217,9 +4326,14 @@ impl PostingListBatchBuilder { capacity, )) }; + let impacts = schema + .field_with_name(IMPACT_COL) + .ok() + .map(|_| ListBuilder::with_capacity(LargeBinaryBuilder::new(), capacity)); Self { schema, postings: ListBuilder::with_capacity(LargeBinaryBuilder::new(), capacity), + impacts, max_scores: Float32Builder::with_capacity(capacity), lengths: UInt32Builder::with_capacity(capacity), positions, @@ -4239,11 +4353,19 @@ impl PostingListBatchBuilder { fn append( &mut self, compressed: LargeBinaryArray, + impacts: Option<&ImpactSkipData>, max_score: f32, length: u32, positions: Option<&CompressedPositionStorage>, ) -> Result<()> { - let posting_bytes = compressed.value_data().len(); + let impact_bytes = if self.impacts.is_some() { + impacts + .map(|impacts| impacts.entries().value_data().len()) + .unwrap_or(0) + } else { + 0 + }; + let posting_bytes = compressed.value_data().len() + impact_bytes; { let values = self.postings.values(); for index in 0..compressed.len() { @@ -4251,6 +4373,19 @@ impl PostingListBatchBuilder { } } self.postings.append(true); + if let Some(impacts_builder) = &mut self.impacts { + let impacts = impacts.ok_or_else(|| { + Error::index(format!( + "impacts builder missing impact data for posting length {}", + length + )) + })?; + let values = impacts_builder.values(); + for index in 0..impacts.entries().len() { + values.append_value(impacts.entries().value(index)); + } + impacts_builder.append(true); + } self.group_accumulator.push(posting_bytes); self.max_scores.append_value(max_score); self.lengths.append_value(length); @@ -4316,6 +4451,9 @@ impl PostingListBatchBuilder { Arc::new(self.max_scores.finish()) as ArrayRef, Arc::new(self.lengths.finish()) as ArrayRef, ]; + if let Some(impacts) = &mut self.impacts { + columns.push(Arc::new(impacts.finish()) as ArrayRef); + } match &mut self.positions { BatchPositionsBuilder::None => {} BatchPositionsBuilder::Legacy(position_lists) => { @@ -4733,6 +4871,7 @@ impl PostingListBuilder { fn build_batch( self, compressed: LargeBinaryArray, + impacts: Option, max_score: f32, schema: SchemaRef, positions: Option, @@ -4751,6 +4890,22 @@ impl PostingListBuilder { length as u32, ))) as ArrayRef, ]; + if schema.field_with_name(IMPACT_COL).is_ok() { + let impacts = impacts.ok_or_else(|| { + Error::index(format!( + "impact column requested without impact data for posting length {}", + length + )) + })?; + let impact_offsets = + OffsetBuffer::new(ScalarBuffer::from(vec![0, impacts.entries().len() as i32])); + columns.push(Arc::new(ListArray::try_new( + Arc::new(Field::new("item", datatypes::DataType::LargeBinary, true)), + impact_offsets, + Arc::new(impacts.entries().clone()), + None, + )?) as ArrayRef); + } columns.extend(Self::build_position_columns(positions)?); let batch = RecordBatch::try_new(schema, columns)?; @@ -4821,13 +4976,19 @@ impl PostingListBuilder { tail_entries: tail_entries.as_slice(), tail_position_block: with_positions.then(|| tail_positions.finish()), }; - let (compressed, shared_positions, max_score) = + let (compressed, shared_positions, max_score, impacts) = Self::build_compressed_with_scores_from_parts(parts, docs)?; let positions = match legacy_positions { Some(positions) => Some(CompressedPositionStorage::LegacyPerDoc(positions)), None => shared_positions.map(CompressedPositionStorage::SharedStream), }; - batch_builder.append(compressed, max_score, len, positions.as_ref()) + batch_builder.append( + compressed, + Some(&impacts), + max_score, + len, + positions.as_ref(), + ) } fn extend_tail_components( @@ -4844,7 +5005,12 @@ impl PostingListBuilder { fn build_compressed_with_scores_from_parts( parts: PostingListParts<'_>, docs: &DocSet, - ) -> Result<(LargeBinaryArray, Option, f32)> { + ) -> Result<( + LargeBinaryArray, + Option, + f32, + ImpactSkipData, + )> { let PostingListParts { with_positions, posting_tail_codec, @@ -4860,6 +5026,9 @@ impl PostingListBuilder { let mut max_score = f32::MIN; let mut doc_ids = Vec::with_capacity(block_size); let mut frequencies = Vec::with_capacity(block_size); + let mut impact_block = Vec::with_capacity(block_size); + let mut impact_builder = + ImpactSkipDataBuilder::with_capacity(length.div_ceil(block_size), block_size); for index in 0..encoded_blocks.len() { let block = encoded_blocks.block(index); @@ -4871,13 +5040,15 @@ impl PostingListBuilder { &mut frequencies, block_size, ); - let block_score = compute_block_score( + let block_score = compute_block_score_and_impact_block( docs, avgdl, idf_scale, doc_ids.iter().copied(), frequencies.iter().copied(), + &mut impact_block, ); + impact_builder.append_block(impact_block.as_slice())?; max_score = max_score.max(block_score); if super::encoding::posting_block_score_prefix_len(block_size) > 0 { encoded_blocks.set_block_score(index, block_score); @@ -4886,13 +5057,15 @@ impl PostingListBuilder { if !tail_entries.is_empty() { Self::extend_tail_components(tail_entries, &mut doc_ids, &mut frequencies); - let block_score = compute_block_score( + let block_score = compute_block_score_and_impact_block( docs, avgdl, idf_scale, doc_ids.iter().copied(), frequencies.iter().copied(), + &mut impact_block, ); + impact_builder.append_block(impact_block.as_slice())?; max_score = max_score.max(block_score); encoded_blocks.append_remainder_block_with_codec( doc_ids.as_slice(), @@ -4912,10 +5085,12 @@ impl PostingListBuilder { } } + let impacts = impact_builder.finish()?; Ok(( encoded_blocks.into_array(), with_positions.then(|| encoded_position_blocks.into_stream()), max_score, + impacts, )) } @@ -4981,10 +5156,11 @@ impl PostingListBuilder { self.posting_tail_codec, self.block_size, )?; - let schema = inverted_list_schema_for_version_with_block_size( + let schema = inverted_list_schema_for_version_with_block_size_and_impacts( self.has_positions(), format_version, self.block_size, + false, ); let legacy_positions = if self.with_positions && !format_version.uses_shared_position_stream() { @@ -5042,7 +5218,7 @@ impl PostingListBuilder { Some(positions) => Some(CompressedPositionStorage::LegacyPerDoc(positions)), None => shared_positions.map(CompressedPositionStorage::SharedStream), }; - builder.build_batch(compressed, max_score, schema, positions) + builder.build_batch(compressed, None, max_score, schema, positions) } pub fn to_batch_with_docs(self, docs: &DocSet, schema: SchemaRef) -> Result { @@ -5084,7 +5260,7 @@ impl PostingListBuilder { tail_entries: tail_entries.as_slice(), tail_position_block: with_positions.then(|| tail_positions.finish()), }; - let (compressed, shared_positions, max_score) = + let (compressed, shared_positions, max_score, impacts) = Self::build_compressed_with_scores_from_parts(parts, docs)?; let builder = Self { with_positions, @@ -5104,7 +5280,7 @@ impl PostingListBuilder { Some(positions) => Some(CompressedPositionStorage::LegacyPerDoc(positions)), None => shared_positions.map(CompressedPositionStorage::SharedStream), }; - builder.build_batch(compressed, max_score, schema, positions) + builder.build_batch(compressed, Some(impacts), max_score, schema, positions) } pub fn remap(&mut self, removed: &[u32]) { @@ -5132,19 +5308,23 @@ impl PostingListBuilder { } } -fn compute_block_score( +fn compute_block_score_and_impact_block( docs: &DocSet, avgdl: f32, idf_scale: f32, doc_ids: impl Iterator, frequencies: impl Iterator, + impact_block: &mut Vec<(u32, u32, u32)>, ) -> f32 { + impact_block.clear(); let mut block_max_score = f32::MIN; for (doc_id, freq) in doc_ids.zip(frequencies) { - let doc_norm = K1 * (1.0 - B + B * docs.num_tokens(doc_id) as f32 / avgdl); - let freq = freq as f32; - let score = freq / (freq + doc_norm); + let doc_len = docs.num_tokens(doc_id); + let doc_norm = K1 * (1.0 - B + B * doc_len as f32 / avgdl); + let freq_f32 = freq as f32; + let score = freq_f32 / (freq_f32 + doc_norm); block_max_score = block_max_score.max(score); + impact_block.push((doc_id, freq, doc_len)); } block_max_score * idf_scale } @@ -6110,7 +6290,10 @@ mod tests { use crate::prefilter::NoFilter; use crate::scalar::ScalarIndex; use crate::scalar::inverted::builder::{ - InnerBuilder, InvertedIndexBuilder, PositionRecorder, inverted_list_schema, + InnerBuilder, InvertedIndexBuilder, PositionRecorder, doc_file_path, inverted_list_schema, + inverted_list_schema_for_version_with_block_size, + inverted_list_schema_for_version_with_block_size_and_impacts, posting_file_path, + token_file_path, }; use crate::scalar::inverted::encoding::{ compress_positions, compress_posting_list_with_tail_codec, @@ -6210,6 +6393,57 @@ mod tests { assert!(err.to_string().contains("block_size")); } + #[test] + fn test_posting_builder_writes_impacts_for_supported_block_sizes() { + for block_size in [128, 256] { + let format_version = default_fts_format_version_for_block_size(block_size).unwrap(); + let num_docs = block_size * 33 + 1; + let mut docs = DocSet::default(); + let mut posting = PostingListBuilder::new_with_posting_tail_codec_and_block_size( + false, + format_version.posting_tail_codec(), + block_size, + ); + for doc_id in 0..num_docs { + docs.append(doc_id as u64, (doc_id % 5 + 1) as u32); + posting.add( + doc_id as u32, + PositionRecorder::Count((doc_id % 3 + 1) as u32), + ); + } + let schema = + inverted_list_schema_for_version_with_block_size(false, format_version, block_size); + let batch = posting.to_batch_with_docs(&docs, schema).unwrap(); + assert!(batch.column_by_name(IMPACT_COL).is_some()); + let max_score = batch[MAX_SCORE_COL].as_primitive::().value(0); + let length = batch[LENGTH_COL].as_primitive::().value(0); + let posting = PostingList::from_batch(&batch, Some(max_score), Some(length)).unwrap(); + let PostingList::Compressed(posting) = posting else { + panic!("expected compressed posting list"); + }; + let impacts = posting.impacts.expect("posting should include impacts"); + assert_eq!(impacts.level0_len(), posting.blocks.len()); + assert_eq!(impacts.level1_len(), posting.blocks.len().div_ceil(32)); + assert_eq!( + impacts.entries().len(), + impacts.level0_len() + impacts.level1_len() + ); + } + } + + #[test] + fn test_posting_builder_without_impact_column_roundtrips_without_impacts() { + let mut posting = PostingListBuilder::new(false); + for doc_id in 0..BLOCK_SIZE + 3 { + posting.add(doc_id as u32, PositionRecorder::Count(1)); + } + let batch = posting.to_batch(vec![1.0, 1.0]).unwrap(); + assert!(batch.column_by_name(IMPACT_COL).is_none()); + let posting = + PostingList::from_batch(&batch, Some(1.0), Some((BLOCK_SIZE + 3) as u32)).unwrap(); + assert!(!posting.has_impacts()); + } + #[tokio::test] async fn test_build_search_uses_configured_posting_block_size() { let tmpdir = TempObjDir::default(); @@ -6261,6 +6495,16 @@ mod tests { }; assert_eq!(posting.block_size, block_size); assert_eq!(posting.blocks.len(), num_docs.div_ceil(block_size)); + let impacts = posting + .impacts + .as_ref() + .expect("newly written posting list should include impacts"); + assert_eq!(impacts.level0_len(), posting.blocks.len()); + assert_eq!(impacts.level1_len(), posting.blocks.len().div_ceil(32)); + assert_eq!( + impacts.entries().len(), + impacts.level0_len() + impacts.level1_len() + ); let tokens = Arc::new(Tokens::new(vec!["needle".to_owned()], DocType::Text)); let params = Arc::new(FtsSearchParams::new().with_limit(Some(10))); @@ -7012,7 +7256,11 @@ mod tests { let (start, end) = inverted_list.group_range_for_token(0).unwrap(); let group = inverted_list .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&PostingListGroupKey { + start, + end, + has_impacts: inverted_list.has_impacts, + }) .await .unwrap(); @@ -7324,7 +7572,11 @@ mod tests { let (start, end) = inverted_list.group_range_for_token(token_id).unwrap(); let group = inverted_list .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&PostingListGroupKey { + start, + end, + has_impacts: inverted_list.has_impacts, + }) .await .unwrap(); let slot = (token_id - start) as usize; @@ -7799,7 +8051,11 @@ mod tests { let (start, end) = inverted_list.group_range_for_token(0).unwrap(); let group = inverted_list .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&PostingListGroupKey { + start, + end, + has_impacts: inverted_list.has_impacts, + }) .await .unwrap(); @@ -7888,6 +8144,7 @@ mod tests { PostingTailCodec::Fixed32, LEGACY_BLOCK_SIZE, None, + None, ); let full_backing = full.get_buffer_memory_size(); @@ -8034,7 +8291,11 @@ mod tests { let (start, end) = inverted_list.group_range_for_token(0).unwrap(); let group = inverted_list .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&PostingListGroupKey { + start, + end, + has_impacts: inverted_list.has_impacts, + }) .await .unwrap(); assert!( @@ -8278,6 +8539,178 @@ mod tests { writer.finish_with_metadata(metadata).await.unwrap(); } + async fn write_test_partition_with_optional_impacts( + store: &Arc, + partition_id: u64, + mut builder: InnerBuilder, + token_set_format: TokenSetFormat, + with_impacts: bool, + ) { + let format_version = InvertedListFormatVersion::V1; + let block_size = LEGACY_BLOCK_SIZE; + let docs = std::mem::take(&mut builder.docs); + let schema = inverted_list_schema_for_version_with_block_size_and_impacts( + false, + format_version, + block_size, + with_impacts, + ); + + let mut posting_writer = store + .new_index_file(&posting_file_path(partition_id), schema.clone()) + .await + .unwrap(); + for posting_list in std::mem::take(&mut builder.posting_lists) { + let batch = posting_list + .to_batch_with_docs(&docs, schema.clone()) + .unwrap(); + posting_writer.write_record_batch(batch).await.unwrap(); + } + posting_writer.finish().await.unwrap(); + + let token_batch = std::mem::take(&mut builder.tokens) + .to_batch(token_set_format) + .unwrap(); + let mut token_writer = store + .new_index_file(&token_file_path(partition_id), token_batch.schema()) + .await + .unwrap(); + token_writer.write_record_batch(token_batch).await.unwrap(); + token_writer.finish().await.unwrap(); + + let doc_batch = docs.to_batch().unwrap(); + let mut doc_writer = store + .new_index_file(&doc_file_path(partition_id), doc_batch.schema()) + .await + .unwrap(); + doc_writer.write_record_batch(doc_batch).await.unwrap(); + doc_writer.finish().await.unwrap(); + } + + #[tokio::test] + async fn test_mixed_impact_and_legacy_partitions_use_global_final_scores() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + let mut impact_builder = InnerBuilder::new_with_format_version( + 0, + false, + TokenSetFormat::default(), + InvertedListFormatVersion::V1, + ); + impact_builder.tokens.add("alpha".to_owned()); + impact_builder + .posting_lists + .push(PostingListBuilder::new_with_posting_tail_codec( + false, + InvertedListFormatVersion::V1.posting_tail_codec(), + )); + impact_builder.posting_lists[0].add(0, PositionRecorder::Count(1)); + impact_builder.docs.append(100, 5_000); + for row_id in 101..111 { + impact_builder.docs.append(row_id, 5_000); + } + write_test_partition_with_optional_impacts( + &store, + 0, + impact_builder, + TokenSetFormat::default(), + true, + ) + .await; + + let mut legacy_builder = InnerBuilder::new_with_format_version( + 1, + false, + TokenSetFormat::default(), + InvertedListFormatVersion::V1, + ); + legacy_builder.tokens.add("alpha".to_owned()); + legacy_builder + .posting_lists + .push(PostingListBuilder::new_with_posting_tail_codec( + false, + InvertedListFormatVersion::V1.posting_tail_codec(), + )); + legacy_builder.posting_lists[0].add(0, PositionRecorder::Count(1)); + legacy_builder.docs.append(200, 1_000); + for row_id in 201..301 { + legacy_builder.docs.append(row_id, 1); + } + write_test_partition_with_optional_impacts( + &store, + 1, + legacy_builder, + TokenSetFormat::default(), + false, + ) + .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 impact_partition = index + .partitions + .iter() + .find(|partition| partition.id() == 0) + .unwrap(); + let legacy_partition = index + .partitions + .iter() + .find(|partition| partition.id() == 1) + .unwrap(); + + let impact_posting = impact_partition + .inverted_list + .posting_list(0, false, &NoOpMetricsCollector) + .await + .unwrap(); + assert!(impact_posting.has_impacts()); + + let legacy_posting = legacy_partition + .inverted_list + .posting_list(0, false, &NoOpMetricsCollector) + .await + .unwrap(); + assert!(!legacy_posting.has_impacts()); + + let tokens = Arc::new(Tokens::new(vec!["alpha".to_string()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(1))); + let (row_ids, scores) = index + .bm25_search( + tokens.clone(), + params.clone(), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + + assert_eq!(row_ids, vec![200]); + assert_eq!(row_ids.len(), scores.len()); + + let scorer = index + .bm25_base_scorer(tokens.as_ref(), params.as_ref()) + .await + .unwrap(); + let expected_score = scorer.query_weight("alpha") * scorer.doc_weight(1, 1_000); + assert!( + (scores[0] - expected_score).abs() < 1e-6, + "score: {}, expected: {}", + scores[0], + expected_score + ); + } + #[tokio::test] async fn test_and_query_returns_empty_when_exact_term_missing() { let tmpdir = TempObjDir::default(); @@ -9463,7 +9896,11 @@ mod tests { assert!( posting_reader .index_cache - .get_with_key(&PostingListGroupKey { start, end }) + .get_with_key(&PostingListGroupKey { + start, + end, + has_impacts: posting_reader.has_impacts, + }) .await .is_some(), "prewarm did not populate group [{start}, {end}) that the read \ @@ -9610,7 +10047,10 @@ mod tests { assert!( posting_reader .index_cache - .get_with_key(&PostingListKey { token_id }) + .get_with_key(&PostingListKey { + token_id, + has_impacts: posting_reader.has_impacts, + }) .await .is_some(), "fallback prewarm should populate per-token entry {token_id}", diff --git a/rust/lance-index/src/scalar/inverted/scorer.rs b/rust/lance-index/src/scalar/inverted/scorer.rs index eb7d78ff397..b6b48ce4149 100644 --- a/rust/lance-index/src/scalar/inverted/scorer.rs +++ b/rust/lance-index/src/scalar/inverted/scorer.rs @@ -3,6 +3,7 @@ use super::InvertedPartition; use std::collections::HashMap; +use std::sync::Arc; // the Scorer trait is used to calculate the score of a token in a document // in general, the score is calculated as: @@ -12,6 +13,16 @@ pub trait Scorer: Send + Sync { fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32; } +impl Scorer for Arc { + fn query_weight(&self, token: &str) -> f32 { + self.as_ref().query_weight(token) + } + + fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32 { + self.as_ref().doc_weight(freq, doc_tokens) + } +} + // BM25 parameters pub const K1: f32 = 1.2; pub const B: f32 = 0.75; diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 9a48a72dfbb..467ada26e4f 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -22,6 +22,7 @@ use crate::metrics::MetricsCollector; use super::{ CompressedPositionStorage, + impact::{IMPACT_LEVEL1_BLOCKS, ImpactScoreCache, ImpactSkipData}, query::Operator, scorer::{K1, idf}, }; @@ -38,6 +39,7 @@ use super::{ use super::{DocInfo, builder::BLOCK_SIZE}; const TERMINATED_DOC_ID: u64 = u64::MAX; +const LINEAR_BLOCK_SKIP_LIMIT: usize = 8; pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { std::env::var("LANCE_FLAT_SEARCH_PERCENT_THRESHOLD") .unwrap_or_else(|_| "10".to_string()) @@ -45,6 +47,33 @@ pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { .unwrap_or(10) }); +#[inline] +fn posting_block_idx(index: usize, block_size: usize) -> usize { + match block_size { + 128 => index >> 7, + 256 => index >> 8, + _ => index / block_size, + } +} + +#[inline] +fn posting_block_offset(index: usize, block_size: usize) -> usize { + match block_size { + 128 => index & 127, + 256 => index & 255, + _ => index % block_size, + } +} + +#[inline] +fn posting_block_start(block_idx: usize, block_size: usize) -> usize { + match block_size { + 128 => block_idx << 7, + 256 => block_idx << 8, + _ => block_idx * block_size, + } +} + pub struct PostingIterator { token: String, token_id: u32, @@ -55,6 +84,7 @@ pub struct PostingIterator { index: usize, // the index of current block, this can be changed by `next() and shallow_next()` block_idx: usize, + current_doc: Option, approximate_upper_bound: f32, // for compressed posting list @@ -71,6 +101,8 @@ struct CompressedState { position_values: Vec, position_offsets: Vec, block_max_window: BlockMaxWindow, + current_block_max_score: Option<(usize, f32)>, + block_least_doc_ids: Vec>, } impl CompressedState { @@ -84,6 +116,8 @@ impl CompressedState { position_values: Vec::new(), position_offsets: Vec::new(), block_max_window: BlockMaxWindow::new(), + current_block_max_score: None, + block_least_doc_ids: Vec::new(), } } @@ -100,7 +134,7 @@ impl CompressedState { self.doc_ids.clear(); self.freqs.clear(); - let remainder = length as usize % block_size; + let remainder = posting_block_offset(length as usize, block_size); if block_idx + 1 == num_blocks && remainder != 0 { decompress_posting_remainder( block, @@ -133,6 +167,7 @@ struct BlockMaxWindow { start_block_idx: usize, next_block_idx: usize, max_scores: VecDeque<(usize, f32)>, + impact_score_cache: ImpactScoreCache, } struct BlockMaxScore { @@ -146,6 +181,7 @@ impl BlockMaxWindow { start_block_idx: 0, next_block_idx: 0, max_scores: VecDeque::new(), + impact_score_cache: ImpactScoreCache::default(), } } @@ -155,12 +191,28 @@ impl BlockMaxWindow { self.max_scores.clear(); } - fn max_score_up_to( + fn max_score_up_to( &mut self, list: &CompressedPostingList, start_block_idx: usize, up_to: u64, + query_weight: f32, + scorer: &S, ) -> BlockMaxScore { + if let Some(impacts) = &list.impacts { + let score = impacts.max_score_up_to_cached( + start_block_idx, + up_to, + query_weight, + scorer, + &mut self.impact_score_cache, + ); + return BlockMaxScore { + score: score.score, + blocks_scanned: score.entries_scanned, + }; + } + if start_block_idx >= list.blocks.len() { self.reset(start_block_idx); return BlockMaxScore { @@ -264,6 +316,96 @@ impl Ord for PostingIterator { } impl PostingIterator { + #[inline] + fn block_least_doc_id(&self, list: &CompressedPostingList, block_idx: usize) -> u32 { + let compressed = unsafe { &mut *self.compressed_state_ptr() }; + if compressed.block_least_doc_ids.len() < list.blocks.len() { + compressed + .block_least_doc_ids + .resize(list.blocks.len(), None); + } + if let Some(doc_id) = compressed.block_least_doc_ids[block_idx] { + return doc_id; + } + let doc_id = list.block_least_doc_id(block_idx); + compressed.block_least_doc_ids[block_idx] = Some(doc_id); + doc_id + } + + fn block_idx_for_doc( + &self, + list: &CompressedPostingList, + mut block_idx: usize, + least_id: u32, + ) -> usize { + let mut linear_skips = 0; + while block_idx + 1 < list.blocks.len() && linear_skips < LINEAR_BLOCK_SKIP_LIMIT { + if self.block_least_doc_id(list, block_idx + 1) > least_id { + return block_idx; + } + block_idx += 1; + linear_skips += 1; + } + + if block_idx + 1 >= list.blocks.len() { + return block_idx; + } + + if let Some(impacts) = list.impacts.as_ref() + && let Some(block_idx) = + self.block_idx_for_doc_with_impacts(list, impacts, block_idx, least_id) + { + return block_idx; + } + + self.block_idx_for_doc_by_least_doc_id(list, block_idx, least_id, list.blocks.len()) + } + + fn block_idx_for_doc_with_impacts( + &self, + list: &CompressedPostingList, + impacts: &ImpactSkipData, + mut block_idx: usize, + least_id: u32, + ) -> Option { + while block_idx + 1 < list.blocks.len() { + let group_idx = (block_idx + 1) / IMPACT_LEVEL1_BLOCKS; + let group_end = ((group_idx + 1) * IMPACT_LEVEL1_BLOCKS).min(list.blocks.len()); + let group_doc_up_to = impacts.level1_doc_up_to(group_idx)?; + if group_doc_up_to < least_id { + block_idx = group_end - 1; + continue; + } + if group_doc_up_to == least_id { + return Some(group_end - 1); + } + return Some( + self.block_idx_for_doc_by_least_doc_id(list, block_idx, least_id, group_end), + ); + } + Some(block_idx) + } + + fn block_idx_for_doc_by_least_doc_id( + &self, + list: &CompressedPostingList, + block_idx: usize, + least_id: u32, + right: usize, + ) -> usize { + let mut left = block_idx + 1; + let mut right = right; + while left < right { + let mid = left + (right - left) / 2; + if self.block_least_doc_id(list, mid) <= least_id { + left = mid + 1; + } else { + right = mid; + } + } + left - 1 + } + #[inline] fn compressed_state_ptr(&self) -> *mut CompressedState { debug_assert!(self.compressed.is_some()); @@ -312,9 +454,12 @@ impl PostingIterator { list: PostingList, num_doc: usize, ) -> Self { - let approximate_upper_bound = match list.max_score() { - Some(max_score) => max_score, - None => idf(list.len(), num_doc) * (K1 + 1.0), + let approximate_upper_bound = match &list { + PostingList::Compressed(posting) if posting.impacts.is_some() => f32::INFINITY, + _ => match list.max_score() { + Some(max_score) => max_score, + None => idf(list.len(), num_doc) * (K1 + 1.0), + }, }; let compressed = match &list { PostingList::Compressed(list) => { @@ -323,7 +468,7 @@ impl PostingIterator { PostingList::Plain(_) => None, }; - Self { + let mut posting = Self { token, token_id, position, @@ -331,9 +476,12 @@ impl PostingIterator { list, index: 0, block_idx: 0, + current_doc: None, approximate_upper_bound, compressed, - } + }; + posting.refresh_current_doc(); + posting } #[inline] @@ -351,8 +499,24 @@ impl PostingIterator { self.approximate_upper_bound } + /// Tightest known list-wide score bound. Impact lists answer from the + /// baked doc-weight slab (the data-driven equivalent of the max_score the + /// non-impact format bakes at build time); everything else falls back to + /// `approximate_upper_bound`. A finite, tight global bound is what lets + /// lagging iterators park in the WAND tail instead of being force-advanced + /// on every candidate. + #[inline] + fn global_upper_bound(&self, scorer: &S) -> f32 { + if let PostingList::Compressed(ref list) = self.list + && let Some(impacts) = list.impacts.as_ref() + { + return self.query_weight * impacts.global_max_doc_weight(scorer); + } + self.approximate_upper_bound + } + #[inline] - fn score(&self, scorer: &S, freq: u32, doc_length: u32) -> f32 { + fn score(&self, scorer: &S, freq: u32, doc_length: u32) -> f32 { self.query_weight * scorer.doc_weight(freq, doc_length) } @@ -368,14 +532,19 @@ impl PostingIterator { #[inline] fn doc(&self) -> Option { + self.current_doc + } + + fn refresh_current_doc(&mut self) { if self.empty() { - return None; + self.current_doc = None; + return; } - match self.list { + let current_doc = match self.list { PostingList::Compressed(ref list) => { - let block_idx = self.index / list.block_size; - let block_offset = self.index % list.block_size; + let block_idx = posting_block_idx(self.index, list.block_size); + let block_offset = posting_block_offset(self.index, list.block_size); let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; // Read from the decompressed block @@ -385,7 +554,8 @@ impl PostingIterator { Some(doc) } PostingList::Plain(ref list) => Some(DocInfo::Located(list.doc(self.index))), - } + }; + self.current_doc = current_doc; } fn position_cursor(&self) -> Option> { @@ -413,8 +583,8 @@ impl PostingIterator { )) } CompressedPositionStorage::SharedStream(stream) => { - let block_idx = self.index / list.block_size; - let block_offset = self.index % list.block_size; + let block_idx = posting_block_idx(self.index, list.block_size); + let block_offset = posting_block_offset(self.index, list.block_size); let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; if compressed.position_block_idx != Some(block_idx) { @@ -454,36 +624,46 @@ impl PostingIterator { PostingList::Compressed(ref list) => { debug_assert!(least_id <= u32::MAX as u64); let least_id = least_id as u32; - let mut block_idx = self.index / list.block_size; - while block_idx + 1 < list.blocks.len() - && list.block_least_doc_id(block_idx + 1) <= least_id - { - block_idx += 1; - } - self.index = self.index.max(block_idx * list.block_size); + let block_idx = self.block_idx_for_doc( + list, + posting_block_idx(self.index, list.block_size), + least_id, + ); + self.index = self + .index + .max(posting_block_start(block_idx, list.block_size)); let length = list.length as usize; while self.index < length { - let block_idx = self.index / list.block_size; - let block_offset = self.index % list.block_size; + let block_idx = posting_block_idx(self.index, list.block_size); + let block_offset = posting_block_offset(self.index, list.block_size); let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; let in_block = &compressed.doc_ids[block_offset..]; let offset_in_block = in_block.partition_point(|&doc_id| doc_id < least_id); let new_offset = block_offset + offset_in_block; if new_offset < compressed.doc_ids.len() { - self.index = block_idx * list.block_size + new_offset; - break; + self.index = posting_block_start(block_idx, list.block_size) + new_offset; + self.block_idx = block_idx; + self.current_doc = Some(DocInfo::Raw(RawDocInfo { + doc_id: compressed.doc_ids[new_offset], + frequency: compressed.freqs[new_offset], + })); + return; } if block_idx + 1 >= list.blocks.len() { self.index = length; + self.block_idx = posting_block_idx(self.index, list.block_size); + self.current_doc = None; break; } - self.index = (block_idx + 1) * list.block_size; + self.index = posting_block_start(block_idx + 1, list.block_size); } - self.block_idx = self.index / list.block_size; + self.block_idx = posting_block_idx(self.index, list.block_size); + self.current_doc = None; } PostingList::Plain(ref list) => { self.index += list.row_ids[self.index..].partition_point(|&id| id < least_id); + self.current_doc = (!self.empty()).then(|| DocInfo::Located(list.doc(self.index))); } } } @@ -493,11 +673,7 @@ impl PostingIterator { PostingList::Compressed(ref list) => { debug_assert!(least_id <= u32::MAX as u64); let least_id = least_id as u32; - while self.block_idx + 1 < list.blocks.len() - && list.block_least_doc_id(self.block_idx + 1) <= least_id - { - self.block_idx += 1; - } + self.block_idx = self.block_idx_for_doc(list, self.block_idx, least_id); } PostingList::Plain(_) => { // we don't have block max score for legacy index, @@ -507,21 +683,48 @@ impl PostingIterator { } #[inline] - fn block_max_score(&self) -> f32 { + fn block_max_score(&self, scorer: &S) -> f32 { match self.list { - PostingList::Compressed(ref list) => list.block_max_score(self.block_idx), + PostingList::Compressed(ref list) => { + if let Some(impacts) = list.impacts.as_ref() { + let compressed = unsafe { &mut *self.compressed_state_ptr() }; + if let Some((block_idx, score)) = compressed.current_block_max_score + && block_idx == self.block_idx + { + return score; + } + + let score = impacts.level0_score_cached( + self.block_idx, + self.query_weight, + scorer, + &mut compressed.block_max_window.impact_score_cache, + ); + compressed.current_block_max_score = Some((self.block_idx, score)); + return score; + } + list.block_max_score(self.block_idx) + } PostingList::Plain(_) => self.approximate_upper_bound, } } #[inline] - fn block_max_score_up_to_with_stats(&mut self, up_to: u64) -> BlockMaxScore { + fn block_max_score_up_to_with_stats( + &mut self, + up_to: u64, + scorer: &S, + ) -> BlockMaxScore { match self.list { PostingList::Compressed(ref list) => { let compressed = unsafe { &mut *self.compressed_state_ptr() }; - compressed - .block_max_window - .max_score_up_to(list, self.block_idx, up_to) + compressed.block_max_window.max_score_up_to( + list, + self.block_idx, + up_to, + self.query_weight, + scorer, + ) } PostingList::Plain(_) => BlockMaxScore { score: self.approximate_upper_bound, @@ -546,7 +749,7 @@ impl PostingIterator { fn block_first_doc(&self) -> Option { match self.list { PostingList::Compressed(ref list) => { - Some(list.block_least_doc_id(self.block_idx) as u64) + Some(self.block_least_doc_id(list, self.block_idx) as u64) } PostingList::Plain(ref plain) => plain.row_ids.get(self.index).cloned(), } @@ -559,7 +762,7 @@ impl PostingIterator { if self.block_idx + 1 >= list.blocks.len() { return None; } - Some(list.block_least_doc_id(self.block_idx + 1) as u64) + Some(self.block_least_doc_id(list, self.block_idx + 1) as u64) } PostingList::Plain(ref plain) => plain.row_ids.get(self.index + 1).cloned(), } @@ -1187,7 +1390,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let remaining_upper_bound = remaining .iter() - .map(|posting| posting.block_max_score()) + .map(|posting| posting.block_max_score(&self.scorer)) .sum::(); first.score(&self.scorer, doc.frequency(), doc_length) + remaining_upper_bound <= self.threshold @@ -1374,7 +1577,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let narrow_max_score = self .lead .iter() - .map(|posting| posting.block_max_score()) + .map(|posting| posting.block_max_score(&self.scorer)) .sum::(); if narrow_max_score >= self.threshold { @@ -1397,7 +1600,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut wide_max_score = 0.0; let mut range_blocks_scanned = 0; for posting in &mut self.lead { - let block_max = posting.block_max_score_up_to_with_stats(lead_up_to); + let block_max = posting.block_max_score_up_to_with_stats(lead_up_to, &self.scorer); wide_max_score += block_max.score; range_blocks_scanned += block_max.blocks_scanned; } @@ -1461,7 +1664,7 @@ impl<'a, S: Scorer> Wand<'a, S> { fn move_head_before_target_to_tail(&mut self, target: u64) { while matches!(self.head_doc(), Some(doc_id) if doc_id < target) { if let Some(posting) = self.head.pop() { - let upper_bound = posting.posting.approximate_upper_bound(); + let upper_bound = posting.posting.global_upper_bound(&self.scorer); if let Some(mut evicted) = self.insert_tail_with_overflow(posting.posting, upper_bound) { @@ -1480,12 +1683,12 @@ impl<'a, S: Scorer> Wand<'a, S> { let lead: f32 = self .lead .iter() - .map(|posting| posting.block_max_score()) + .map(|posting| posting.block_max_score(&self.scorer)) .sum(); let head: f32 = self .head .iter() - .map(|posting| posting.posting.block_max_score()) + .map(|posting| posting.posting.block_max_score(&self.scorer)) .sum(); lead + head + self.tail_max_score } @@ -1498,12 +1701,12 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut sum = self .lead .iter() - .map(|posting| posting.block_max_score()) + .map(|posting| posting.block_max_score(&self.scorer)) .sum::(); let mut possible_matches = self.lead.len(); for posting in &self.tail { if matches!(posting.posting.block_first_doc(), Some(block_doc) if block_doc <= target) { - sum += posting.posting.block_max_score(); + sum += posting.posting.block_max_score(&self.scorer); possible_matches += 1; } } @@ -1569,7 +1772,7 @@ impl<'a, S: Scorer> Wand<'a, S> { Some(block_doc) if block_doc <= target => { tail_posting .posting - .block_max_score_up_to_with_stats(up_to) + .block_max_score_up_to_with_stats(up_to, &self.scorer) .score } _ => 0.0, @@ -1706,9 +1909,9 @@ impl<'a, S: Scorer> Wand<'a, S> { && posting.is_compressed() && self.up_to.is_some_and(|up_to| target <= up_to) { - posting.block_max_score() + posting.block_max_score(&self.scorer) } else { - posting.approximate_upper_bound() + posting.global_upper_bound(&self.scorer) } } @@ -1990,13 +2193,17 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; + use super::super::impact::build_impact_skip_data; use super::*; use crate::scalar::inverted::scorer::IndexBM25Scorer; use crate::{ metrics::{MetricsCollector, NoOpMetricsCollector}, scalar::inverted::{ - CompressedPostingList, PlainPostingList, PostingListBuilder, builder::PositionRecorder, - encoding::compress_posting_list, + CompressedPostingList, PlainPostingList, PostingListBuilder, + builder::PositionRecorder, + encoding::{ + compress_posting_list, compress_posting_list_with_tail_codec_and_block_size, + }, }, }; @@ -2161,6 +2368,7 @@ mod tests { crate::scalar::inverted::PostingTailCodec::VarintDelta, crate::scalar::inverted::LEGACY_BLOCK_SIZE, None, + None, )) } else { PostingList::Plain(PlainPostingList::new( @@ -2172,6 +2380,75 @@ mod tests { } } + fn generate_impact_posting_list_with_freqs( + doc_ids: Vec, + freqs: Vec, + doc_lengths: Vec, + ) -> PostingList { + generate_impact_posting_list_with_freqs_and_block_size( + doc_ids, + freqs, + doc_lengths, + crate::scalar::inverted::LEGACY_BLOCK_SIZE, + ) + } + + fn generate_impact_posting_list_with_freqs_and_block_size( + doc_ids: Vec, + freqs: Vec, + doc_lengths: Vec, + block_size: usize, + ) -> PostingList { + assert_eq!(doc_ids.len(), freqs.len()); + assert_eq!(doc_ids.len(), doc_lengths.len()); + let block_max_scores = vec![0.0; doc_ids.len().div_ceil(block_size)]; + let blocks = compress_posting_list_with_tail_codec_and_block_size( + doc_ids.len(), + doc_ids.iter(), + freqs.iter(), + block_max_scores.into_iter(), + crate::scalar::inverted::PostingTailCodec::VarintDelta, + block_size, + ) + .unwrap(); + let impact_blocks = doc_ids + .chunks(block_size) + .zip(freqs.chunks(block_size)) + .zip(doc_lengths.chunks(block_size)) + .map(|((doc_ids, freqs), doc_lengths)| { + doc_ids + .iter() + .copied() + .zip(freqs.iter().copied()) + .zip(doc_lengths.iter().copied()) + .map(|((doc_id, freq), doc_length)| (doc_id, freq, doc_length)) + .collect::>() + }) + .collect::>(); + let impacts = build_impact_skip_data(impact_blocks.as_slice()).unwrap(); + PostingList::Compressed(CompressedPostingList::new( + blocks, + 0.0, + doc_ids.len() as u32, + crate::scalar::inverted::PostingTailCodec::VarintDelta, + block_size, + None, + Some(impacts), + )) + } + + fn generate_contiguous_impact_posting_list_with_block_size( + total: usize, + block_size: usize, + ) -> PostingList { + generate_impact_posting_list_with_freqs_and_block_size( + (0..total as u32).collect(), + vec![1; total], + vec![1; total], + block_size, + ) + } + fn generate_posting_list_with_positions( doc_ids: Vec, positions_by_doc: Vec>, @@ -2996,7 +3273,7 @@ mod tests { posting.shallow_next(0); assert_eq!( posting - .block_max_score_up_to_with_stats((3 * BLOCK_SIZE - 1) as u64) + .block_max_score_up_to_with_stats((3 * BLOCK_SIZE - 1) as u64, &UnitScorer) .score, 4.0 ); @@ -3004,7 +3281,7 @@ mod tests { posting.shallow_next((2 * BLOCK_SIZE) as u64); assert_eq!( posting - .block_max_score_up_to_with_stats((4 * BLOCK_SIZE - 1) as u64) + .block_max_score_up_to_with_stats((4 * BLOCK_SIZE - 1) as u64, &UnitScorer) .score, 5.0 ); @@ -3012,12 +3289,127 @@ mod tests { posting.shallow_next((4 * BLOCK_SIZE) as u64); assert_eq!( posting - .block_max_score_up_to_with_stats((5 * BLOCK_SIZE - 1) as u64) + .block_max_score_up_to_with_stats((5 * BLOCK_SIZE - 1) as u64, &UnitScorer) .score, 3.0 ); } + #[test] + fn test_impact_level1_skip_keeps_boundary_equality_in_group() { + for block_size in [crate::scalar::inverted::LEGACY_BLOCK_SIZE, 256] { + let total = (IMPACT_LEVEL1_BLOCKS + 1) * block_size; + let mut posting = PostingIterator::new( + String::from("term"), + 0, + 0, + generate_contiguous_impact_posting_list_with_block_size(total, block_size), + total, + ); + let target = (IMPACT_LEVEL1_BLOCKS * block_size - 1) as u64; + + posting.shallow_next(target); + assert_eq!(posting.block_idx, IMPACT_LEVEL1_BLOCKS - 1); + + posting.next(target); + assert_eq!(posting.block_idx, IMPACT_LEVEL1_BLOCKS - 1); + assert_eq!(posting.doc().map(|doc| doc.doc_id()), Some(target)); + } + } + + #[test] + fn test_impact_level1_skip_handles_partial_final_group() { + for block_size in [crate::scalar::inverted::LEGACY_BLOCK_SIZE, 256] { + let total = (IMPACT_LEVEL1_BLOCKS + 3) * block_size + 17; + let mut posting = PostingIterator::new( + String::from("term"), + 0, + 0, + generate_contiguous_impact_posting_list_with_block_size(total, block_size), + total, + ); + let target = (total - 1) as u64; + let expected_block = total.div_ceil(block_size) - 1; + + posting.shallow_next(target); + assert_eq!(posting.block_idx, expected_block); + + posting.next(target); + assert_eq!(posting.block_idx, expected_block); + assert_eq!(posting.doc().map(|doc| doc.doc_id()), Some(target)); + } + } + + #[test] + fn test_impact_level1_skip_reaches_far_target_doc() { + for block_size in [crate::scalar::inverted::LEGACY_BLOCK_SIZE, 256] { + let total = (IMPACT_LEVEL1_BLOCKS * 3 + 5) * block_size; + let target_block = IMPACT_LEVEL1_BLOCKS * 2 + 2; + let target = (target_block * block_size + 17) as u64; + let mut posting = PostingIterator::new( + String::from("term"), + 0, + 0, + generate_contiguous_impact_posting_list_with_block_size(total, block_size), + total, + ); + + posting.shallow_next(target); + assert_eq!(posting.block_idx, target_block); + + posting.next(target); + assert_eq!(posting.block_idx, target_block); + assert_eq!(posting.doc().map(|doc| doc.doc_id()), Some(target)); + } + } + + #[test] + fn test_compressed_impact_block_max_score_memoizes_current_block() { + let total = 2 * BLOCK_SIZE as u32; + let doc_ids = (0..total).collect::>(); + let freqs = doc_ids + .iter() + .map(|doc_id| if *doc_id < BLOCK_SIZE as u32 { 1 } else { 2 }) + .collect::>(); + let doc_lengths = vec![1; total as usize]; + let posting_list = generate_impact_posting_list_with_freqs(doc_ids, freqs, doc_lengths); + let mut posting = + PostingIterator::new(String::from("term"), 0, 0, posting_list, total as usize); + let scored = Arc::new(AtomicUsize::new(0)); + let scorer = CountingScorer { + scored: scored.clone(), + }; + + let first_score = posting.block_max_score(&scorer); + assert_eq!(first_score, 1.0); + // Baking the shared doc-weight bounds visits every frontier pair once + // (two level0 entries plus one level1 entry for this list). + let baked = scored.load(Ordering::Relaxed); + assert!(baked >= 2); + { + let compressed = unsafe { &mut *posting.compressed_state_ptr() }; + assert_eq!(compressed.current_block_max_score, Some((0, first_score))); + compressed.block_max_window.impact_score_cache = ImpactScoreCache::default(); + } + + let second_score = posting.block_max_score(&scorer); + assert_eq!(second_score, first_score); + assert_eq!( + scored.load(Ordering::Relaxed), + baked, + "repeated block max scores must not recompute doc weights" + ); + + posting.shallow_next(BLOCK_SIZE as u64); + let next_block_score = posting.block_max_score(&scorer); + assert_eq!(next_block_score, 2.0); + assert_eq!( + scored.load(Ordering::Relaxed), + baked, + "other blocks answer from the baked bounds without rescoring" + ); + } + #[test] fn test_and_candidate_prune_scores_first_term_before_full_score() { let total_docs = 2 * BLOCK_SIZE as u32 + 1; @@ -3354,7 +3746,7 @@ mod tests { let posting = PostingIterator::new(String::from("test"), 0, 0, posting_list, 1); - let actual = posting.block_max_score(); + let actual = posting.block_max_score(&UnitScorer); assert!( (actual - expected).abs() < 1e-6, "block max score should match stored value" From acac9d885391709b57a1f71bbf5dbb6255408896 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 3 Jul 2026 13:13:21 +0800 Subject: [PATCH 3/7] perf(fts): bulk MAXSCORE search path for top-k disjunctions Port of Lucene's MaxScoreBulkScorer, opt-in via LANCE_FTS_MAXSCORE=1: per outer window (bounded by the essential clauses' blocks with adaptive growth), clauses split into a non-essential prefix and essential rest by window max score vs the running threshold. Essential clauses bulk-stream decompressed blocks (single-essential windows stream with no accumulator); non-essential clauses are only probed for candidates that can still beat the threshold. Dead ranges with one live clause skip by scanning the baked per-block bound slab. Candidate emission matches the classic path (must beat the running threshold), so results are score-identical. Measured on a 200M-doc warm benchmark at 24 partitions: 3-word OR match k10 0.137s -> 0.035s, k100 0.250s -> 0.064s; hot single-term 250ms -> 3ms. --- .../lance-index/src/scalar/inverted/impact.rs | 95 +- rust/lance-index/src/scalar/inverted/wand.rs | 1030 +++++++++++++++-- 2 files changed, 952 insertions(+), 173 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/impact.rs b/rust/lance-index/src/scalar/inverted/impact.rs index b465d0e7cfe..bb710694a1e 100644 --- a/rust/lance-index/src/scalar/inverted/impact.rs +++ b/rust/lance-index/src/scalar/inverted/impact.rs @@ -56,30 +56,14 @@ impl PartialEq for ImpactSkipData { } } +/// Test-only score/scan-count pair returned by [`ImpactSkipData::max_score_up_to`]. +#[cfg(test)] #[derive(Debug, Clone, Copy)] pub struct ImpactScore { pub score: f32, pub entries_scanned: usize, } -#[derive(Debug, Default, Clone)] -pub struct ImpactScoreCache {} - -impl ImpactScoreCache { - fn entry_score( - &mut self, - impacts: &ImpactSkipData, - entry_idx: usize, - query_weight: f32, - scorer: &S, - ) -> f32 { - if query_weight <= 0.0 { - return 0.0; - } - query_weight * impacts.doc_weight_bounds(scorer)[entry_idx] - } -} - #[derive(Debug, Clone, Copy)] struct ImpactEntryHeader { doc_up_to: u32, @@ -160,6 +144,12 @@ impl ImpactSkipData { &self.baked_bounds(scorer).0 } + /// Baked per-block max doc weights (level0 entries only), for bulk skip + /// scans over dead ranges without per-block window bookkeeping. + pub(crate) fn level0_doc_weight_bounds(&self, scorer: &S) -> &[f32] { + &self.baked_bounds(scorer).0[..self.level0_len] + } + /// List-wide max doc weight, from the baked bounds. The tightest valid /// global score bound for this list is `query_weight * this`, matching /// what the non-impact format stores as `max_score` at build time — but @@ -192,11 +182,20 @@ impl ImpactSkipData { } } + /// Last doc id covered by the level0 entry of `block_idx`, or `None` when + /// the entry is missing or malformed. + pub(crate) fn level0_doc_up_to(&self, block_idx: usize) -> Option { + if block_idx >= self.level0_len { + return None; + } + match self.entry_doc_up_tos[block_idx] { + u32::MAX => None, + doc_up_to => Some(doc_up_to), + } + } + /// Max score of the docs covered by the level0 entry of `block_idx`, /// answered from the baked bounds slab. - // Only tests exercise the uncached form until the maxscore rework - // (stacked follow-up) anchors its block-max caches on it. - #[cfg_attr(not(test), allow(dead_code))] pub fn level0_score( &self, block_idx: usize, @@ -209,17 +208,18 @@ impl ImpactSkipData { query_weight * self.doc_weight_bounds(scorer)[block_idx] } - pub fn level0_score_cached( + /// Max score of the docs covered by the level1 entry of `group_idx`, + /// answered from the baked bounds slab. + pub fn level1_score( &self, - block_idx: usize, + group_idx: usize, query_weight: f32, scorer: &S, - cache: &mut ImpactScoreCache, ) -> f32 { - if block_idx >= self.level0_len { + if group_idx >= level1_len(self.level0_len) || query_weight <= 0.0 { return 0.0; } - cache.entry_score(self, block_idx, query_weight, scorer) + query_weight * self.doc_weight_bounds(scorer)[self.level0_len + group_idx] } #[cfg(test)] @@ -240,22 +240,7 @@ impl ImpactSkipData { }) } - pub fn max_score_up_to_cached( - &self, - start_block_idx: usize, - up_to: u64, - query_weight: f32, - scorer: &S, - cache: &mut ImpactScoreCache, - ) -> ImpactScore - where - S: Scorer + ?Sized, - { - self.max_score_up_to_with(start_block_idx, up_to, |impacts, entry_idx| { - cache.entry_score(impacts, entry_idx, query_weight, scorer) - }) - } - + #[cfg(test)] fn max_score_up_to_with( &self, start_block_idx: usize, @@ -638,25 +623,6 @@ mod tests { assert_eq!(impacts.level1_doc_up_to(0), None); } - #[test] - fn impact_score_cache_matches_uncached_scores() { - let blocks = (0..40) - .map(|block| vec![(block as u32, 1 + block as u32 % 3, 10)]) - .collect::>(); - let impacts = build_impact_skip_data(&blocks).unwrap(); - let scorer = MemBM25Scorer::new(400, 40, HashMap::from([(String::from("token"), 40usize)])); - let mut cache = ImpactScoreCache::default(); - - let uncached_level0 = impacts.level0_score(3, 1.0, &scorer); - let cached_level0 = impacts.level0_score_cached(3, 1.0, &scorer, &mut cache); - assert_eq!(cached_level0, uncached_level0); - - let uncached = impacts.max_score_up_to(0, 31, |idx| idx as u32, 1.0, &scorer); - let cached = impacts.max_score_up_to_cached(0, 31, 1.0, &scorer, &mut cache); - assert_eq!(cached.score, uncached.score); - assert_eq!(cached.entries_scanned, uncached.entries_scanned); - } - #[test] fn impact_entries_are_decoded_lazily() { let level0_0 = encode_impact_entry(&[(0, 1, 10)], ImpactFormat::FixedU32).unwrap(); @@ -675,11 +641,8 @@ mod tests { assert!(score.score.is_finite()); assert_eq!(score.entries_scanned, 1); - let mut cache = ImpactScoreCache::default(); - assert_eq!( - impacts.level0_score_cached(1, 1.0, &scorer, &mut cache), - f32::INFINITY - ); + // The malformed level0 entry degrades to an unskippable bound. + assert_eq!(impacts.level0_score(1, 1.0, &scorer), f32::INFINITY); } #[test] diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 467ada26e4f..70ee2015cda 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -22,7 +22,7 @@ use crate::metrics::MetricsCollector; use super::{ CompressedPositionStorage, - impact::{IMPACT_LEVEL1_BLOCKS, ImpactScoreCache, ImpactSkipData}, + impact::{IMPACT_LEVEL1_BLOCKS, ImpactSkipData}, query::Operator, scorer::{K1, idf}, }; @@ -39,6 +39,9 @@ use super::{ 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)>>; const LINEAR_BLOCK_SKIP_LIMIT: usize = 8; pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { std::env::var("LANCE_FLAT_SEARCH_PERCENT_THRESHOLD") @@ -46,6 +49,11 @@ pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { .parse::() .unwrap_or(10) }); +// Bulk MAXSCORE path for top-k disjunctions (Lucene MaxScoreBulkScorer +// style). Off by default while it bakes; LANCE_FTS_MAXSCORE=1 opts in. Its +// results are score-identical to the classic WAND loop. +static USE_MAXSCORE_SEARCH: LazyLock = + LazyLock::new(|| std::env::var("LANCE_FTS_MAXSCORE").as_deref() == Ok("1")); #[inline] fn posting_block_idx(index: usize, block_size: usize) -> usize { @@ -101,8 +109,11 @@ struct CompressedState { position_values: Vec, position_offsets: Vec, block_max_window: BlockMaxWindow, - current_block_max_score: Option<(usize, f32)>, - block_least_doc_ids: Vec>, + // Lucene-style anchored impact score caches: one slot per level, keyed by + // the entry the block cursor currently sits in. Each holds + // (entry_idx, doc_up_to, max_score). See `impact_level0`/`impact_level1`. + level0_cache: Option<(usize, u32, f32)>, + level1_cache: Option<(usize, u32, f32)>, } impl CompressedState { @@ -116,8 +127,8 @@ impl CompressedState { position_values: Vec::new(), position_offsets: Vec::new(), block_max_window: BlockMaxWindow::new(), - current_block_max_score: None, - block_least_doc_ids: Vec::new(), + level0_cache: None, + level1_cache: None, } } @@ -164,10 +175,11 @@ impl CompressedState { struct BlockMaxWindow { // Sliding block range used for Lucene-style getMaxScore(upTo). The deque is // monotonic by score and covers blocks in [start_block_idx, next_block_idx). + // Only used for compressed lists without impact skip data; impact lists + // answer window max scores from the anchored level caches instead. start_block_idx: usize, next_block_idx: usize, max_scores: VecDeque<(usize, f32)>, - impact_score_cache: ImpactScoreCache, } struct BlockMaxScore { @@ -181,7 +193,6 @@ impl BlockMaxWindow { start_block_idx: 0, next_block_idx: 0, max_scores: VecDeque::new(), - impact_score_cache: ImpactScoreCache::default(), } } @@ -199,20 +210,6 @@ impl BlockMaxWindow { query_weight: f32, scorer: &S, ) -> BlockMaxScore { - if let Some(impacts) = &list.impacts { - let score = impacts.max_score_up_to_cached( - start_block_idx, - up_to, - query_weight, - scorer, - &mut self.impact_score_cache, - ); - return BlockMaxScore { - score: score.score, - blocks_scanned: score.entries_scanned, - }; - } - if start_block_idx >= list.blocks.len() { self.reset(start_block_idx); return BlockMaxScore { @@ -242,7 +239,10 @@ impl BlockMaxWindow { while self.next_block_idx < list.blocks.len() && list.block_least_doc_id(self.next_block_idx) as u64 <= up_to { - let score = list.block_max_score(self.next_block_idx); + let score = match list.impacts.as_ref() { + Some(impacts) => impacts.level0_score(self.next_block_idx, query_weight, scorer), + None => list.block_max_score(self.next_block_idx), + }; while matches!(self.max_scores.back(), Some((_, old_score)) if *old_score <= score) { self.max_scores.pop_back(); } @@ -318,18 +318,7 @@ impl Ord for PostingIterator { impl PostingIterator { #[inline] fn block_least_doc_id(&self, list: &CompressedPostingList, block_idx: usize) -> u32 { - let compressed = unsafe { &mut *self.compressed_state_ptr() }; - if compressed.block_least_doc_ids.len() < list.blocks.len() { - compressed - .block_least_doc_ids - .resize(list.blocks.len(), None); - } - if let Some(doc_id) = compressed.block_least_doc_ids[block_idx] { - return doc_id; - } - let doc_id = list.block_least_doc_id(block_idx); - compressed.block_least_doc_ids[block_idx] = Some(doc_id); - doc_id + list.block_least_doc_id(block_idx) } fn block_idx_for_doc( @@ -406,6 +395,30 @@ impl PostingIterator { left - 1 } + #[inline] + fn block_end_doc(&self) -> u64 { + self.next_block_first_doc() + .map(|doc| doc.saturating_sub(1)) + .unwrap_or(TERMINATED_DOC_ID) + } + + /// Level1 bound of the group holding the current block, for group-wide + /// skipping: (group doc_up_to, group max score). `None` when the list has + /// no impact skip data or the group entry is missing/malformed. + fn impact_group_bound(&self, scorer: &S) -> Option<(u64, f32)> { + match self.list { + PostingList::Compressed(ref list) => { + let impacts = list.impacts.as_ref()?; + let (doc_up_to, score) = self.impact_level1(impacts, scorer); + if doc_up_to == u32::MAX { + return None; + } + Some((u64::from(doc_up_to), score)) + } + PostingList::Plain(_) => None, + } + } + #[inline] fn compressed_state_ptr(&self) -> *mut CompressedState { debug_assert!(self.compressed.is_some()); @@ -454,8 +467,15 @@ impl PostingIterator { list: PostingList, num_doc: usize, ) -> Self { + // BM25's doc weight is bounded by K1 + 1 for any freq and doc length, + // so query_weight * (K1 + 1) is a valid global bound even when index + // stats drift after appends. Keeping it finite matters: an INFINITY + // bound can never park the iterator in the WAND tail, forcing a deep + // advance on every candidate. let approximate_upper_bound = match &list { - PostingList::Compressed(posting) if posting.impacts.is_some() => f32::INFINITY, + PostingList::Compressed(posting) if posting.impacts.is_some() => { + query_weight * (K1 + 1.0) + } _ => match list.max_score() { Some(max_score) => max_score, None => idf(list.len(), num_doc) * (K1 + 1.0), @@ -502,9 +522,8 @@ impl PostingIterator { /// Tightest known list-wide score bound. Impact lists answer from the /// baked doc-weight slab (the data-driven equivalent of the max_score the /// non-impact format bakes at build time); everything else falls back to - /// `approximate_upper_bound`. A finite, tight global bound is what lets - /// lagging iterators park in the WAND tail instead of being force-advanced - /// on every candidate. + /// `approximate_upper_bound`. A finite, tight global bound lets lagging + /// iterators park in the WAND tail instead of being force-advanced. #[inline] fn global_upper_bound(&self, scorer: &S) -> f32 { if let PostingList::Compressed(ref list) = self.list @@ -682,26 +701,55 @@ impl PostingIterator { } } + /// Anchored level0 impact bound of the current block: (doc_up_to, max_score), + /// memoized until the block cursor moves. Malformed entries degrade to + /// (u32::MAX, INFINITY), which keeps pruning safe by making the block look + /// unskippable. + #[inline] + fn impact_level0( + &self, + impacts: &ImpactSkipData, + scorer: &S, + ) -> (u32, f32) { + let compressed = unsafe { &mut *self.compressed_state_ptr() }; + if let Some((block_idx, doc_up_to, score)) = compressed.level0_cache + && block_idx == self.block_idx + { + return (doc_up_to, score); + } + let doc_up_to = impacts.level0_doc_up_to(self.block_idx).unwrap_or(u32::MAX); + let score = impacts.level0_score(self.block_idx, self.query_weight, scorer); + compressed.level0_cache = Some((self.block_idx, doc_up_to, score)); + (doc_up_to, score) + } + + /// Anchored level1 impact bound of the group holding the current block, + /// memoized until the cursor crosses a group boundary. + #[inline] + fn impact_level1( + &self, + impacts: &ImpactSkipData, + scorer: &S, + ) -> (u32, f32) { + let group_idx = self.block_idx / IMPACT_LEVEL1_BLOCKS; + let compressed = unsafe { &mut *self.compressed_state_ptr() }; + if let Some((cached_group_idx, doc_up_to, score)) = compressed.level1_cache + && cached_group_idx == group_idx + { + return (doc_up_to, score); + } + let doc_up_to = impacts.level1_doc_up_to(group_idx).unwrap_or(u32::MAX); + let score = impacts.level1_score(group_idx, self.query_weight, scorer); + compressed.level1_cache = Some((group_idx, doc_up_to, score)); + (doc_up_to, score) + } + #[inline] fn block_max_score(&self, scorer: &S) -> f32 { match self.list { PostingList::Compressed(ref list) => { if let Some(impacts) = list.impacts.as_ref() { - let compressed = unsafe { &mut *self.compressed_state_ptr() }; - if let Some((block_idx, score)) = compressed.current_block_max_score - && block_idx == self.block_idx - { - return score; - } - - let score = impacts.level0_score_cached( - self.block_idx, - self.query_weight, - scorer, - &mut compressed.block_max_window.impact_score_cache, - ); - compressed.current_block_max_score = Some((self.block_idx, score)); - return score; + return self.impact_level0(impacts, scorer).1; } list.block_max_score(self.block_idx) } @@ -709,14 +757,27 @@ impl PostingIterator { } } + /// Tight max-score bound over `[current block, up_to]`. The common case — + /// a window ending inside the current block — answers from the anchored + /// level0 memo; wider windows fall back to the sliding block-max deque, + /// which scores each block once as it slides forward. #[inline] fn block_max_score_up_to_with_stats( - &mut self, + &self, up_to: u64, scorer: &S, ) -> BlockMaxScore { match self.list { PostingList::Compressed(ref list) => { + if let Some(impacts) = list.impacts.as_ref() { + let (level0_up_to, level0_score) = self.impact_level0(impacts, scorer); + if up_to <= u64::from(level0_up_to) { + return BlockMaxScore { + score: level0_score, + blocks_scanned: 0, + }; + } + } let compressed = unsafe { &mut *self.compressed_state_ptr() }; compressed.block_max_window.max_score_up_to( list, @@ -733,6 +794,16 @@ impl PostingIterator { } } + fn window_max_score(&self, up_to: Option, scorer: &S) -> f32 { + if let Some(up_to) = up_to + && let PostingList::Compressed(ref list) = self.list + && list.impacts.is_some() + { + return self.block_max_score_up_to_with_stats(up_to, scorer).score; + } + self.block_max_score(scorer) + } + #[inline] fn is_compressed(&self) -> bool { matches!(self.list, PostingList::Compressed(_)) @@ -755,6 +826,87 @@ impl PostingIterator { } } + /// Bulk-score every posting in `[current doc, up_to]` into the window + /// accumulator (slot = doc - window_min) and leave the iterator on the + /// first doc beyond `up_to`. This is the Lucene `nextDocsAndScores` + /// equivalent: it walks the decompressed block arrays directly, with no + /// per-doc heap traffic. + fn collect_window_scores( + &mut self, + window_min: u64, + up_to: u64, + clause_idx: usize, + docs: &DocSet, + scorer: &S, + acc: &mut WindowAccumulator, + ) { + if self.doc().is_some_and(|doc| doc.doc_id() < window_min) { + self.next(window_min); + } + match self.list { + PostingList::Compressed(ref list) => { + 'blocks: while let Some(doc) = self.current_doc { + if doc.doc_id() > up_to { + break; + } + let block_idx = posting_block_idx(self.index, list.block_size); + let block_offset = posting_block_offset(self.index, list.block_size); + let compressed = + unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; + for offset in block_offset..compressed.doc_ids.len() { + let doc_id = compressed.doc_ids[offset]; + if u64::from(doc_id) > up_to { + self.index = posting_block_start(block_idx, list.block_size) + offset; + self.block_idx = block_idx; + self.current_doc = Some(DocInfo::Raw(RawDocInfo { + doc_id, + frequency: compressed.freqs[offset], + })); + 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); + let slot = (u64::from(doc_id) - window_min) as usize; + acc.add(clause_idx, slot, score, freq); + } + // Block exhausted: step into the next block (or finish). + let next_start = posting_block_start(block_idx + 1, list.block_size); + if next_start >= list.length as usize { + self.index = list.length as usize; + self.block_idx = posting_block_idx(self.index, list.block_size); + self.current_doc = None; + break; + } + self.index = next_start; + self.block_idx = block_idx + 1; + let compressed = + unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx + 1) }; + self.current_doc = Some(DocInfo::Raw(RawDocInfo { + doc_id: compressed.doc_ids[0], + frequency: compressed.freqs[0], + })); + } + } + PostingList::Plain(_) => { + while let Some(doc) = self.doc() { + let doc_id = doc.doc_id(); + if doc_id > up_to { + break; + } + let doc_length = match &doc { + DocInfo::Raw(raw) => docs.scoring_num_tokens(raw.doc_id), + DocInfo::Located(located) => docs.num_tokens_by_row_id(located.row_id), + }; + let score = self.score(scorer, doc.frequency(), doc_length); + let slot = (doc_id - window_min) as usize; + acc.add(clause_idx, slot, score, doc.frequency()); + self.next(doc_id + 1); + } + } + } + } + #[inline] fn next_block_first_doc(&self) -> Option { match self.list { @@ -769,6 +921,50 @@ impl PostingIterator { } } +/// Inner window span (in doc ids) of the bulk MAXSCORE path. Same as Lucene's +/// `MaxScoreBulkScorer.INNER_WINDOW_SIZE`. +const MAXSCORE_INNER_WINDOW: usize = 1 << 12; + +/// Per-window score/frequency accumulator for the bulk MAXSCORE path. Slot i +/// covers doc id `window_min + i`; `freqs` is laid out clause-major so accept +/// paths can recover (term, freq) pairs of the essential clauses. +struct WindowAccumulator { + scores: Vec, + freqs: Vec, + words: Vec, + num_clauses: usize, +} + +impl WindowAccumulator { + fn new(num_clauses: usize) -> Self { + Self { + scores: vec![0.0; MAXSCORE_INNER_WINDOW], + freqs: vec![0; num_clauses * MAXSCORE_INNER_WINDOW], + words: vec![0; MAXSCORE_INNER_WINDOW / 64], + num_clauses, + } + } + + #[inline] + fn add(&mut self, clause_idx: usize, slot: usize, score: f32, freq: u32) { + self.scores[slot] += score; + // Doc-major layout: one slot's clause frequencies share a cache line. + self.freqs[slot * self.num_clauses + clause_idx] = freq; + self.words[slot >> 6] |= 1u64 << (slot & 63); + } + + #[inline] + fn clause_freq(&self, clause_idx: usize, slot: usize) -> u32 { + self.freqs[slot * self.num_clauses + clause_idx] + } + + #[inline] + fn clear_slot(&mut self, slot: usize) { + self.scores[slot] = 0.0; + self.freqs[slot * self.num_clauses..(slot + 1) * self.num_clauses].fill(0); + } +} + /// How wand identified a candidate: either it already had the real /// row_id (DocSet carried row_ids), or only the partition-local /// doc_id (deferred-row_id path; the caller must resolve via @@ -1066,6 +1262,22 @@ impl<'a, S: Scorer> Wand<'a, S> { _ => {} } + // Top-k disjunctions over compressed lists can opt into the bulk + // MAXSCORE path (Lucene MaxScoreBulkScorer style): it streams whole + // blocks of the essential clauses into a window accumulator instead of + // advancing doc-at-a-time through a heap. + if *USE_MAXSCORE_SEARCH + && self.operator == Operator::Or + && params.phrase_slop.is_none() + && !self.head.is_empty() + && self + .head + .iter() + .all(|posting| posting.posting.is_compressed()) + { + return self.maxscore_search(params, mask, metrics); + } + // Deferred-row_id path: when the DocSet was built without // row_ids, wand emits candidates carrying just the // partition-local doc_id; the outer caller resolves them to @@ -1352,6 +1564,484 @@ impl<'a, S: Scorer> Wand<'a, S> { .collect()) } + /// Bulk MAXSCORE top-k disjunction, mirroring Lucene's MaxScoreBulkScorer. + /// + /// Per outer window (bounded by the essential clauses' block boundaries): + /// clauses are partitioned into a non-essential prefix — sorted by window + /// max score, as many as fit under the threshold — and the essential rest. + /// Essential clauses stream their postings into a window accumulator in + /// bulk; only accumulated candidates that could still beat the threshold + /// probe the non-essential clauses. No per-doc heap maintenance happens + /// anywhere on this path. + fn maxscore_search( + &mut self, + params: &FtsSearchParams, + mask: Arc, + metrics: &dyn MetricsCollector, + ) -> Result> { + struct MaxScoreClause { + posting: Box, + bound: f32, + prefix_bound: f32, + } + + #[inline] + fn float_sum_upper(sum: f32, num_terms: usize) -> f32 { + // Mirror of Lucene MathUtil.sumUpperBound: widen the float sum so + // it stays a true upper bound of the exact sum. + sum + sum.abs() * (num_terms as f32) * f32::EPSILON + } + + let limit = params.limit.unwrap_or(usize::MAX); + let docs_has_row_ids = self.docs.has_row_ids(); + let mut clauses = std::mem::take(&mut self.head) + .into_vec() + .into_iter() + .map(|head| MaxScoreClause { + posting: head.posting, + bound: 0.0, + prefix_bound: 0.0, + }) + .collect::>(); + + let mut acc = WindowAccumulator::new(clauses.len()); + let mut candidates: TopKHeap = + BinaryHeap::with_capacity(std::cmp::min(limit, BLOCK_SIZE * 10)); + 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. + let mut min_window_size = 1u64; + let mut num_windows = 0u64; + let mut prev_first_essential = 0usize; + + let mut window_min = clauses + .iter() + .filter_map(|clause| clause.posting.doc().map(|doc| doc.doc_id())) + .min() + .unwrap_or(TERMINATED_DOC_ID); + + while window_min < TERMINATED_DOC_ID { + clauses.retain(|clause| clause.posting.doc().is_some()); + if clauses.is_empty() { + break; + } + self.raise_to_shared_floor(params.wand_factor); + + // Window boundary from the previous window's essential clauses + // only: dense non-essential clauses must not fragment the window. + let first_window_lead = prev_first_essential.min(clauses.len() - 1); + let mut window_max = TERMINATED_DOC_ID; + for clause in &mut clauses { + let doc = clause + .posting + .doc() + .map(|doc| doc.doc_id()) + .expect("exhausted clauses were retained out"); + clause.posting.shallow_next(doc.max(window_min)); + } + for clause in &clauses[first_window_lead..] { + window_max = window_max.min(clause.posting.block_end_doc()); + } + if clauses.len() > 1 { + // Target at least 32 candidates per clause per window on + // average before shrinking windows back to block granularity. + if (num_comparisons as u64) < num_windows * 32 * clauses.len() as u64 { + min_window_size = (min_window_size * 2).min(MAXSCORE_INNER_WINDOW as u64); + } else { + min_window_size = 1; + } + window_max = window_max.max(window_min.saturating_add(min_window_size - 1)); + } + + for clause in &mut clauses { + let doc = clause + .posting + .doc() + .map(|doc| doc.doc_id()) + .expect("exhausted clauses were retained out"); + clause.bound = if doc > window_max { + 0.0 + } else { + clause + .posting + .block_max_score_up_to_with_stats(window_max, &self.scorer) + .score + }; + } + clauses.sort_unstable_by(|a, b| a.bound.total_cmp(&b.bound)); + let mut first_essential = 0; + let mut prefix = 0.0_f32; + if self.threshold > 0.0 { + for (i, clause) in clauses.iter_mut().enumerate() { + let widened = float_sum_upper(prefix + clause.bound, i + 1); + if widened > self.threshold { + break; + } + prefix += clause.bound; + clause.prefix_bound = prefix; + first_essential = i + 1; + } + } + prev_first_essential = first_essential; + num_windows += 1; + + if first_essential == clauses.len() { + // No clause combination inside this window can beat the + // threshold: skip it wholesale. + window_min = match window_max { + TERMINATED_DOC_ID => TERMINATED_DOC_ID, + max => max + 1, + }; + // Single live clause: instead of re-running the window + // machinery once per block, scan the baked per-block bounds + // for the next block that can beat the threshold. This is the + // slab form of Lucene's getSkipUpTo and turns dead stretches + // of a dominant term into a tight load-mul-compare loop. + if clauses.len() == 1 && window_min != TERMINATED_DOC_ID && self.threshold > 0.0 { + let posting = &clauses[0].posting; + if let PostingList::Compressed(ref list) = posting.list + && let Some(impacts) = list.impacts.as_ref() + { + let bounds = impacts.level0_doc_weight_bounds(&self.scorer); + let query_weight = posting.query_weight; + // Position by binary search on the first-doc slab: the + // deep cursor lags arbitrarily far behind during long + // skip runs, and walking from it re-scans the same + // blocks on every dead window. + let first_docs = list.block_first_docs(); + let mut block_idx = first_docs + .partition_point(|&first| u64::from(first) <= window_min) + .saturating_sub(1); + while block_idx < bounds.len() + && query_weight * bounds[block_idx] <= self.threshold + { + block_idx += 1; + } + window_min = if block_idx < bounds.len() { + window_min.max(u64::from(list.block_least_doc_id(block_idx))) + } else { + TERMINATED_DOC_ID + }; + } + } + continue; + } + + let total_non_essential_bound = if first_essential > 0 { + clauses[first_essential - 1].prefix_bound + } else { + 0.0 + }; + + // Single essential clause (the common case once the threshold is + // competitive): stream it directly against the non-essential + // prefix, skipping the accumulator entirely. + if first_essential + 1 == clauses.len() { + let (non_essential, essential) = clauses.split_at_mut(first_essential); + let posting = &mut essential[0].posting; + if posting.doc().is_some_and(|doc| doc.doc_id() < window_min) { + posting.next(window_min); + } + let essential_term = posting.term_index(); + let essential_weight = posting.query_weight; + + macro_rules! consider_candidate { + ($doc:expr, $freq:expr) => {{ + 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); + if !(self.threshold > 0.0 + && score + total_non_essential_bound <= self.threshold) + { + let row_id = if docs_has_row_ids { + self.docs.row_id(doc as u32) + } else { + doc + }; + let masked_out = docs_has_row_ids + && (row_id == RowAddress::TOMBSTONE_ROW || !mask.selected(row_id)); + if !masked_out { + let mut total = score; + let mut rejected = false; + for i in (0..non_essential.len()).rev() { + if self.threshold > 0.0 + && total + non_essential[i].prefix_bound <= self.threshold + { + rejected = true; + break; + } + let probe = &mut non_essential[i].posting; + if probe.doc().is_some_and(|d| d.doc_id() < doc) { + probe.next(doc); + } + if let Some(d) = probe.doc() + && d.doc_id() == doc + { + total += + probe.score(&self.scorer, d.frequency(), doc_length); + } + } + + // Match the classic path's emission rule: a + // candidate must beat the running threshold, + // 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 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(), + )); + } + } + if full { + candidates.pop(); + } + candidates.push(Reverse(( + ScoredDoc::new(row_id, total), + freqs, + doc_length, + doc, + ))); + if candidates.len() == limit { + let kth = candidates.peek().unwrap().0.0.score.0; + self.update_threshold(kth, params.wand_factor); + } + } + } + } + } + }}; + } + + match posting.list { + PostingList::Compressed(ref list) => { + 'stream: while let Some(cur) = posting.current_doc { + if cur.doc_id() > window_max { + break; + } + let block_idx = posting_block_idx(posting.index, list.block_size); + let block_offset = posting_block_offset(posting.index, list.block_size); + let compressed = unsafe { + &mut *posting.ensure_compressed_block_ptr(list, block_idx) + }; + for offset in block_offset..compressed.doc_ids.len() { + let doc_id = compressed.doc_ids[offset]; + if u64::from(doc_id) > window_max { + posting.index = + posting_block_start(block_idx, list.block_size) + offset; + posting.block_idx = block_idx; + posting.current_doc = Some(DocInfo::Raw(RawDocInfo { + doc_id, + frequency: compressed.freqs[offset], + })); + break 'stream; + } + consider_candidate!(u64::from(doc_id), compressed.freqs[offset]); + } + let next_start = posting_block_start(block_idx + 1, list.block_size); + if next_start >= list.length as usize { + posting.index = list.length as usize; + posting.block_idx = + posting_block_idx(posting.index, list.block_size); + posting.current_doc = None; + break; + } + posting.index = next_start; + posting.block_idx = block_idx + 1; + let compressed = unsafe { + &mut *posting.ensure_compressed_block_ptr(list, block_idx + 1) + }; + posting.current_doc = Some(DocInfo::Raw(RawDocInfo { + doc_id: compressed.doc_ids[0], + frequency: compressed.freqs[0], + })); + } + } + PostingList::Plain(_) => { + while let Some(cur) = posting.doc() { + let doc = cur.doc_id(); + if doc > window_max { + break; + } + consider_candidate!(doc, cur.frequency()); + posting.next(doc + 1); + } + } + } + + window_min = match window_max { + TERMINATED_DOC_ID => TERMINATED_DOC_ID, + max => max + 1, + }; + continue; + } + + // Stream the essential clauses through inner windows. + let mut inner_min = window_min; + loop { + let mut next_essential_doc = TERMINATED_DOC_ID; + for clause in &clauses[first_essential..] { + if let Some(doc) = clause.posting.doc() { + next_essential_doc = next_essential_doc.min(doc.doc_id()); + } + } + inner_min = inner_min.max(next_essential_doc); + if inner_min == TERMINATED_DOC_ID || inner_min > window_max { + break; + } + let inner_max = + window_max.min(inner_min.saturating_add(MAXSCORE_INNER_WINDOW as u64 - 1)); + + for (clause_idx, clause) in clauses.iter_mut().enumerate().skip(first_essential) { + clause.posting.collect_window_scores( + inner_min, + inner_max, + clause_idx, + self.docs, + &self.scorer, + &mut acc, + ); + } + + // Drain candidates in doc order, completing them with the + // non-essential clauses ordered by descending bound. + for word_idx in 0..acc.words.len() { + let mut word = acc.words[word_idx]; + if word == 0 { + continue; + } + acc.words[word_idx] = 0; + while word != 0 { + let bit = word.trailing_zeros() as usize; + word &= word - 1; + let slot = (word_idx << 6) | bit; + let doc = inner_min + slot as u64; + let mut score = acc.scores[slot]; + num_comparisons += 1; + + if self.threshold > 0.0 + && score + total_non_essential_bound <= self.threshold + { + acc.clear_slot(slot); + continue; + } + + let row_id = if docs_has_row_ids { + self.docs.row_id(doc as u32) + } else { + doc + }; + if docs_has_row_ids + && (row_id == RowAddress::TOMBSTONE_ROW || !mask.selected(row_id)) + { + acc.clear_slot(slot); + continue; + } + + let doc_length = self.docs.scoring_num_tokens(doc as u32); + let mut rejected = false; + for i in (0..first_essential).rev() { + if self.threshold > 0.0 + && score + clauses[i].prefix_bound <= self.threshold + { + rejected = true; + break; + } + let posting = &mut clauses[i].posting; + if posting.doc().is_some_and(|d| d.doc_id() < doc) { + posting.next(doc); + } + if let Some(d) = posting.doc() + && d.doc_id() == doc + { + score += posting.score(&self.scorer, d.frequency(), doc_length); + } + } + + 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 freqs = 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()) + }) + }) + } + }) + .collect::>(); + if full { + candidates.pop(); + } + candidates.push(Reverse(( + ScoredDoc::new(row_id, score), + freqs, + doc_length, + doc, + ))); + if candidates.len() == limit { + let kth = candidates.peek().unwrap().0.0.score.0; + self.update_threshold(kth, params.wand_factor); + } + } + } + acc.clear_slot(slot); + } + } + if inner_max >= window_max { + break; + } + inner_min = inner_max + 1; + } + + window_min = match window_max { + TERMINATED_DOC_ID => TERMINATED_DOC_ID, + max => max + 1, + }; + } + + metrics.record_comparisons(num_comparisons); + + let to_addr = |row_id_slot: u64| { + if docs_has_row_ids { + CandidateAddr::RowId(row_id_slot) + } else { + CandidateAddr::Pending(row_id_slot as u32) + } + }; + Ok(candidates + .into_iter() + .map( + |Reverse((doc, freqs, doc_length, posting_doc_id))| DocCandidate { + addr: to_addr(doc.row_id), + posting_doc_id, + freqs, + doc_length, + }, + ) + .collect()) + } + // calculate the score of the current document fn score(&self, doc_length: u32) -> f32 { let mut score = 0.0; @@ -1430,10 +2120,16 @@ impl<'a, S: Scorer> Wand<'a, S> { if self.threshold > 0.0 && self.or_block_window_max() <= self.threshold { // On the final block `up_to` is the `u64::MAX` sentinel; step once // there to avoid seeking past the valid doc id range. - let skip_to = match self.up_to { + let mut skip_to = match self.up_to { Some(up_to) if up_to < u32::MAX as u64 => up_to + 1, _ => target + 1, }; + // The narrow window is dead; if the whole level1 group is dead + // too, hop over it in one advance. + let group_skip = self.or_group_skip_to(); + if let Some(group_skip_to) = group_skip { + skip_to = skip_to.max(group_skip_to); + } self.push_back_leads(skip_to); continue; } @@ -1683,12 +2379,12 @@ impl<'a, S: Scorer> Wand<'a, S> { let lead: f32 = self .lead .iter() - .map(|posting| posting.block_max_score(&self.scorer)) + .map(|posting| posting.window_max_score(self.up_to, &self.scorer)) .sum(); let head: f32 = self .head .iter() - .map(|posting| posting.posting.block_max_score(&self.scorer)) + .map(|posting| posting.posting.window_max_score(self.up_to, &self.scorer)) .sum(); lead + head + self.tail_max_score } @@ -1701,12 +2397,12 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut sum = self .lead .iter() - .map(|posting| posting.block_max_score(&self.scorer)) + .map(|posting| posting.window_max_score(self.up_to, &self.scorer)) .sum::(); let mut possible_matches = self.lead.len(); for posting in &self.tail { if matches!(posting.posting.block_first_doc(), Some(block_doc) if block_doc <= target) { - sum += posting.posting.block_max_score(&self.scorer); + sum += posting.posting.window_max_score(self.up_to, &self.scorer); possible_matches += 1; } } @@ -1720,72 +2416,121 @@ impl<'a, S: Scorer> Wand<'a, S> { fn update_max_scores(&mut self, target: u64) { // Refresh the block-max window for the current target. The resulting // `up_to` is the furthest doc id for which this block-max view remains - // valid. + // valid. Like Lucene's WANDScorer, the boundary comes from the cheap + // clauses only, and the refresh avoids allocating: heaps are recycled + // through their backing vectors (shallow_next never changes the doc a + // head entry is ordered by, so heapify restores the same shape). let lead_cost = self .lead .iter() .map(|posting| posting.cost()) .min() .unwrap_or(usize::MAX); - let mut up_to = TERMINATED_DOC_ID; + let mut narrow_up_to = TERMINATED_DOC_ID; for posting in &mut self.lead { posting.shallow_next(target); - let block_end = posting - .next_block_first_doc() - .map(|doc| doc.saturating_sub(1)) - .unwrap_or(TERMINATED_DOC_ID); - up_to = up_to.min(block_end); - } - let head = std::mem::take(&mut self.head); - let mut rebuilt_head = BinaryHeap::with_capacity(head.len()); - for mut posting in head.into_vec() { - if posting.posting.cost() <= lead_cost { - posting.posting.shallow_next(posting.doc_id()); - let block_end = posting - .posting - .next_block_first_doc() - .map(|doc| doc.saturating_sub(1)) - .unwrap_or(TERMINATED_DOC_ID); - up_to = up_to.min(block_end); - } - rebuilt_head.push(posting); + narrow_up_to = narrow_up_to.min(posting.block_end_doc()); + } + + let mut head_postings = std::mem::take(&mut self.head).into_vec(); + for posting in &mut head_postings { + // Unlike Lucene, every head clause participates in the boundary: + // the refresh is allocation-free and answers from the anchored + // level caches, so frequent refreshes are cheap, while keeping the + // window inside every clause's current block keeps all the bounds + // at tight level0 values. + let doc_id = posting.doc_id(); + posting.posting.shallow_next(doc_id); + narrow_up_to = narrow_up_to.min(posting.posting.block_end_doc()); } - self.head = rebuilt_head; - if up_to == TERMINATED_DOC_ID - && let Some(top) = self.tail.peek() - && top.cost <= lead_cost + + let mut tail_postings = std::mem::take(&mut self.tail).into_vec(); + for tail_posting in &mut tail_postings { + tail_posting.posting.shallow_next(target); + } + + if narrow_up_to == TERMINATED_DOC_ID + && let Some(top) = tail_postings + .iter() + .min_by_key(|posting| posting.posting.cost()) + && top.posting.cost() <= lead_cost { - let block_end = top - .posting - .next_block_first_doc() - .map(|doc| doc.saturating_sub(1)) - .unwrap_or(TERMINATED_DOC_ID); - up_to = up_to.min(block_end.max(target)); + narrow_up_to = narrow_up_to.min(top.posting.block_end_doc().max(target)); } - self.up_to = Some(up_to); - let tail = std::mem::take(&mut self.tail); + self.up_to = Some(narrow_up_to); + self.head = BinaryHeap::from(head_postings); + self.tail_max_score = 0.0; - for mut tail_posting in tail.into_vec() { - tail_posting.posting.shallow_next(target); - let upper_bound = match tail_posting.posting.block_first_doc() { + for tail_posting in tail_postings { + let posting = tail_posting.posting; + let upper_bound = match posting.block_first_doc() { Some(block_doc) if block_doc <= target => { - tail_posting - .posting - .block_max_score_up_to_with_stats(up_to, &self.scorer) - .score + posting.window_max_score(self.up_to, &self.scorer) } _ => 0.0, }; - if let Some(mut evicted) = - self.insert_tail_with_overflow(tail_posting.posting, upper_bound) - { + if let Some(mut evicted) = self.insert_tail_with_overflow(posting, upper_bound) { evicted.next(target); self.push_head(evicted); } } } + /// After the narrow window proved skippable, try widening the skip to the + /// level1 group boundary, in the spirit of Lucene's `getSkipUpTo`. All + /// bounds come from the anchored level caches, so a failed attempt costs a + /// few loads and float adds — unlike an eager wide-window probe. + /// + /// The group boundary is the minimum current-group end over every live + /// iterator, so each iterator's level1 score is a valid bound over + /// `[target, group_up_to]`. + fn or_group_skip_to(&self) -> Option { + let mut group_up_to = TERMINATED_DOC_ID; + for posting in &self.lead { + let (doc_up_to, _) = posting.impact_group_bound(&self.scorer)?; + group_up_to = group_up_to.min(doc_up_to); + } + for posting in self.head.iter() { + let (doc_up_to, _) = posting.posting.impact_group_bound(&self.scorer)?; + group_up_to = group_up_to.min(doc_up_to); + } + for tail_posting in self.tail.iter() { + let (doc_up_to, _) = tail_posting.posting.impact_group_bound(&self.scorer)?; + group_up_to = group_up_to.min(doc_up_to); + } + if self.up_to.is_some_and(|up_to| group_up_to <= up_to) { + // No gain over the narrow window skip. + return None; + } + + // Second pass over the memoized bounds: sum only the iterators that + // can produce a doc inside the skipped range. + let mut bounds_sum = 0.0_f32; + for posting in &self.lead { + let (_, score) = posting.impact_group_bound(&self.scorer)?; + bounds_sum += score; + } + for posting in self.head.iter() { + if posting.doc_id() > group_up_to { + continue; + } + let (_, score) = posting.posting.impact_group_bound(&self.scorer)?; + bounds_sum += score; + } + for tail_posting in self.tail.iter() { + if !matches!( + tail_posting.posting.block_first_doc(), + Some(block_doc) if block_doc <= group_up_to + ) { + continue; + } + let (_, score) = tail_posting.posting.impact_group_bound(&self.scorer)?; + bounds_sum += score; + } + (bounds_sum <= self.threshold).then_some(group_up_to.saturating_add(1)) + } + fn refine_or_candidate(&mut self, target: u64, doc_length: u32) -> bool { if self.threshold <= 0.0 { return true; @@ -1909,7 +2654,7 @@ impl<'a, S: Scorer> Wand<'a, S> { && posting.is_compressed() && self.up_to.is_some_and(|up_to| target <= up_to) { - posting.block_max_score(&self.scorer) + posting.window_max_score(self.up_to, &self.scorer) } else { posting.global_upper_bound(&self.scorer) } @@ -3363,6 +4108,75 @@ mod tests { } } + #[test] + fn test_or_impact_level1_window_skips_low_group_with_single_score() { + let total = (IMPACT_LEVEL1_BLOCKS + 1) * BLOCK_SIZE; + let target = (IMPACT_LEVEL1_BLOCKS * BLOCK_SIZE) as u64; + let mut docs = DocSet::default(); + for doc_id in 0..total as u64 { + docs.append(doc_id, 1); + } + + let doc_ids = (0..total as u32).collect::>(); + let freqs = doc_ids + .iter() + .map(|doc_id| if u64::from(*doc_id) < target { 1 } else { 10 }) + .collect::>(); + let posting_list = generate_impact_posting_list_with_freqs(doc_ids, freqs, vec![1; total]); + let mut probe = PostingIterator::with_query_weight( + String::from("term"), + 0, + 0, + 1.0, + posting_list.clone(), + docs.len(), + ); + probe.shallow_next(0); + let counting_scorer = CountingScorer { + scored: Arc::new(AtomicUsize::new(0)), + }; + let (group_up_to, group_score) = probe.impact_group_bound(&counting_scorer).unwrap(); + assert_eq!(group_up_to, target - 1); + assert_eq!(group_score, 1.0); + // A window ending past the current block must answer from level1. + assert_eq!( + probe.window_max_score(Some(target - 1), &counting_scorer), + 1.0 + ); + + let posting = PostingIterator::with_query_weight( + String::from("term"), + 0, + 0, + 1.0, + posting_list, + docs.len(), + ); + let scored = Arc::new(AtomicUsize::new(0)); + let mut wand = Wand::new( + Operator::Or, + std::iter::once(posting), + &docs, + CountingScorer { + scored: scored.clone(), + }, + ); + wand.threshold = 2.0; + + let (candidate, score) = wand.next().unwrap().unwrap(); + assert_eq!(candidate.doc_id(), target); + assert_eq!(score, 10.0); + // The doc-weight bounds bake exactly once (one doc_weight call per + // frontier pair across all entries); beyond that only the returned + // candidate is scored. + let total_entries = (IMPACT_LEVEL1_BLOCKS + 1) + 2; + assert!( + scored.load(Ordering::Relaxed) <= total_entries + 8, + "bounds should be baked once instead of recomputed per window; scored={}", + scored.load(Ordering::Relaxed) + ); + } + #[test] fn test_compressed_impact_block_max_score_memoizes_current_block() { let total = 2 * BLOCK_SIZE as u32; @@ -3382,14 +4196,16 @@ mod tests { let first_score = posting.block_max_score(&scorer); assert_eq!(first_score, 1.0); - // Baking the shared doc-weight bounds visits every frontier pair once - // (two level0 entries plus one level1 entry for this list). + // Baking the doc-weight bounds visits every frontier pair once: two + // level0 entries plus one level1 entry for this list. let baked = scored.load(Ordering::Relaxed); assert!(baked >= 2); { let compressed = unsafe { &mut *posting.compressed_state_ptr() }; - assert_eq!(compressed.current_block_max_score, Some((0, first_score))); - compressed.block_max_window.impact_score_cache = ImpactScoreCache::default(); + assert_eq!( + compressed.level0_cache, + Some((0, BLOCK_SIZE as u32 - 1, first_score)) + ); } let second_score = posting.block_max_score(&scorer); From fec0a88f604828d40ec2d7d80c29dab60d5f4287 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 3 Jul 2026 13:13:21 +0800 Subject: [PATCH 4/7] perf(fts): enable the bulk MAXSCORE path by default With right-sized partitions the bulk path reaches Lucene-parity latency (k10 0.034-0.037s/213-235qps vs Lucene 10.4's 0.037s/216qps on the same 200M-doc corpus, queries, and warm protocol) and its results are score-identical to the classic WAND loop. LANCE_FTS_MAXSCORE=0 opts back into the classic loop. --- rust/lance-index/src/scalar/inverted/wand.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 70ee2015cda..253e1b0b9be 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -50,10 +50,11 @@ pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { .unwrap_or(10) }); // Bulk MAXSCORE path for top-k disjunctions (Lucene MaxScoreBulkScorer -// style). Off by default while it bakes; LANCE_FTS_MAXSCORE=1 opts in. Its -// results are score-identical to the classic WAND loop. +// style). Default on: with right-sized partitions it wins by a wide margin +// (Lucene-parity latency) and its results are score-identical to the classic +// WAND loop. LANCE_FTS_MAXSCORE=0 opts back into the classic loop. static USE_MAXSCORE_SEARCH: LazyLock = - LazyLock::new(|| std::env::var("LANCE_FTS_MAXSCORE").as_deref() == Ok("1")); + LazyLock::new(|| std::env::var("LANCE_FTS_MAXSCORE").as_deref() != Ok("0")); #[inline] fn posting_block_idx(index: usize, block_size: usize) -> usize { From 22e3522a80bba89fb9fb3bc7bdabb754b3d9a02f Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 3 Jul 2026 21:18:00 +0800 Subject: [PATCH 5/7] perf(index): bulk fts conjunction path and per-doc position seeking AND and phrase queries previously leapfrogged doc-at-a-time through boxed PostingIterator::next calls (~61% of the AND profile) and phrase checks decoded a whole 256-doc position block per candidate (~39% of the phrase profile). - and_bulk_search: block-max window pruning plus a k-pointer merge over decompressed block slices; per-candidate advance cost drops to a few loads. Results are identical to the classic loop (LANCE_FTS_BULK_AND=0 opts out). Phrase queries ride the same path. - seek_packed_doc_positions: PackedDelta full groups are self-describing ([num_bits u8][16*num_bits bytes]), so group offsets are recovered by hopping headers; decode only the 1-2 groups overlapping the candidate doc's delta range, with a lazily-built group index, memoized unpacked group, and a decoded-tail cache per block. - check_exact_positions_bulk: allocation-free slop=0 alignment check on the decoded scratch slices for parked lead clauses. Warm mmlb benchmarks, 8 concurrent queries: AND\@200M k10 0.114->0.060s, k100 0.240->0.118s; phrase\@50M 3-word k10 0.335->0.210s, 2-word k10 0.098->0.042s. All steps verified score-identical to the classic path. Co-Authored-By: Claude Fable 5 --- rust/lance-index/src/scalar/inverted.rs | 2 +- .../src/scalar/inverted/encoding.rs | 166 +++++ rust/lance-index/src/scalar/inverted/wand.rs | 607 +++++++++++++++++- 3 files changed, 749 insertions(+), 26 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 0ba7e52e530..c7a3cdade51 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -2,9 +2,9 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors pub mod builder; -mod impact; mod cache_codec; mod encoding; +mod impact; mod index; mod iter; pub mod json; diff --git a/rust/lance-index/src/scalar/inverted/encoding.rs b/rust/lance-index/src/scalar/inverted/encoding.rs index 6296de077c0..131b54dcad4 100644 --- a/rust/lance-index/src/scalar/inverted/encoding.rs +++ b/rust/lance-index/src/scalar/inverted/encoding.rs @@ -779,6 +779,98 @@ fn decode_position_stream_packed_block( Ok(()) } +/// Decode one document's positions out of a PackedDelta position block +/// without decoding the rest of the block. Full 128-delta groups are +/// self-describing (`[num_bits u8][16 * num_bits packed bytes]`), so group +/// byte offsets are recovered by hopping headers — no format change is +/// involved. `delta_range` is the doc's range in the block-wide delta stream +/// (from the frequency prefix sums); per-doc deltas reset at document +/// boundaries, so decoding starts cleanly at `delta_range.start`. +/// +/// The caller passes per-block scratch state that this function maintains: +/// `group_offsets` (lazily extended header index, seeded with `[0]`), +/// `unpacked_group`/`unpacked_group_idx` (the last unpacked group), and +/// `tail_cache` (the varint tail, decoded in full on first touch). All are +/// reset by the caller when the block cursor moves. +#[allow(clippy::too_many_arguments)] +pub(super) fn seek_packed_doc_positions( + src: &[u8], + total_deltas: usize, + delta_range: std::ops::Range, + group_offsets: &mut Vec, + unpacked_group: &mut [u32; BLOCK_SIZE], + unpacked_group_idx: &mut Option, + tail_cache: &mut Vec, + dst: &mut Vec, +) -> Result<()> { + dst.clear(); + if delta_range.is_empty() { + return Ok(()); + } + let num_full_groups = total_deltas / BLOCK_SIZE; + let packed_deltas_end = num_full_groups * BLOCK_SIZE; + + // Extend the header index far enough for this range (tail needs the + // offset one past the last full group). + let last_needed_group = if delta_range.end > packed_deltas_end { + num_full_groups + } else { + (delta_range.end - 1) / BLOCK_SIZE + }; + while group_offsets.len() <= last_needed_group { + let last = *group_offsets + .last() + .expect("group offset index is seeded with the first group"); + let num_bits = *src.get(last).ok_or_else(|| { + Error::index("unexpected EOF while seeking packed position stream".to_owned()) + })? as usize; + group_offsets.push(last + 1 + num_bits * BLOCK_SIZE / 8); + } + + let mut previous = 0u32; + let mut first = true; + let mut push_delta = |delta: u32, dst: &mut Vec| -> Result<()> { + let position = if first { + first = false; + delta + } else { + previous + .checked_add(delta) + .ok_or_else(|| Error::index("position stream overflow while decoding".to_owned()))? + }; + dst.push(position); + previous = position; + Ok(()) + }; + + for index in delta_range.start..delta_range.end.min(packed_deltas_end) { + let group = index / BLOCK_SIZE; + if *unpacked_group_idx != Some(group) { + let offset = group_offsets[group]; + let num_bits = src[offset]; + BitPacker4x::new().decompress(&src[offset + 1..], unpacked_group, num_bits); + *unpacked_group_idx = Some(group); + } + push_delta(unpacked_group[index % BLOCK_SIZE], dst)?; + } + + if delta_range.end > packed_deltas_end { + let tail_len = total_deltas - packed_deltas_end; + if tail_cache.len() != tail_len { + tail_cache.clear(); + tail_cache.reserve(tail_len); + let mut offset = group_offsets[num_full_groups]; + for _ in 0..tail_len { + tail_cache.push(decode_varint_u32(src, &mut offset)?); + } + } + for index in delta_range.start.max(packed_deltas_end)..delta_range.end { + push_delta(tail_cache[index - packed_deltas_end], dst)?; + } + } + Ok(()) +} + #[cfg(test)] pub fn encode_position_stream_block_into( positions: &[u32], @@ -1243,6 +1335,80 @@ mod tests { Ok(()) } + /// Per-doc seek decoding of a PackedDelta position block must return + /// exactly the same positions as decoding the whole block, for every doc, + /// across group-boundary-straddling docs and varint tails. + #[test] + fn test_packed_position_doc_seek_matches_block_decode() -> Result<()> { + let mut rng = rand::rng(); + // Frequency shapes: tiny blocks (tail only), exactly one group, a doc + // straddling group boundaries, and a large multi-group block. + let freq_shapes: Vec> = vec![ + vec![1], + vec![3, 1, 5], + vec![64, 64], + vec![100, 60, 40], + (0..256u32).map(|i| (i % 7) + 1).collect(), + vec![300, 2, 129, 1, 77], + ]; + for frequencies in freq_shapes { + let total: usize = frequencies.iter().map(|&f| f as usize).sum(); + // Positions ascend within each doc; docs are independent. + let mut positions = Vec::with_capacity(total); + for &freq in &frequencies { + let mut current = rng.random_range(0..1000u32); + for _ in 0..freq { + positions.push(current); + current += rng.random_range(1..50u32); + } + } + + let mut encoded = Vec::new(); + encode_position_stream_block_into( + &positions, + &frequencies, + PositionStreamCodec::PackedDelta, + &mut encoded, + )?; + + let mut whole = Vec::new(); + decode_position_stream_block( + &encoded, + &frequencies, + PositionStreamCodec::PackedDelta, + &mut whole, + )?; + assert_eq!(whole, positions); + + let mut group_offsets = vec![0usize]; + let mut unpacked_group = Box::new([0u32; BLOCK_SIZE]); + let mut unpacked_group_idx = None; + let mut tail_cache = Vec::new(); + let mut scratch = Vec::new(); + let mut delta_start = 0usize; + for &freq in &frequencies { + let delta_end = delta_start + freq as usize; + seek_packed_doc_positions( + &encoded, + total, + delta_start..delta_end, + &mut group_offsets, + &mut unpacked_group, + &mut unpacked_group_idx, + &mut tail_cache, + &mut scratch, + )?; + assert_eq!( + scratch, + whole[delta_start..delta_end], + "doc positions mismatch for range {delta_start}..{delta_end} freqs={frequencies:?}" + ); + delta_start = delta_end; + } + } + Ok(()) + } + #[test] fn test_encode_position_stream_block_roundtrip() -> Result<()> { let frequencies = vec![1, 3, 2, 4]; diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 253e1b0b9be..01525fc2832 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -23,6 +23,7 @@ use crate::metrics::MetricsCollector; use super::{ CompressedPositionStorage, impact::{IMPACT_LEVEL1_BLOCKS, ImpactSkipData}, + index::PositionStreamCodec, query::Operator, scorer::{K1, idf}, }; @@ -31,7 +32,7 @@ use super::{ builder::ScoredDoc, encoding::{ MAX_POSTING_BLOCK_SIZE, decode_position_stream_block, decompress_positions, - decompress_posting_block, decompress_posting_remainder, + decompress_posting_block, decompress_posting_remainder, seek_packed_doc_positions, }, query::FtsSearchParams, scorer::Scorer, @@ -55,6 +56,12 @@ pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { // WAND loop. LANCE_FTS_MAXSCORE=0 opts back into the classic loop. static USE_MAXSCORE_SEARCH: LazyLock = LazyLock::new(|| std::env::var("LANCE_FTS_MAXSCORE").as_deref() != Ok("0")); +// Bulk conjunction path for top-k AND / phrase queries: block-max window +// skipping plus a slice-level merge over decompressed blocks, replacing the +// per-doc `next()` leapfrog. Results are identical to the classic AND loop. +// LANCE_FTS_BULK_AND=0 opts back into the classic loop. +static USE_BULK_AND_SEARCH: LazyLock = + LazyLock::new(|| std::env::var("LANCE_FTS_BULK_AND").as_deref() != Ok("0")); #[inline] fn posting_block_idx(index: usize, block_size: usize) -> usize { @@ -109,6 +116,17 @@ struct CompressedState { position_block_idx: Option, position_values: Vec, position_offsets: Vec, + // Seek state for PackedDelta position blocks: the lazily-built group + // header index, the last unpacked group (memoized), the decoded varint + // tail, the block's total delta count, and the scratch the current doc's + // positions land in. Together these let a phrase check decode just the + // candidate doc's positions instead of the whole 256-doc position block. + position_group_offsets: Vec, + position_unpacked_group: Box<[u32; BLOCK_SIZE]>, + position_unpacked_group_idx: Option, + position_tail: Vec, + position_total_deltas: usize, + position_doc_scratch: Vec, block_max_window: BlockMaxWindow, // Lucene-style anchored impact score caches: one slot per level, keyed by // the entry the block cursor currently sits in. Each holds @@ -127,6 +145,12 @@ impl CompressedState { position_block_idx: None, position_values: Vec::new(), position_offsets: Vec::new(), + position_group_offsets: Vec::new(), + position_unpacked_group: Box::new([0; BLOCK_SIZE]), + position_unpacked_group_idx: None, + position_tail: Vec::new(), + position_total_deltas: 0, + position_doc_scratch: Vec::new(), block_max_window: BlockMaxWindow::new(), level0_cache: None, level1_cache: None, @@ -607,32 +631,79 @@ impl PostingIterator { let block_offset = posting_block_offset(self.index, list.block_size); let compressed = unsafe { &mut *self.ensure_compressed_block_ptr(list, block_idx) }; - if compressed.position_block_idx != Some(block_idx) { - decode_position_stream_block( - stream.block(block_idx), - compressed.freqs.as_slice(), - stream.codec(), - &mut compressed.position_values, - ) - .expect("shared position stream decoding should succeed"); - compressed.position_offsets.clear(); - compressed - .position_offsets - .reserve(compressed.freqs.len() + 1); - compressed.position_offsets.push(0); - let mut offset = 0usize; - for &freq in &compressed.freqs { - offset += freq as usize; - compressed.position_offsets.push(offset); + match stream.codec() { + PositionStreamCodec::PackedDelta => { + // Seekable layout: decode only the candidate doc's + // positions. Per-block seek state resets when the + // block cursor moves; the group header index and + // varint tail fill in lazily as candidates touch + // them. + if compressed.position_block_idx != Some(block_idx) { + compressed.position_group_offsets.clear(); + compressed.position_group_offsets.push(0); + compressed.position_tail.clear(); + compressed.position_unpacked_group_idx = None; + compressed.position_offsets.clear(); + compressed + .position_offsets + .reserve(compressed.freqs.len() + 1); + compressed.position_offsets.push(0); + let mut offset = 0usize; + for &freq in &compressed.freqs { + offset += freq as usize; + compressed.position_offsets.push(offset); + } + compressed.position_total_deltas = offset; + compressed.position_block_idx = Some(block_idx); + } + let delta_start = compressed.position_offsets[block_offset]; + let delta_end = compressed.position_offsets[block_offset + 1]; + seek_packed_doc_positions( + stream.block(block_idx), + compressed.position_total_deltas, + delta_start..delta_end, + &mut compressed.position_group_offsets, + &mut compressed.position_unpacked_group, + &mut compressed.position_unpacked_group_idx, + &mut compressed.position_tail, + &mut compressed.position_doc_scratch, + ) + .expect("shared position stream doc decoding should succeed"); + Some(PositionCursor::new( + PositionValues::Borrowed(&compressed.position_doc_scratch[..]), + self.position as i32, + )) + } + PositionStreamCodec::VarintDocDelta => { + if compressed.position_block_idx != Some(block_idx) { + compressed.position_values.clear(); + decode_position_stream_block( + stream.block(block_idx), + compressed.freqs.as_slice(), + stream.codec(), + &mut compressed.position_values, + ) + .expect("shared position stream decoding should succeed"); + compressed.position_offsets.clear(); + compressed + .position_offsets + .reserve(compressed.freqs.len() + 1); + compressed.position_offsets.push(0); + let mut offset = 0usize; + for &freq in &compressed.freqs { + offset += freq as usize; + compressed.position_offsets.push(offset); + } + compressed.position_block_idx = Some(block_idx); + } + let start = compressed.position_offsets[block_offset]; + let end = compressed.position_offsets[block_offset + 1]; + Some(PositionCursor::new( + PositionValues::Borrowed(&compressed.position_values[start..end]), + self.position as i32, + )) } - compressed.position_block_idx = Some(block_idx); } - let start = compressed.position_offsets[block_offset]; - let end = compressed.position_offsets[block_offset + 1]; - Some(PositionCursor::new( - PositionValues::Borrowed(&compressed.position_values[start..end]), - self.position as i32, - )) } }, } @@ -1138,6 +1209,9 @@ pub struct Wand<'a, S: Scorer> { and_last_doc: Option, and_window_stats: AndWindowStats, and_candidates_pruned_before_return: usize, + // Test-only escape hatch to run the classic conjunction loop even when the + // bulk path is enabled process-wide. + bulk_and_disabled: bool, docs: &'a DocSet, scorer: S, // Shared cross-partition top-k floor. Each partition publishes its local @@ -1202,12 +1276,21 @@ impl<'a, S: Scorer> Wand<'a, S> { and_last_doc: None, and_window_stats: AndWindowStats::default(), and_candidates_pruned_before_return: 0, + bulk_and_disabled: false, docs, scorer, shared_threshold: None, } } + /// Test hook: force the classic doc-at-a-time AND loop so parity tests can + /// compare it against the bulk conjunction path within one process. + #[cfg(test)] + pub(crate) fn disable_bulk_and(mut self) -> Self { + self.bulk_and_disabled = true; + self + } + /// Share one cross-partition top-k floor across a query's partitions. pub(crate) fn with_shared_threshold(mut self, shared: Arc) -> Self { self.shared_threshold = Some(shared); @@ -1279,6 +1362,19 @@ impl<'a, S: Scorer> Wand<'a, S> { return self.maxscore_search(params, mask, metrics); } + // Top-k conjunctions (AND and phrase) over compressed lists use the + // bulk path: the same block-max window pruning, but candidates come + // from a slice-level merge over decompressed blocks instead of per-doc + // `next()` leapfrogging through boxed iterators. + if *USE_BULK_AND_SEARCH + && !self.bulk_and_disabled + && self.operator == Operator::And + && !self.lead.is_empty() + && self.lead.iter().all(|posting| posting.is_compressed()) + { + return self.and_bulk_search(params, mask, metrics); + } + // Deferred-row_id path: when the DocSet was built without // row_ids, wand emits candidates carrying just the // partition-local doc_id; the outer caller resolves them to @@ -2248,6 +2344,306 @@ impl<'a, S: Scorer> Wand<'a, S> { .max(target) } + /// Bulk conjunction search. The window ends at the nearest next-block + /// boundary across the clauses, so within a window every clause + /// contributes exactly one decompressed block and the intersection is a + /// plain merge over `u32` slices — the per-candidate cost drops from a + /// full `PostingIterator::next` call per clause to a couple of loads. + /// Window skipping, per-candidate pruning, scoring, and heap semantics + /// mirror the classic loop exactly, so results are identical. + fn and_bulk_search( + &mut self, + params: &FtsSearchParams, + mask: Arc, + metrics: &dyn MetricsCollector, + ) -> Result> { + let limit = params.limit.unwrap_or(usize::MAX); + if limit == 0 { + return Ok(vec![]); + } + let docs_has_row_ids = self.docs.has_row_ids(); + let num_lists = self.lead.len(); + let phrase_slop = params.phrase_slop; + + // Per-window view of one clause's current block. Raw pointers into the + // clause's `CompressedState`; valid for the whole window because the + // block cursor does not move within a window (position decoding writes + // to separate fields of the same state). + struct WindowList { + docs: *const u32, + freqs: *const u32, + // Cursor and exclusive end, as offsets within the block. + pos: usize, + end: usize, + // Absolute posting index of the block's first entry. + block_start: usize, + } + + let mut candidates: TopKHeap = + BinaryHeap::with_capacity(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, + ..Default::default() + }; + let mut wins: Vec = Vec::with_capacity(num_lists); + + // The conjunction can only start at the max of the clauses' first docs. + let mut target: u64 = 0; + for posting in &self.lead { + match posting.doc() { + Some(doc) => target = target.max(doc.doc_id()), + None => return Ok(vec![]), + } + } + + 'window: loop { + self.raise_to_shared_floor(params.wand_factor); + if self.threshold > 0.0 { + let advanced = self.and_advance_target(target); + if advanced == TERMINATED_DOC_ID { + break; + } + target = advanced; + } + debug_assert!(target <= u32::MAX as u64); + let target32 = target as u32; + + // Position every clause's block cursor at the block that can hold + // `target`, and end the window at the nearest next-block boundary. + let mut win_end = TERMINATED_DOC_ID; + for j in 0..num_lists { + let (block_idx, block_up_to) = { + let posting = &self.lead[j]; + let PostingList::Compressed(ref list) = posting.list else { + unreachable!("bulk AND requires compressed postings"); + }; + let block_idx = posting.block_idx_for_doc(list, posting.block_idx, target32); + let block_up_to = if block_idx + 1 < list.blocks.len() { + u64::from(posting.block_least_doc_id(list, block_idx + 1)).saturating_sub(1) + } else { + TERMINATED_DOC_ID + }; + (block_idx, block_up_to.max(target)) + }; + self.lead[j].block_idx = block_idx; + win_end = win_end.min(block_up_to); + } + let win_end32 = u32::try_from(win_end).unwrap_or(u32::MAX); + + // Decompress each clause's block and slice it to [target, win_end]. + wins.clear(); + let mut skip_window = false; + let mut exhausted = false; + for posting in &self.lead { + let PostingList::Compressed(ref list) = posting.list else { + unreachable!("bulk AND requires compressed postings"); + }; + let block_idx = posting.block_idx; + let state = unsafe { &mut *posting.ensure_compressed_block_ptr(list, block_idx) }; + let lo = state.doc_ids.partition_point(|&doc| doc < target32); + let hi = if win_end32 == u32::MAX { + state.doc_ids.len() + } else { + lo + state.doc_ids[lo..].partition_point(|&doc| doc <= win_end32) + }; + if lo == hi { + // No docs of this clause in the window: the whole window + // has no conjunction match. If this was the clause's last + // block and it is fully behind the target, the clause is + // exhausted and the conjunction is done. + if block_idx + 1 >= list.blocks.len() + && state.doc_ids.last().is_none_or(|&doc| doc < target32) + { + exhausted = true; + } + skip_window = true; + break; + } + wins.push(WindowList { + docs: state.doc_ids.as_ptr(), + freqs: state.freqs.as_ptr(), + pos: lo, + end: hi, + block_start: posting_block_start(block_idx, list.block_size), + }); + } + if exhausted { + break 'window; + } + if !skip_window { + // Constant within the window (block-anchored); mirrors + // `and_candidate_cannot_beat_threshold`'s remaining-clause + // bound of first-clause-exact + rest-block-max. + let others_block_max: f32 = self.lead[1..] + .iter() + .map(|posting| posting.block_max_score(&self.scorer)) + .sum(); + + while wins[0].pos < wins[0].end { + let doc = unsafe { *wins[0].docs.add(wins[0].pos) }; + let mut advance_to: Option = None; + let mut window_done = false; + for win in wins.iter_mut().skip(1) { + while win.pos < win.end && unsafe { *win.docs.add(win.pos) } < doc { + win.pos += 1; + } + if win.pos >= win.end { + window_done = true; + break; + } + let clause_doc = unsafe { *win.docs.add(win.pos) }; + if clause_doc > doc { + advance_to = Some(clause_doc); + break; + } + } + if window_done { + break; + } + if let Some(advance_to) = advance_to { + let win = &mut wins[0]; + while win.pos < win.end && unsafe { *win.docs.add(win.pos) } < advance_to { + win.pos += 1; + } + continue; + } + + // All clauses sit on `doc`: a conjunction candidate. + let doc_length = self.docs.scoring_num_tokens(doc); + if self.threshold > 0.0 && num_lists >= 2 { + let first_score = self.lead[0].score( + &self.scorer, + unsafe { *wins[0].freqs.add(wins[0].pos) }, + doc_length, + ); + if first_score + others_block_max <= self.threshold { + self.and_candidates_pruned_before_return += 1; + wins[0].pos += 1; + continue; + } + } + stats.candidates_seen += 1; + self.and_window_stats.candidates_returned += 1; + num_comparisons += 1; + + let row_id = if docs_has_row_ids { + self.docs.row_id(doc) + } else { + u64::from(doc) + }; + if docs_has_row_ids + && (row_id == RowAddress::TOMBSTONE_ROW || !mask.selected(row_id)) + { + wins[0].pos += 1; + continue; + } + + if let Some(slop) = phrase_slop { + // Park every clause's iterator on this doc so + // `position_cursor` reads the right posting entry. The + // window block is already decompressed; position blocks + // decode lazily and are cached per block. + for (win, posting) in wins.iter().zip(self.lead.iter_mut()) { + posting.index = win.block_start + win.pos; + posting.current_doc = Some(DocInfo::Raw(RawDocInfo { + doc_id: doc, + frequency: unsafe { *win.freqs.add(win.pos) }, + })); + } + let matched = if slop == 0 { + self.check_exact_positions_bulk() + } else { + self.check_positions(slop as i32) + }; + if !matched { + wins[0].pos += 1; + continue; + } + } + stats.full_scores += 1; + + let mut score = 0.0f32; + for (win, posting) in wins.iter().zip(self.lead.iter()) { + let freq = unsafe { *win.freqs.add(win.pos) }; + score += posting.score(&self.scorer, freq, doc_length); + } + + let insert = if candidates.len() < limit { + true + } else { + score > candidates.peek().unwrap().0.0.score.0 + }; + if insert { + stats.freqs_collected += 1; + let freqs = wins + .iter() + .zip(self.lead.iter()) + .map(|(win, posting)| { + (posting.term_index(), unsafe { *win.freqs.add(win.pos) }) + }) + .collect(); + if candidates.len() >= limit { + candidates.pop(); + } + candidates.push(Reverse(( + ScoredDoc::new(row_id, score), + freqs, + doc_length, + u64::from(doc), + ))); + if candidates.len() == limit { + let kth = candidates.peek().unwrap().0.0.score.0; + self.update_threshold(kth, params.wand_factor); + } + } + wins[0].pos += 1; + } + } + + if win_end == TERMINATED_DOC_ID { + break; + } + target = win_end + 1; + } + + tracing::debug!( + and_windows_wide = self.and_window_stats.windows_wide, + and_windows_narrow = self.and_window_stats.windows_narrow, + and_windows_skipped = self.and_window_stats.windows_skipped, + and_range_blocks_scanned = self.and_window_stats.range_blocks_scanned, + and_candidates_returned = self.and_window_stats.candidates_returned, + "fts conjunction block-max window stats (bulk)" + ); + metrics.record_comparisons(num_comparisons); + let pruned_before_return = self + .and_candidates_pruned_before_return + .saturating_sub(stats.pruned_before_return_start); + metrics.record_and_candidates_seen(stats.candidates_seen); + metrics.record_and_candidates_pruned_before_return(pruned_before_return); + metrics.record_and_full_scores(stats.full_scores); + metrics.record_freqs_collected(stats.freqs_collected); + + let to_addr = |row_id_slot: u64| { + if docs_has_row_ids { + CandidateAddr::RowId(row_id_slot) + } else { + CandidateAddr::Pending(row_id_slot as u32) + } + }; + Ok(candidates + .into_iter() + .map( + |Reverse((doc, freqs, doc_length, posting_doc_id))| DocCandidate { + addr: to_addr(doc.row_id), + posting_doc_id, + freqs, + doc_length, + }, + ) + .collect()) + } + fn and_move_to_next_block(&mut self, target: u64) { if self.threshold <= 0.0 { self.up_to = Some(target); @@ -2812,6 +3208,57 @@ impl<'a, S: Scorer> Wand<'a, S> { } } + /// Allocation-free exact-phrase check for the bulk conjunction path, + /// where every clause is a parked `lead` iterator. Semantically identical + /// to [`Self::check_exact_positions`] — some base position must align all + /// clauses at their query offsets — without the per-candidate cursor vec + /// and sort. + fn check_exact_positions_bulk(&self) -> bool { + const MAX_INLINE_CLAUSES: usize = 16; + let num_clauses = self.lead.len(); + if num_clauses > MAX_INLINE_CLAUSES { + return self.check_exact_positions(); + } + // Cursors stay alive in the stack array so owned position buffers + // (legacy per-doc storage) remain valid while we scan. + let mut cursors: [Option>; MAX_INLINE_CLAUSES] = + std::array::from_fn(|_| None); + let mut anchor_idx = 0usize; + let mut anchor_len = usize::MAX; + for (index, (slot, posting)) in cursors.iter_mut().zip(self.lead.iter()).enumerate() { + let cursor = posting.position_cursor().expect("positions must exist"); + if cursor.len() < anchor_len { + anchor_len = cursor.len(); + anchor_idx = index; + } + *slot = Some(cursor); + } + + let anchor = cursors[anchor_idx] + .as_ref() + .expect("anchor cursor was just populated"); + let anchor_offset = anchor.position_in_query as u32; + 'anchor: for &anchor_position in anchor.positions.as_slice() { + let Some(base) = anchor_position.checked_sub(anchor_offset) else { + continue; + }; + for (index, slot) in cursors[..num_clauses].iter().enumerate() { + if index == anchor_idx { + continue; + } + let cursor = slot.as_ref().expect("clause cursor was just populated"); + let Some(target) = base.checked_add(cursor.position_in_query as u32) else { + return false; + }; + if cursor.positions.as_slice().binary_search(&target).is_err() { + continue 'anchor; + } + } + return true; + } + false + } + fn check_exact_positions(&self) -> bool { let mut position_iters = self .current_doc_postings() @@ -4703,4 +5150,114 @@ mod tests { assert_eq!(second.0.doc_id(), 1); assert!(wand.check_positions(0)); } + + /// The bulk conjunction path must return exactly the classic loop's + /// results — same docs, freqs, and doc lengths — for both plain AND and + /// phrase queries, across multi-block lists with heap/threshold pruning + /// in play. + #[rstest] + #[case::and_k10(false, 10)] + #[case::and_k3(false, 3)] + #[case::phrase_k10(true, 10)] + #[case::phrase_k3(true, 3)] + fn test_bulk_and_matches_classic(#[case] phrase: bool, #[case] limit: usize) { + let num_docs = (BLOCK_SIZE * 8 + 37) as u32; + let mut docs = DocSet::default(); + for doc_id in 0..num_docs { + docs.append(u64::from(doc_id), 32 + doc_id % 57); + } + + // Three clauses with different densities; membership comes from a + // cheap deterministic mix so docs scatter across blocks. + let clause_docs = |modulus: u32, salt: u32| -> Vec { + (0..num_docs) + .filter(|doc| (doc.wrapping_mul(2654435761).wrapping_add(salt)) % modulus < 2) + .collect() + }; + let clauses = [clause_docs(3, 7), clause_docs(4, 13), clause_docs(5, 29)]; + + let build_postings = || { + clauses + .iter() + .enumerate() + .map(|(term_pos, doc_ids)| { + let list = if phrase { + // Roughly half of each clause's docs put the token at + // position base+term_pos (forming the phrase); the + // rest scatter so the position check has misses. + let positions = doc_ids + .iter() + .map(|&doc| { + if doc % 2 == 0 { + vec![5 + term_pos as u32, 40 + (doc % 3)] + } else { + vec![20 + (term_pos as u32) * 4] + } + }) + .collect::>(); + generate_posting_list_with_positions(doc_ids.clone(), positions, 1.0, true) + } else { + generate_posting_list(doc_ids.clone(), 1.0, None, true) + }; + PostingIterator::with_query_weight( + format!("t{term_pos}"), + term_pos as u32, + term_pos as u32, + 1.0 + term_pos as f32 * 0.5, + list, + docs.len(), + ) + }) + .collect::>() + }; + + let mut params = FtsSearchParams::default().with_limit(Some(limit)); + if phrase { + params.phrase_slop = Some(0); + } + + let normalize = |result: Vec| { + let mut rows = result + .into_iter() + .map(|candidate| { + ( + candidate.posting_doc_id, + candidate.doc_length, + candidate.freqs, + match candidate.addr { + CandidateAddr::RowId(row_id) => row_id, + CandidateAddr::Pending(doc_id) => u64::from(doc_id), + }, + ) + }) + .collect::>(); + rows.sort_unstable(); + rows + }; + + let run = |disable_bulk: bool| { + let mut wand = Wand::new( + Operator::And, + build_postings().into_iter(), + &docs, + UnitScorer, + ); + if disable_bulk { + wand = wand.disable_bulk_and(); + } + normalize( + wand.search( + ¶ms, + Arc::new(RowAddrMask::default()), + &NoOpMetricsCollector, + ) + .unwrap(), + ) + }; + + let bulk = run(false); + let classic = run(true); + assert!(!bulk.is_empty(), "test corpus should produce matches"); + assert_eq!(bulk, classic); + } } From 3c60d9afc56a230b6135de3866598fa1b5b84705 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Sat, 4 Jul 2026 01:26:10 +0800 Subject: [PATCH 6/7] perf(index): specialized SIMD merge kernels and freq-bound pruning for bulk fts conjunctions Three changes to the bulk AND/phrase candidate pipeline, all score-identical to the classic loop: - Clause-count-specialized merge kernels (2 and 3 clauses) keep cursors and bounds in registers, with an AVX2 find_next_geq catch-up scan (the analogue of Lucene's VectorUtil.findNextGEQ) replacing the mispredict- heavy scalar walk. Generic kernel covers other widths. - Two-pass batched scoring: kernels only record (doc, offsets) matches; doc lengths are then gathered back-to-back so their cache misses overlap, and the prune/score/insert pass runs in doc order with the classic semantics. - A per-clause frequency-bucketed score-bound LUT lets kernels skip lead docs before any follower advance when even the doc-length-free bound cannot beat the threshold (score-first ordering, Lucene-style); the bound is monotone so the skipped set is a subset of the exact prune. Warm mmlb AND @200M, 8 concurrent: k10 0.060->0.049s (161qps), k100 0.118->0.098s (81qps); phrase\@50M 3w k10 0.210->0.204s, 2w 0.042->0.040s. Co-Authored-By: Claude Fable 5 --- rust/lance-index/src/scalar/inverted/wand.rs | 430 ++++++++++++++++--- 1 file changed, 382 insertions(+), 48 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 01525fc2832..3ef39b776bc 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -23,7 +23,7 @@ use crate::metrics::MetricsCollector; use super::{ CompressedPositionStorage, impact::{IMPACT_LEVEL1_BLOCKS, ImpactSkipData}, - index::PositionStreamCodec, + index::{PositionStreamCodec, dequantize_doc_length}, query::Operator, scorer::{K1, idf}, }; @@ -63,6 +63,55 @@ static USE_MAXSCORE_SEARCH: LazyLock = static USE_BULK_AND_SEARCH: LazyLock = LazyLock::new(|| std::env::var("LANCE_FTS_BULK_AND").as_deref() != Ok("0")); +#[cfg(target_arch = "x86_64")] +static HAS_AVX2: LazyLock = LazyLock::new(|| std::arch::is_x86_feature_detected!("avx2")); + +/// First index in `[pos, end)` where `docs[index] >= target` (scalar). +/// Posting-block doc ids stay below 2^31, which the AVX2 variant relies on. +#[inline] +unsafe fn find_next_geq_scalar(docs: *const u32, mut pos: usize, end: usize, target: u32) -> usize { + unsafe { + while pos < end && *docs.add(pos) < target { + pos += 1; + } + } + pos +} + +/// AVX2 `find_next_geq` (the analogue of Lucene's VectorUtil.findNextGEQ): +/// branchless 8-wide compare+movemask kills the mispredicted exits that +/// dominate the scalar catch-up scan on irregular doc gaps. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +unsafe fn find_next_geq_avx2(docs: *const u32, mut pos: usize, end: usize, target: u32) -> usize { + use core::arch::x86_64::*; + debug_assert!(target <= i32::MAX as u32); + unsafe { + let target_lanes = _mm256_set1_epi32(target as i32); + while pos + 8 <= end { + let docs_lanes = _mm256_loadu_si256(docs.add(pos) as *const __m256i); + // lane mask of docs[i] < target; doc ids < 2^31 keep the signed + // compare equivalent to unsigned. + let below = _mm256_cmpgt_epi32(target_lanes, docs_lanes); + let mask = _mm256_movemask_ps(_mm256_castsi256_ps(below)) as u32; + if mask != 0xFF { + return pos + mask.trailing_ones() as usize; + } + pos += 8; + } + find_next_geq_scalar(docs, pos, end, target) + } +} + +#[inline] +unsafe fn find_next_geq(docs: *const u32, pos: usize, end: usize, target: u32) -> usize { + #[cfg(target_arch = "x86_64")] + if *HAS_AVX2 { + return unsafe { find_next_geq_avx2(docs, pos, end, target) }; + } + unsafe { find_next_geq_scalar(docs, pos, end, target) } +} + #[inline] fn posting_block_idx(index: usize, block_size: usize) -> usize { match block_size { @@ -2379,6 +2428,180 @@ impl<'a, S: Scorer> Wand<'a, S> { block_start: usize, } + // Merge kernels: intersect the window's per-clause slices and record + // each match as (doc, per-clause block offsets). Hand-specialized for + // two and three clauses so cursors and bounds stay in registers; the + // generic kernel covers other widths. Offsets fit u8 because a block + // holds at most `MAX_POSTING_BLOCK_SIZE` (256) entries. The macro + // stamps a scalar and an AVX2 variant — `#[target_feature]` must + // cover the whole kernel for the vector catch-up scans to inline. + // Kernels prune a lead doc before any follower advance when its + // frequency-bucketed score bound (plus the other clauses' block + // maxes) cannot beat the threshold — the score-first ordering + // Lucene's conjunction scorer uses, with doc length dropped from the + // bound so no doc-length load is involved. The bound is monotone + // (doc length only shrinks a BM25 weight, and the last bucket holds + // the clause sup), so every skipped doc would also fail the exact + // per-candidate prune: results are unchanged, the work never happens. + macro_rules! merge_kernels { + ($name2:ident, $name3:ident, $geq:ident $(, #[$feat:meta])?) => { + $(#[$feat])? + unsafe fn $name2( + wins: &[WindowList], + lut: &[f32; FREQ_LUT_BUCKETS], + others_block_max: f32, + threshold: f32, + docs_out: &mut Vec, + offs_out: &mut Vec, + ) { + let (d0, mut p0, e0) = (wins[0].docs, wins[0].pos, wins[0].end); + let (d1, mut p1, e1) = (wins[1].docs, wins[1].pos, wins[1].end); + let f0 = wins[0].freqs; + let prune = threshold > f32::NEG_INFINITY; + unsafe { + while p0 < e0 { + let doc = *d0.add(p0); + if prune { + let freq = (*f0.add(p0) as usize).min(FREQ_LUT_BUCKETS - 1); + if lut[freq] + others_block_max <= threshold { + p0 += 1; + continue; + } + } + p1 = $geq(d1, p1, e1, doc); + if p1 >= e1 { + return; + } + let second = *d1.add(p1); + if second > doc { + p0 = $geq(d0, p0 + 1, e0, second); + continue; + } + docs_out.push(doc); + offs_out.push(p0 as u8); + offs_out.push(p1 as u8); + p0 += 1; + } + } + } + + $(#[$feat])? + unsafe fn $name3( + wins: &[WindowList], + lut: &[f32; FREQ_LUT_BUCKETS], + others_block_max: f32, + threshold: f32, + docs_out: &mut Vec, + offs_out: &mut Vec, + ) { + let (d0, mut p0, e0) = (wins[0].docs, wins[0].pos, wins[0].end); + let (d1, mut p1, e1) = (wins[1].docs, wins[1].pos, wins[1].end); + let (d2, mut p2, e2) = (wins[2].docs, wins[2].pos, wins[2].end); + let f0 = wins[0].freqs; + let prune = threshold > f32::NEG_INFINITY; + unsafe { + 'outer: while p0 < e0 { + let doc = *d0.add(p0); + if prune { + let freq = (*f0.add(p0) as usize).min(FREQ_LUT_BUCKETS - 1); + if lut[freq] + others_block_max <= threshold { + p0 += 1; + continue 'outer; + } + } + p1 = $geq(d1, p1, e1, doc); + if p1 >= e1 { + return; + } + let second = *d1.add(p1); + if second > doc { + p0 = $geq(d0, p0 + 1, e0, second); + continue 'outer; + } + p2 = $geq(d2, p2, e2, doc); + if p2 >= e2 { + return; + } + let third = *d2.add(p2); + if third > doc { + p0 = $geq(d0, p0 + 1, e0, third); + continue 'outer; + } + docs_out.push(doc); + offs_out.push(p0 as u8); + offs_out.push(p1 as u8); + offs_out.push(p2 as u8); + p0 += 1; + } + } + } + }; + } + merge_kernels!(merge_window_2, merge_window_3, find_next_geq_scalar); + #[cfg(target_arch = "x86_64")] + merge_kernels!( + merge_window_2_avx2, + merge_window_3_avx2, + find_next_geq_avx2, + #[target_feature(enable = "avx2")] + ); + + #[inline] + fn merge_window_1(wins: &[WindowList], docs_out: &mut Vec, offs_out: &mut Vec) { + let win = &wins[0]; + for pos in win.pos..win.end { + docs_out.push(unsafe { *win.docs.add(pos) }); + offs_out.push(pos as u8); + } + } + + #[inline] + #[allow(clippy::too_many_arguments)] + fn merge_window_n( + wins: &[WindowList], + lut: &[f32; FREQ_LUT_BUCKETS], + others_block_max: f32, + threshold: f32, + cursors: &mut Vec, + docs_out: &mut Vec, + offs_out: &mut Vec, + ) { + let prune = threshold > f32::NEG_INFINITY; + cursors.clear(); + cursors.extend(wins.iter().map(|win| win.pos)); + 'outer: while cursors[0] < wins[0].end { + let doc = unsafe { *wins[0].docs.add(cursors[0]) }; + if prune { + let freq = unsafe { *wins[0].freqs.add(cursors[0]) as usize } + .min(FREQ_LUT_BUCKETS - 1); + if lut[freq] + others_block_max <= threshold { + cursors[0] += 1; + continue 'outer; + } + } + for j in 1..wins.len() { + let win = &wins[j]; + let pos = unsafe { find_next_geq(win.docs, cursors[j], win.end, doc) }; + cursors[j] = pos; + if pos >= win.end { + return; + } + let clause_doc = unsafe { *win.docs.add(pos) }; + if clause_doc > doc { + cursors[0] = unsafe { + find_next_geq(wins[0].docs, cursors[0] + 1, wins[0].end, clause_doc) + }; + continue 'outer; + } + } + docs_out.push(doc); + for &pos in cursors.iter() { + offs_out.push(pos as u8); + } + cursors[0] += 1; + } + } + let mut candidates: TopKHeap = BinaryHeap::with_capacity(std::cmp::min(limit, BLOCK_SIZE * 10)); let mut num_comparisons: usize = 0; @@ -2387,6 +2610,37 @@ impl<'a, S: Scorer> Wand<'a, S> { ..Default::default() }; let mut wins: Vec = Vec::with_capacity(num_lists); + // Per-window candidate batch. The merge kernel only records matches; + // scoring then runs in two passes so the doc-length gather issues + // independent loads (their cache misses overlap) instead of + // serializing behind per-candidate branching. A window spans at most + // one block per clause, so a batch holds at most one block's worth. + 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 cursor_scratch: Vec = Vec::with_capacity(num_lists); + + // Per-window prune LUT for the merge kernels: an upper bound of the + // first (rarest) clause's score by clamped frequency. Lead docs whose + // bound plus the remaining clauses' block maxes cannot beat the + // threshold are skipped before any follower advances — the same + // score-first ordering Lucene's conjunction scorer uses, with the + // doc-length dropped from the bound so no doc-length load is needed. + // The last bucket holds the frequency-independent sup, so clamping + // stays a valid upper bound. Skips only avoid work; every emitted + // candidate still goes through the exact per-candidate prune below. + const FREQ_LUT_BUCKETS: usize = 64; + let mut freq_bound_lut = [f32::INFINITY; FREQ_LUT_BUCKETS]; + for (freq, slot) in freq_bound_lut + .iter_mut() + .enumerate() + .take(FREQ_LUT_BUCKETS - 1) + { + *slot = self.lead[0].score(&self.scorer, freq as u32, 0); + } + // The clamp bucket must bound every frequency it absorbs; the + // clause-wide sup does. + freq_bound_lut[FREQ_LUT_BUCKETS - 1] = self.lead[0].approximate_upper_bound(); // The conjunction can only start at the max of the clauses' first docs. let mut target: u64 = 0; @@ -2480,46 +2734,105 @@ impl<'a, S: Scorer> Wand<'a, S> { .map(|posting| posting.block_max_score(&self.scorer)) .sum(); - while wins[0].pos < wins[0].end { - let doc = unsafe { *wins[0].docs.add(wins[0].pos) }; - let mut advance_to: Option = None; - let mut window_done = false; - for win in wins.iter_mut().skip(1) { - while win.pos < win.end && unsafe { *win.docs.add(win.pos) } < doc { - win.pos += 1; - } - if win.pos >= win.end { - window_done = true; - break; - } - let clause_doc = unsafe { *win.docs.add(win.pos) }; - if clause_doc > doc { - advance_to = Some(clause_doc); - break; + batch_docs.clear(); + batch_offs.clear(); + // NEG_INFINITY disables the kernel-level freq-bound prune + // (single clause, or no threshold yet). + let kernel_threshold = if self.threshold > 0.0 && num_lists >= 2 { + self.threshold + } else { + f32::NEG_INFINITY + }; + #[cfg(target_arch = "x86_64")] + let use_avx2 = *HAS_AVX2; + #[cfg(not(target_arch = "x86_64"))] + let use_avx2 = false; + match (num_lists, use_avx2) { + (1, _) => merge_window_1(&wins, &mut batch_docs, &mut batch_offs), + #[cfg(target_arch = "x86_64")] + (2, true) => unsafe { + merge_window_2_avx2( + &wins, + &freq_bound_lut, + others_block_max, + kernel_threshold, + &mut batch_docs, + &mut batch_offs, + ) + }, + #[cfg(target_arch = "x86_64")] + (3, true) => unsafe { + merge_window_3_avx2( + &wins, + &freq_bound_lut, + others_block_max, + kernel_threshold, + &mut batch_docs, + &mut batch_offs, + ) + }, + (2, _) => unsafe { + merge_window_2( + &wins, + &freq_bound_lut, + others_block_max, + kernel_threshold, + &mut batch_docs, + &mut batch_offs, + ) + }, + (3, _) => unsafe { + merge_window_3( + &wins, + &freq_bound_lut, + others_block_max, + kernel_threshold, + &mut batch_docs, + &mut batch_offs, + ) + }, + _ => merge_window_n( + &wins, + &freq_bound_lut, + others_block_max, + kernel_threshold, + &mut cursor_scratch, + &mut batch_docs, + &mut batch_offs, + ), + } + + // 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. + batch_lens.clear(); + match self.docs.scoring_norms() { + Some(norms) => { + for &doc in batch_docs.iter() { + batch_lens.push(dequantize_doc_length(norms[doc as usize])); } } - if window_done { - break; - } - if let Some(advance_to) = advance_to { - let win = &mut wins[0]; - while win.pos < win.end && unsafe { *win.docs.add(win.pos) } < advance_to { - win.pos += 1; + None => { + for &doc in batch_docs.iter() { + batch_lens.push(self.docs.scoring_num_tokens(doc)); } - continue; } + } - // All clauses sit on `doc`: a conjunction candidate. - let doc_length = 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 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(wins[0].pos) }, + unsafe { *wins[0].freqs.add(offs[0] as usize) }, doc_length, ); if first_score + others_block_max <= self.threshold { self.and_candidates_pruned_before_return += 1; - wins[0].pos += 1; continue; } } @@ -2535,7 +2848,6 @@ impl<'a, S: Scorer> Wand<'a, S> { if docs_has_row_ids && (row_id == RowAddress::TOMBSTONE_ROW || !mask.selected(row_id)) { - wins[0].pos += 1; continue; } @@ -2544,11 +2856,13 @@ impl<'a, S: Scorer> Wand<'a, S> { // `position_cursor` reads the right posting entry. The // window block is already decompressed; position blocks // decode lazily and are cached per block. - for (win, posting) in wins.iter().zip(self.lead.iter_mut()) { - posting.index = win.block_start + win.pos; + for ((win, posting), &off) in + wins.iter().zip(self.lead.iter_mut()).zip(offs.iter()) + { + posting.index = win.block_start + off as usize; posting.current_doc = Some(DocInfo::Raw(RawDocInfo { doc_id: doc, - frequency: unsafe { *win.freqs.add(win.pos) }, + frequency: unsafe { *win.freqs.add(off as usize) }, })); } let matched = if slop == 0 { @@ -2557,15 +2871,15 @@ impl<'a, S: Scorer> Wand<'a, S> { self.check_positions(slop as i32) }; if !matched { - wins[0].pos += 1; continue; } } stats.full_scores += 1; let mut score = 0.0f32; - for (win, posting) in wins.iter().zip(self.lead.iter()) { - let freq = unsafe { *win.freqs.add(win.pos) }; + 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); } @@ -2579,8 +2893,11 @@ impl<'a, S: Scorer> Wand<'a, S> { let freqs = wins .iter() .zip(self.lead.iter()) - .map(|(win, posting)| { - (posting.term_index(), unsafe { *win.freqs.add(win.pos) }) + .zip(offs.iter()) + .map(|((win, posting), &off)| { + (posting.term_index(), unsafe { + *win.freqs.add(off as usize) + }) }) .collect(); if candidates.len() >= limit { @@ -2597,7 +2914,6 @@ impl<'a, S: Scorer> Wand<'a, S> { self.update_threshold(kth, params.wand_factor); } } - wins[0].pos += 1; } } @@ -4725,8 +5041,11 @@ mod tests { let addrs = result.into_iter().map(|doc| doc.addr).collect::>(); assert!(matches!(addrs.as_slice(), [CandidateAddr::RowId(0)])); let scored = scored.load(Ordering::Relaxed); + // The bulk path evaluates 63 doc weights up front to fill its + // frequency-bound prune LUT; those are bound computations, not + // per-candidate scoring. assert!( - scored <= BLOCK_SIZE + 1, + scored <= BLOCK_SIZE + 1 + 63, "expected candidate pruning to avoid full scoring in the first block, scored {scored}" ); } @@ -5156,25 +5475,40 @@ mod tests { /// phrase queries, across multi-block lists with heap/threshold pruning /// in play. #[rstest] - #[case::and_k10(false, 10)] - #[case::and_k3(false, 3)] - #[case::phrase_k10(true, 10)] - #[case::phrase_k3(true, 3)] - fn test_bulk_and_matches_classic(#[case] phrase: bool, #[case] limit: usize) { + #[case::and_k10(false, 10, 3)] + #[case::and_k3(false, 3, 3)] + #[case::and_two_clauses(false, 10, 2)] + #[case::and_four_clauses(false, 10, 4)] + #[case::phrase_k10(true, 10, 3)] + #[case::phrase_k3(true, 3, 3)] + #[case::phrase_two_clauses(true, 10, 2)] + #[case::phrase_four_clauses(true, 10, 4)] + fn test_bulk_and_matches_classic( + #[case] phrase: bool, + #[case] limit: usize, + #[case] num_clauses: usize, + ) { let num_docs = (BLOCK_SIZE * 8 + 37) as u32; let mut docs = DocSet::default(); for doc_id in 0..num_docs { docs.append(u64::from(doc_id), 32 + doc_id % 57); } - // Three clauses with different densities; membership comes from a - // cheap deterministic mix so docs scatter across blocks. + // Clauses with different densities; membership comes from a cheap + // deterministic mix so docs scatter across blocks. Clause count picks + // the specialized (1-3) or generic (4+) merge kernel. let clause_docs = |modulus: u32, salt: u32| -> Vec { (0..num_docs) .filter(|doc| (doc.wrapping_mul(2654435761).wrapping_add(salt)) % modulus < 2) .collect() }; - let clauses = [clause_docs(3, 7), clause_docs(4, 13), clause_docs(5, 29)]; + let clauses = [ + clause_docs(3, 7), + clause_docs(4, 13), + clause_docs(5, 29), + clause_docs(3, 41), + ][..num_clauses] + .to_vec(); let build_postings = || { clauses From a1aea58ba40c865422a937e4e214b2db150d4811 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Sat, 4 Jul 2026 18:57:44 +0800 Subject: [PATCH 7/7] 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 b6b48ce4149..482be1e7ae5 100644 --- a/rust/lance-index/src/scalar/inverted/scorer.rs +++ b/rust/lance-index/src/scalar/inverted/scorer.rs @@ -11,6 +11,16 @@ use std::sync::Arc; pub trait Scorer: Send + Sync { fn query_weight(&self, token: &str) -> f32; fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32; + + /// 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 { @@ -21,6 +31,18 @@ impl Scorer for Arc { fn doc_weight(&self, freq: u32, doc_tokens: u32) -> f32 { self.as_ref().doc_weight(freq, doc_tokens) } + + 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 @@ -88,10 +110,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())) } } @@ -152,10 +178,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)) } } diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 3ef39b776bc..8be17920cc8 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, 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") @@ -952,6 +984,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, @@ -959,6 +994,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) { @@ -986,8 +1022,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); } @@ -1346,6 +1390,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) { @@ -1431,6 +1495,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, @@ -1504,31 +1569,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); @@ -1572,10 +1637,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, }, ) @@ -1619,7 +1684,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); @@ -1667,25 +1733,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); @@ -1700,10 +1766,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, }, ) @@ -1753,6 +1819,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. @@ -1896,8 +1967,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 + total_non_essential_bound <= self.threshold) { @@ -1925,8 +2009,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), + ), + }; } } @@ -1939,26 +2035,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; @@ -2057,6 +2156,7 @@ impl<'a, S: Scorer> Wand<'a, S> { clause_idx, self.docs, &self.scorer, + norm_k_ref, &mut acc, ); } @@ -2096,7 +2196,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 @@ -2112,7 +2217,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) + } + }; } } @@ -2120,10 +2237,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)) @@ -2134,16 +2251,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; @@ -2178,10 +2295,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, }, ) @@ -2618,7 +2735,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 @@ -2802,35 +2927,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; @@ -2880,7 +3022,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 { @@ -2890,24 +3037,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; @@ -2950,10 +3095,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, }, )