From ff6209ba2d76e6ea33cbfca8bcaa7ed17230a64c Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 3 Jul 2026 21:18:00 +0800 Subject: [PATCH 1/4] 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 --- .../src/scalar/inverted/encoding.rs | 166 +++++ rust/lance-index/src/scalar/inverted/wand.rs | 607 +++++++++++++++++- 2 files changed, 748 insertions(+), 25 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/encoding.rs b/rust/lance-index/src/scalar/inverted/encoding.rs index 8fb7f8f4032..a6f2107e27f 100644 --- a/rust/lance-index/src/scalar/inverted/encoding.rs +++ b/rust/lance-index/src/scalar/inverted/encoding.rs @@ -746,6 +746,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], @@ -1185,6 +1277,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 550f07cde8f..7c7d125dff8 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, ImpactScoreCache, 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 conservative_bm25_upper_bound(query_weight: f32) -> f32 { @@ -105,6 +112,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 @@ -123,6 +141,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, @@ -630,32 +654,79 @@ impl PostingIterator { let block_offset = self.index & list.block_mask(); 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, - )) } }, } @@ -1194,6 +1265,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 @@ -1258,12 +1332,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); @@ -1335,6 +1418,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 @@ -2323,6 +2419,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(list.block_least_doc_id(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: block_idx << list.block_shift(), + }); + } + 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); @@ -2905,6 +3301,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() @@ -4972,4 +5419,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 63113532ce8972d877115b2ee8bc5ff7e6746e8c Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Wed, 15 Jul 2026 18:06:14 +0800 Subject: [PATCH 2/4] fix(fts): gate bulk conjunctions and recycle position buffers --- rust/lance-index/src/scalar/inverted/wand.rs | 282 ++++++++++++++++--- 1 file changed, 240 insertions(+), 42 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 7c7d125dff8..acd4fc29f76 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -5,7 +5,7 @@ use std::ops::Deref; use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::{Arc, LazyLock}; use std::{ - cell::UnsafeCell, + cell::{RefCell, UnsafeCell}, collections::{BinaryHeap, VecDeque}, }; use std::{cmp::Reverse, fmt::Debug}; @@ -59,9 +59,60 @@ static USE_MAXSCORE_SEARCH: LazyLock = // 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")); +// LANCE_FTS_BULK_AND accepts auto (default), on/1, or off/0. Auto enables the +// bulk path only for its consistently faster two- and three-clause kernels. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +enum BulkAndMode { + #[default] + Auto, + On, + Off, +} + +impl BulkAndMode { + fn parse(value: &str) -> Option { + let value = value.trim(); + if value.eq_ignore_ascii_case("auto") { + Some(Self::Auto) + } else if value.eq_ignore_ascii_case("on") || value == "1" { + Some(Self::On) + } else if value.eq_ignore_ascii_case("off") || value == "0" { + Some(Self::Off) + } else { + None + } + } + + const fn enabled_for(self, num_clauses: usize) -> bool { + match self { + Self::Auto => matches!(num_clauses, 2 | 3), + Self::On => true, + Self::Off => false, + } + } +} + +fn bulk_and_mode_from_env() -> BulkAndMode { + match std::env::var("LANCE_FTS_BULK_AND") { + Ok(value) => BulkAndMode::parse(&value).unwrap_or_else(|| { + log::warn!( + "Invalid LANCE_FTS_BULK_AND value {value:?}; expected auto, on/1, or off/0; \ + falling back to auto" + ); + BulkAndMode::Auto + }), + Err(std::env::VarError::NotPresent) => BulkAndMode::Auto, + Err(std::env::VarError::NotUnicode(value)) => { + log::warn!( + "Invalid non-Unicode LANCE_FTS_BULK_AND value {value:?}; expected auto, on/1, \ + or off/0; falling back to auto" + ); + BulkAndMode::Auto + } + } +} + +static BULK_AND_MODE: LazyLock = LazyLock::new(bulk_and_mode_from_env); #[inline] fn conservative_bm25_upper_bound(query_weight: f32) -> f32 { @@ -98,6 +149,10 @@ pub struct PostingIterator { block_idx: usize, current_doc: Option, approximate_upper_bound: f32, + // Position cursors temporarily own this buffer and return it on drop. This + // keeps repeated cursor creation allocation-free without lending a slice + // out of the interior-mutable compressed state. + position_scratch: RefCell>>, // for compressed posting list compressed: Option>, @@ -114,15 +169,14 @@ struct CompressedState { 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. + // tail, and the block's total delta count. 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 @@ -146,7 +200,6 @@ impl CompressedState { 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, @@ -533,6 +586,7 @@ impl PostingIterator { block_idx: 0, current_doc: None, approximate_upper_bound, + position_scratch: RefCell::new(Some(Vec::new())), compressed, }; posting.refresh_current_doc(); @@ -681,6 +735,11 @@ impl PostingIterator { } let delta_start = compressed.position_offsets[block_offset]; let delta_end = compressed.position_offsets[block_offset + 1]; + let mut position_values = self + .position_scratch + .borrow_mut() + .take() + .unwrap_or_default(); seek_packed_doc_positions( stream.block(block_idx), compressed.position_total_deltas, @@ -689,11 +748,14 @@ impl PostingIterator { &mut compressed.position_unpacked_group, &mut compressed.position_unpacked_group_idx, &mut compressed.position_tail, - &mut compressed.position_doc_scratch, + &mut position_values, ) .expect("shared position stream doc decoding should succeed"); Some(PositionCursor::new( - PositionValues::Borrowed(&compressed.position_doc_scratch[..]), + PositionValues::Recycled(RecycledPositionValues::new( + position_values, + &self.position_scratch, + )), self.position as i32, )) } @@ -721,8 +783,19 @@ impl PostingIterator { } let start = compressed.position_offsets[block_offset]; let end = compressed.position_offsets[block_offset + 1]; + let mut position_values = self + .position_scratch + .borrow_mut() + .take() + .unwrap_or_default(); + position_values.clear(); + position_values + .extend_from_slice(&compressed.position_values[start..end]); Some(PositionCursor::new( - PositionValues::Borrowed(&compressed.position_values[start..end]), + PositionValues::Recycled(RecycledPositionValues::new( + position_values, + &self.position_scratch, + )), self.position as i32, )) } @@ -1265,9 +1338,11 @@ 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, + // Test-only override for comparing bulk and classic conjunctions without + // mutating the process-wide environment. + bulk_and_mode_override: Option, + #[cfg(test)] + bulk_and_searches: usize, docs: &'a DocSet, scorer: S, // Shared cross-partition top-k floor. Each partition publishes its local @@ -1332,18 +1407,20 @@ 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, + bulk_and_mode_override: None, + #[cfg(test)] + bulk_and_searches: 0, 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. + /// Test hook: force one conjunction mode so parity tests can compare bulk + /// and classic search within one process. #[cfg(test)] - pub(crate) fn disable_bulk_and(mut self) -> Self { - self.bulk_and_disabled = true; + fn with_bulk_and_mode(mut self, mode: BulkAndMode) -> Self { + self.bulk_and_mode_override = Some(mode); self } @@ -1422,12 +1499,18 @@ impl<'a, S: Scorer> Wand<'a, S> { // 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 + if self.operator == Operator::And && !self.lead.is_empty() && self.lead.iter().all(|posting| posting.is_compressed()) + && self + .bulk_and_mode_override + .unwrap_or_else(|| *BULK_AND_MODE) + .enabled_for(self.lead.len()) { + #[cfg(test)] + { + self.bulk_and_searches += 1; + } return self.and_bulk_search(params, mask, metrics); } @@ -3398,16 +3481,50 @@ impl<'a, S: Scorer> Wand<'a, S> { } } +#[derive(Debug)] +struct RecycledPositionValues<'a> { + values: Option>, + pool: &'a RefCell>>, +} + +impl<'a> RecycledPositionValues<'a> { + fn new(values: Vec, pool: &'a RefCell>>) -> Self { + Self { + values: Some(values), + pool, + } + } + + fn as_slice(&self) -> &[u32] { + self.values + .as_deref() + .expect("position values are present until drop") + } +} + +impl Drop for RecycledPositionValues<'_> { + fn drop(&mut self) { + let values = self + .values + .take() + .expect("position values are present until drop"); + let mut pool = self.pool.borrow_mut(); + if pool.is_none() { + *pool = Some(values); + } + } +} + #[derive(Debug)] enum PositionValues<'a> { - Borrowed(&'a [u32]), + Recycled(RecycledPositionValues<'a>), Owned(Vec), } impl<'a> PositionValues<'a> { fn as_slice(&self) -> &[u32] { match self { - Self::Borrowed(values) => values, + Self::Recycled(values) => values.as_slice(), Self::Owned(values) => values.as_slice(), } } @@ -3530,6 +3647,36 @@ mod tests { } } + #[rstest] + #[case::auto("auto", Some(BulkAndMode::Auto))] + #[case::auto_case_and_whitespace(" AUTO ", Some(BulkAndMode::Auto))] + #[case::on("on", Some(BulkAndMode::On))] + #[case::on_legacy("1", Some(BulkAndMode::On))] + #[case::off("off", Some(BulkAndMode::Off))] + #[case::off_legacy("0", Some(BulkAndMode::Off))] + #[case::invalid("true", None)] + #[case::empty("", None)] + fn test_bulk_and_mode_parse(#[case] value: &str, #[case] expected: Option) { + assert_eq!(BulkAndMode::parse(value), expected); + } + + #[rstest] + #[case::auto_one(BulkAndMode::Auto, 1, false)] + #[case::auto_two(BulkAndMode::Auto, 2, true)] + #[case::auto_three(BulkAndMode::Auto, 3, true)] + #[case::auto_four(BulkAndMode::Auto, 4, false)] + #[case::on_one(BulkAndMode::On, 1, true)] + #[case::on_five(BulkAndMode::On, 5, true)] + #[case::off_two(BulkAndMode::Off, 2, false)] + #[case::off_five(BulkAndMode::Off, 5, false)] + fn test_bulk_and_mode_enabled_for( + #[case] mode: BulkAndMode, + #[case] num_clauses: usize, + #[case] expected: bool, + ) { + assert_eq!(mode.enabled_for(num_clauses), expected); + } + struct PanicQueryWeightScorer; impl Scorer for PanicQueryWeightScorer { @@ -3797,6 +3944,31 @@ mod tests { } } + #[test] + fn test_packed_position_cursors_use_independent_scratch() { + let posting_list = + generate_posting_list_with_positions(vec![0], vec![vec![1_u32, 3, 10]], 1.0, true); + let PostingList::Compressed(ref list) = posting_list else { + unreachable!("the helper was asked for a compressed posting list"); + }; + let Some(CompressedPositionStorage::SharedStream(stream)) = list.positions.as_ref() else { + panic!("compressed positions should use shared stream storage"); + }; + assert_eq!(stream.codec(), PositionStreamCodec::PackedDelta); + + let posting = PostingIterator::new(String::from("term"), 0, 0, posting_list, 1); + let first = posting.position_cursor().expect("positions must exist"); + let second = posting.position_cursor().expect("positions must exist"); + + assert_eq!(second.positions.as_slice(), &[1, 3, 10]); + assert_eq!(first.positions.as_slice(), &[1, 3, 10]); + assert!(posting.position_scratch.borrow().is_none()); + drop(second); + assert!(posting.position_scratch.borrow().is_some()); + drop(first); + assert!(posting.position_scratch.borrow().is_some()); + } + fn sorted_candidate_row_ids(candidates: Vec) -> Vec { let mut row_ids = candidates .into_iter() @@ -5425,25 +5597,46 @@ 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::and_five_clauses(false, 10, 5)] + #[case::and_six_clauses(false, 10, 6)] + #[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)] + #[case::phrase_five_clauses(true, 10, 5)] + #[case::phrase_six_clauses(true, 10, 6)] + 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 dedicated (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), + clause_docs(3, 7), + clause_docs(4, 13), + ][..num_clauses] + .to_vec(); let build_postings = || { clauses @@ -5504,29 +5697,34 @@ mod tests { rows }; - let run = |disable_bulk: bool| { + let run = |mode| { let mut wand = Wand::new( Operator::And, build_postings().into_iter(), &docs, UnitScorer, - ); - if disable_bulk { - wand = wand.disable_bulk_and(); - } - normalize( + ) + .with_bulk_and_mode(mode); + let rows = normalize( wand.search( ¶ms, Arc::new(RowAddrMask::default()), &NoOpMetricsCollector, ) .unwrap(), - ) + ); + let used_bulk = wand.bulk_and_searches > 0; + (rows, used_bulk) }; - let bulk = run(false); - let classic = run(true); + let (bulk, bulk_used) = run(BulkAndMode::On); + let (classic, classic_used) = run(BulkAndMode::Off); + let (auto, auto_used) = run(BulkAndMode::Auto); + assert!(bulk_used, "on should use bulk conjunction search"); + assert!(!classic_used, "off should use classic conjunction search"); + assert_eq!(auto_used, matches!(num_clauses, 2 | 3)); assert!(!bulk.is_empty(), "test corpus should produce matches"); assert_eq!(bulk, classic); + assert_eq!(auto, classic); } } From b3baf656212b2f570cb1c72d420cdd0f6871758a Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Wed, 15 Jul 2026 18:36:02 +0800 Subject: [PATCH 3/4] fix(fts): propagate position decoding errors --- .../src/scalar/inverted/encoding.rs | 113 +++++++-- rust/lance-index/src/scalar/inverted/wand.rs | 226 ++++++++++++------ 2 files changed, 252 insertions(+), 87 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/encoding.rs b/rust/lance-index/src/scalar/inverted/encoding.rs index a6f2107e27f..6efa6c0b7b3 100644 --- a/rust/lance-index/src/scalar/inverted/encoding.rs +++ b/rust/lance-index/src/scalar/inverted/encoding.rs @@ -702,15 +702,10 @@ fn decode_position_stream_packed_block( let mut deltas = Vec::with_capacity(total_positions); for _ in 0..full_delta_blocks { - if packed_offset >= src.len() { - return Err(Error::index( - "unexpected EOF while decoding packed position stream".to_owned(), - )); - } - let num_bits = src[packed_offset]; - packed_offset += 1; - let consumed = compressor.decompress(&src[packed_offset..], &mut packed_values, num_bits); - packed_offset += consumed; + let (num_bits, payload, next_offset) = packed_position_group(src, packed_offset)?; + let consumed = compressor.decompress(payload, &mut packed_values, num_bits); + debug_assert_eq!(consumed, payload.len()); + packed_offset = next_offset; deltas.extend_from_slice(&packed_values); } @@ -746,6 +741,39 @@ fn decode_position_stream_packed_block( Ok(()) } +fn packed_position_group(src: &[u8], offset: usize) -> Result<(u8, &[u8], usize)> { + let num_bits = *src.get(offset).ok_or_else(|| { + Error::index(format!( + "unexpected EOF reading packed position group header at byte offset {offset}; \ + stream length is {}", + src.len() + )) + })?; + if num_bits > u32::BITS as u8 { + return Err(Error::index(format!( + "invalid packed position group bit width {num_bits} at byte offset {offset}; \ + expected at most {}", + u32::BITS + ))); + } + + let payload_start = offset + .checked_add(1) + .ok_or_else(|| Error::index("packed position group offset overflow".to_owned()))?; + let payload_len = usize::from(num_bits) * BLOCK_SIZE / 8; + let payload_end = payload_start + .checked_add(payload_len) + .ok_or_else(|| Error::index("packed position group length overflow".to_owned()))?; + let payload = src.get(payload_start..payload_end).ok_or_else(|| { + Error::index(format!( + "unexpected EOF reading packed position group payload at byte offset {offset}; \ + need {payload_len} bytes after the header but stream length is {}", + src.len() + )) + })?; + Ok((num_bits, payload, payload_end)) +} + /// 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 @@ -771,6 +799,12 @@ pub(super) fn seek_packed_doc_positions( dst: &mut Vec, ) -> Result<()> { dst.clear(); + if delta_range.start > delta_range.end || delta_range.end > total_deltas { + return Err(Error::index(format!( + "invalid packed position delta range {}..{} for {total_deltas} total deltas", + delta_range.start, delta_range.end + ))); + } if delta_range.is_empty() { return Ok(()); } @@ -787,11 +821,9 @@ pub(super) fn seek_packed_doc_positions( 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); + .ok_or_else(|| Error::index("packed position group offsets are empty".to_owned()))?; + let (_, _, next_offset) = packed_position_group(src, last)?; + group_offsets.push(next_offset); } let mut previous = 0u32; @@ -813,9 +845,14 @@ pub(super) fn seek_packed_doc_positions( 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); + let offset = *group_offsets.get(group).ok_or_else(|| { + Error::index(format!( + "missing packed position group offset for group {group}; have {} offsets", + group_offsets.len() + )) + })?; + let (num_bits, payload, _) = packed_position_group(src, offset)?; + BitPacker4x::new().decompress(payload, unpacked_group, num_bits); *unpacked_group_idx = Some(group); } push_delta(unpacked_group[index % BLOCK_SIZE], dst)?; @@ -826,7 +863,11 @@ pub(super) fn seek_packed_doc_positions( if tail_cache.len() != tail_len { tail_cache.clear(); tail_cache.reserve(tail_len); - let mut offset = group_offsets[num_full_groups]; + let mut offset = *group_offsets.get(num_full_groups).ok_or_else(|| { + Error::index(format!( + "missing packed position tail offset after {num_full_groups} full groups" + )) + })?; for _ in 0..tail_len { tail_cache.push(decode_varint_u32(src, &mut offset)?); } @@ -1351,6 +1392,42 @@ mod tests { Ok(()) } + #[test] + fn test_packed_position_decoders_reject_malformed_groups() { + let frequencies = [BLOCK_SIZE as u32]; + for encoded in [&[1_u8][..], &[33_u8][..]] { + let mut decoded = Vec::new(); + assert!( + decode_position_stream_block( + encoded, + &frequencies, + PositionStreamCodec::PackedDelta, + &mut decoded, + ) + .is_err() + ); + + 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(); + assert!( + seek_packed_doc_positions( + encoded, + BLOCK_SIZE, + 0..BLOCK_SIZE, + &mut group_offsets, + &mut unpacked_group, + &mut unpacked_group_idx, + &mut tail_cache, + &mut scratch, + ) + .is_err() + ); + } + } + #[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 acd4fc29f76..8d510343678 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -14,8 +14,8 @@ use arrow::array::AsArray; use arrow::datatypes::Int32Type; use arrow_array::Array; use itertools::Itertools; -use lance_core::Result; use lance_core::utils::address::RowAddress; +use lance_core::{Error, Result}; use lance_select::RowAddrMask; use crate::metrics::MetricsCollector; @@ -679,12 +679,18 @@ impl PostingIterator { self.current_doc = current_doc; } - fn position_cursor(&self) -> Option> { + fn position_cursor(&self) -> Result> { match self.list { - PostingList::Plain(ref list) => list.positions.as_ref().map(|positions| { + PostingList::Plain(ref list) => { + let positions = list.positions.as_ref().ok_or_else(|| { + Error::index(format!( + "positions are missing for token {:?} (token id {}, query position {})", + self.token, self.token_id, self.position + )) + })?; let start = positions.value_offsets()[self.index] as usize; let end = positions.value_offsets()[self.index + 1] as usize; - PositionCursor::new( + Ok(PositionCursor::new( PositionValues::Owned( positions.values().as_primitive::().values()[start..end] .iter() @@ -692,13 +698,18 @@ impl PostingIterator { .collect(), ), self.position as i32, - ) - }), - PostingList::Compressed(ref list) => match list.positions.as_ref()? { + )) + } + PostingList::Compressed(ref list) => match list.positions.as_ref().ok_or_else(|| { + Error::index(format!( + "positions are missing for token {:?} (token id {}, query position {})", + self.token, self.token_id, self.position + )) + })? { CompressedPositionStorage::LegacyPerDoc(positions) => { let positions = positions.value(self.index); let positions = decompress_positions(positions.as_binary()); - Some(PositionCursor::new( + Ok(PositionCursor::new( PositionValues::Owned(positions), self.position as i32, )) @@ -740,7 +751,7 @@ impl PostingIterator { .borrow_mut() .take() .unwrap_or_default(); - seek_packed_doc_positions( + if let Err(error) = seek_packed_doc_positions( stream.block(block_idx), compressed.position_total_deltas, delta_start..delta_end, @@ -749,9 +760,14 @@ impl PostingIterator { &mut compressed.position_unpacked_group_idx, &mut compressed.position_tail, &mut position_values, - ) - .expect("shared position stream doc decoding should succeed"); - Some(PositionCursor::new( + ) { + *self.position_scratch.borrow_mut() = Some(position_values); + return Err(Error::index(format!( + "failed to decode positions for token {:?} (token id {}, query position {}) at posting index {}: {error}", + self.token, self.token_id, self.position, self.index + ))); + } + Ok(PositionCursor::new( PositionValues::Recycled(RecycledPositionValues::new( position_values, &self.position_scratch, @@ -768,7 +784,12 @@ impl PostingIterator { stream.codec(), &mut compressed.position_values, ) - .expect("shared position stream decoding should succeed"); + .map_err(|error| { + Error::index(format!( + "failed to decode positions for token {:?} (token id {}, query position {}) in block {block_idx}: {error}", + self.token, self.token_id, self.position + )) + })?; compressed.position_offsets.clear(); compressed .position_offsets @@ -791,7 +812,7 @@ impl PostingIterator { position_values.clear(); position_values .extend_from_slice(&compressed.position_values[start..end]); - Some(PositionCursor::new( + Ok(PositionCursor::new( PositionValues::Recycled(RecycledPositionValues::new( position_values, &self.position_scratch, @@ -1574,7 +1595,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let score = if self.operator == Operator::Or { self.advance_all_tail(doc.doc_id(), Some(doc_length), Some(&mut score)); if params.phrase_slop.is_some() - && !self.check_positions(params.phrase_slop.unwrap() as i32) + && !self.check_positions(params.phrase_slop.unwrap() as i32)? { self.push_back_leads(doc.doc_id() + 1); continue; @@ -1583,7 +1604,7 @@ impl<'a, S: Scorer> Wand<'a, S> { } else { self.advance_all_tail(doc.doc_id(), None, None); if params.phrase_slop.is_some() - && !self.check_positions(params.phrase_slop.unwrap() as i32) + && !self.check_positions(params.phrase_slop.unwrap() as i32)? { continue; } @@ -1733,7 +1754,7 @@ impl<'a, S: Scorer> Wand<'a, S> { // check positions if params.phrase_slop.is_some() - && !self.check_positions(params.phrase_slop.unwrap() as i32) + && !self.check_positions(params.phrase_slop.unwrap() as i32)? { self.advance_lead_to_head(doc_id + 1); continue; @@ -2710,9 +2731,9 @@ impl<'a, S: Scorer> Wand<'a, S> { })); } let matched = if slop == 0 { - self.check_exact_positions_bulk() + self.check_exact_positions_bulk()? } else { - self.check_positions(slop as i32) + self.check_positions(slop as i32)? }; if !matched { wins[0].pos += 1; @@ -3340,7 +3361,7 @@ impl<'a, S: Scorer> Wand<'a, S> { .collect() } - fn check_positions(&self, slop: i32) -> bool { + fn check_positions(&self, slop: i32) -> Result { if slop == 0 { return self.check_exact_positions(); } @@ -3348,8 +3369,8 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut position_iters = self .current_doc_postings() .into_iter() - .map(|posting| posting.position_cursor().expect("positions must exist")) - .collect::>(); + .map(PostingIterator::position_cursor) + .collect::>>()?; position_iters.sort_unstable_by_key(|iter| iter.position_in_query); loop { @@ -3359,7 +3380,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let last = window[0].relative_position(); let next = window[1].relative_position(); let (Some(last), Some(next)) = (last, next) else { - return false; + return Ok(false); }; let move_to = if last > next { @@ -3375,7 +3396,7 @@ impl<'a, S: Scorer> Wand<'a, S> { } if all_same { - return true; + return Ok(true); } position_iters.iter_mut().for_each(|iter| { @@ -3389,7 +3410,7 @@ impl<'a, S: Scorer> Wand<'a, S> { /// 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 { + fn check_exact_positions_bulk(&self) -> Result { const MAX_INLINE_CLAUSES: usize = 16; let num_clauses = self.lead.len(); if num_clauses > MAX_INLINE_CLAUSES { @@ -3402,7 +3423,7 @@ impl<'a, S: Scorer> Wand<'a, S> { 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"); + let cursor = posting.position_cursor()?; if cursor.len() < anchor_len { anchor_len = cursor.len(); anchor_idx = index; @@ -3424,32 +3445,32 @@ impl<'a, S: Scorer> Wand<'a, S> { } 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; + return Ok(false); }; if cursor.positions.as_slice().binary_search(&target).is_err() { continue 'anchor; } } - return true; + return Ok(true); } - false + Ok(false) } - fn check_exact_positions(&self) -> bool { + fn check_exact_positions(&self) -> Result { let mut position_iters = self .current_doc_postings() .into_iter() - .map(|posting| posting.position_cursor().expect("positions must exist")) - .collect::>(); + .map(PostingIterator::position_cursor) + .collect::>>()?; position_iters.sort_unstable_by_key(|iter| iter.len()); let Some(lead) = position_iters.first() else { - return false; + return Ok(false); }; let lead_position = lead.position_in_query; loop { let Some(anchor) = position_iters[0].absolute_position() else { - return false; + return Ok(false); }; let Some(base) = anchor.checked_sub(lead_position as u32) else { position_iters[0].advance_next(); @@ -3460,10 +3481,10 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut matched = true; for follower in position_iters.iter_mut().skip(1) { let Some(target) = base.checked_add(follower.position_in_query as u32) else { - return false; + return Ok(false); }; let Some(position) = follower.advance_to_absolute(target) else { - return false; + return Ok(false); }; if position != target { next_lead_relative = Some(position as i32 - follower.position_in_query); @@ -3473,7 +3494,7 @@ impl<'a, S: Scorer> Wand<'a, S> { } if matched { - return true; + return Ok(true); } position_iters[0].advance_to_relative(next_lead_relative.unwrap()); @@ -3602,10 +3623,11 @@ mod tests { use crate::{ metrics::{MetricsCollector, NoOpMetricsCollector}, scalar::inverted::{ - CompressedPostingList, PlainPostingList, PostingListBuilder, + CompressedPostingList, PlainPostingList, PostingListBuilder, SharedPositionStream, builder::PositionRecorder, encoding::{ compress_posting_list, compress_posting_list_with_tail_codec_and_block_size, + encode_position_stream_block_into, }, }, }; @@ -3944,21 +3966,26 @@ mod tests { } } - #[test] - fn test_packed_position_cursors_use_independent_scratch() { - let posting_list = + #[rstest] + #[case::packed_delta(PositionStreamCodec::PackedDelta)] + #[case::varint_doc_delta(PositionStreamCodec::VarintDocDelta)] + fn test_shared_position_cursors_use_independent_scratch( + #[case] codec: PositionStreamCodec, + ) -> Result<()> { + let mut posting_list = generate_posting_list_with_positions(vec![0], vec![vec![1_u32, 3, 10]], 1.0, true); - let PostingList::Compressed(ref list) = posting_list else { + let PostingList::Compressed(ref mut list) = posting_list else { unreachable!("the helper was asked for a compressed posting list"); }; - let Some(CompressedPositionStorage::SharedStream(stream)) = list.positions.as_ref() else { - panic!("compressed positions should use shared stream storage"); - }; - assert_eq!(stream.codec(), PositionStreamCodec::PackedDelta); + let mut encoded = Vec::new(); + encode_position_stream_block_into(&[1, 3, 10], &[3], codec, &mut encoded)?; + list.positions = Some(CompressedPositionStorage::SharedStream( + SharedPositionStream::new(codec, vec![0], bytes::Bytes::from(encoded)), + )); let posting = PostingIterator::new(String::from("term"), 0, 0, posting_list, 1); - let first = posting.position_cursor().expect("positions must exist"); - let second = posting.position_cursor().expect("positions must exist"); + let first = posting.position_cursor()?; + let second = posting.position_cursor()?; assert_eq!(second.positions.as_slice(), &[1, 3, 10]); assert_eq!(first.positions.as_slice(), &[1, 3, 10]); @@ -3967,6 +3994,65 @@ mod tests { assert!(posting.position_scratch.borrow().is_some()); drop(first); assert!(posting.position_scratch.borrow().is_some()); + Ok(()) + } + + #[test] + fn test_phrase_search_propagates_corrupt_packed_positions() { + let mut docs = DocSet::default(); + docs.append(0, BLOCK_SIZE as u32 + 1); + + let mut corrupt_list = generate_posting_list_with_positions( + vec![0], + vec![(0..BLOCK_SIZE as u32).collect()], + 1.0, + true, + ); + let PostingList::Compressed(ref mut list) = corrupt_list else { + unreachable!("the helper was asked for a compressed posting list"); + }; + // A bit width of one requires a 16-byte payload. Keep only the header + // to verify malformed on-disk data becomes a search error, not a panic. + list.positions = Some(CompressedPositionStorage::SharedStream( + SharedPositionStream::new( + PositionStreamCodec::PackedDelta, + vec![0], + bytes::Bytes::from_static(&[1]), + ), + )); + + let postings = vec![ + PostingIterator::new(String::from("corrupt"), 0, 0, corrupt_list, docs.len()), + PostingIterator::new( + String::from("valid"), + 1, + 1, + generate_posting_list_with_positions( + vec![0], + vec![vec![BLOCK_SIZE as u32]], + 1.0, + true, + ), + docs.len(), + ), + ]; + let mut wand = Wand::new(Operator::And, postings.into_iter(), &docs, UnitScorer); + let mut params = FtsSearchParams::default().with_limit(Some(10)); + params.phrase_slop = Some(0); + + let error = wand + .search( + ¶ms, + Arc::new(RowAddrMask::default()), + &NoOpMetricsCollector, + ) + .expect_err("corrupt packed positions should fail the phrase search"); + let message = error.to_string(); + assert!( + message.contains("packed position group payload"), + "{message}" + ); + assert!(message.contains("corrupt"), "{message}"); } fn sorted_candidate_row_ids(candidates: Vec) -> Vec { @@ -5506,8 +5592,8 @@ mod tests { let bm25 = IndexBM25Scorer::new(std::iter::empty()); let wand = Wand::new(Operator::And, postings.into_iter(), &docs, bm25); - assert!(wand.check_exact_positions()); - assert!(wand.check_positions(0)); + assert!(wand.check_exact_positions().unwrap()); + assert!(wand.check_positions(0).unwrap()); } #[rstest] @@ -5544,8 +5630,8 @@ mod tests { let bm25 = IndexBM25Scorer::new(std::iter::empty()); let wand = Wand::new(Operator::And, postings.into_iter(), &docs, bm25); - assert!(wand.check_exact_positions()); - assert!(wand.check_positions(0)); + assert!(wand.check_exact_positions().unwrap()); + assert!(wand.check_positions(0).unwrap()); } #[rstest] @@ -5584,12 +5670,12 @@ mod tests { let mut wand = Wand::new(Operator::And, postings.into_iter(), &docs, UnitScorer); let first = wand.next().unwrap().unwrap(); assert_eq!(first.0.doc_id(), 0); - assert!(!wand.check_positions(0)); + assert!(!wand.check_positions(0).unwrap()); wand.threshold = 1.5; let second = wand.next().unwrap().unwrap(); assert_eq!(second.0.doc_id(), 1); - assert!(wand.check_positions(0)); + assert!(wand.check_positions(0).unwrap()); } /// The bulk conjunction path must return exactly the classic loop's @@ -5597,20 +5683,22 @@ mod tests { /// phrase queries, across multi-block lists with heap/threshold pruning /// in play. #[rstest] - #[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::and_five_clauses(false, 10, 5)] - #[case::and_six_clauses(false, 10, 6)] - #[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)] - #[case::phrase_five_clauses(true, 10, 5)] - #[case::phrase_six_clauses(true, 10, 6)] + #[case::and_k10(false, 0, 10, 3)] + #[case::and_k3(false, 0, 3, 3)] + #[case::and_two_clauses(false, 0, 10, 2)] + #[case::and_four_clauses(false, 0, 10, 4)] + #[case::and_five_clauses(false, 0, 10, 5)] + #[case::and_six_clauses(false, 0, 10, 6)] + #[case::phrase_k10(true, 0, 10, 3)] + #[case::phrase_k3(true, 0, 3, 3)] + #[case::phrase_slop_three(true, 3, 10, 3)] + #[case::phrase_two_clauses(true, 0, 10, 2)] + #[case::phrase_four_clauses(true, 0, 10, 4)] + #[case::phrase_five_clauses(true, 0, 10, 5)] + #[case::phrase_six_clauses(true, 0, 10, 6)] fn test_bulk_and_matches_classic( #[case] phrase: bool, + #[case] slop: u32, #[case] limit: usize, #[case] num_clauses: usize, ) { @@ -5657,9 +5745,9 @@ mod tests { } }) .collect::>(); - generate_posting_list_with_positions(doc_ids.clone(), positions, 1.0, true) + generate_posting_list_with_positions(doc_ids.clone(), positions, 8.0, true) } else { - generate_posting_list(doc_ids.clone(), 1.0, None, true) + generate_posting_list(doc_ids.clone(), 8.0, None, true) }; PostingIterator::with_query_weight( format!("t{term_pos}"), @@ -5675,7 +5763,7 @@ mod tests { let mut params = FtsSearchParams::default().with_limit(Some(limit)); if phrase { - params.phrase_slop = Some(0); + params.phrase_slop = Some(slop); } let normalize = |result: Vec| { From 110540f421ab5b1015d05c6d9b6458f78abfc386 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Wed, 15 Jul 2026 18:09:58 +0800 Subject: [PATCH 4/4] perf(index): specialize SIMD merge kernels for bulk fts conjunctions --- rust/lance-index/src/scalar/inverted/wand.rs | 399 +++++++++++++++++-- 1 file changed, 359 insertions(+), 40 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 8d510343678..c887c342c99 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, ImpactScoreCache, ImpactSkipData}, - index::PositionStreamCodec, + index::{PositionStreamCodec, dequantize_doc_length}, query::Operator, scorer::{K1, idf}, }; @@ -114,6 +114,55 @@ fn bulk_and_mode_from_env() -> BulkAndMode { static BULK_AND_MODE: LazyLock = LazyLock::new(bulk_and_mode_from_env); +#[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 conservative_bm25_upper_bound(query_weight: f32) -> f32 { if query_weight <= 0.0 { @@ -2558,6 +2607,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; @@ -2566,6 +2789,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; @@ -2659,46 +2913,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; } } @@ -2714,7 +3027,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; } @@ -2723,11 +3035,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 { @@ -2736,15 +3050,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); } @@ -2758,8 +3072,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 { @@ -2776,7 +3093,6 @@ impl<'a, S: Scorer> Wand<'a, S> { self.update_threshold(kth, params.wand_factor); } } - wins[0].pos += 1; } } @@ -5172,8 +5488,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}" ); }