diff --git a/rust/lance-index/src/lib.rs b/rust/lance-index/src/lib.rs index cf60ece2fcf..c95e2ce609c 100644 --- a/rust/lance-index/src/lib.rs +++ b/rust/lance-index/src/lib.rs @@ -1,5 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +#![cfg_attr(coverage, feature(coverage_attribute))] //! Lance secondary index library //! diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 9b1496f7140..d00d270c1a2 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -61,9 +61,8 @@ use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; use super::{InvertedIndexBuilder, InvertedIndexParams, wand::*}; use super::{ builder::{ - BLOCK_SIZE, ScoredDoc, doc_file_path, - inverted_list_schema_for_version_with_block_size_and_impacts, posting_file_path, - token_file_path, + BLOCK_SIZE, doc_file_path, inverted_list_schema_for_version_with_block_size_and_impacts, + posting_file_path, token_file_path, }, iter::PlainPostingListIterator, query::*, @@ -82,6 +81,7 @@ use crate::scalar::{ OldIndexDataFilter, RowIdRemapper, ScalarIndex, ScalarIndexParams, SearchResult, TokenQuery, UpdateCriteria, }; +use crate::vector::graph::OrderedFloat; use crate::{FtsPrewarmOptions, Index}; use crate::{prefilter::PreFilter, scalar::inverted::iter::take_fst_keys}; use std::str::FromStr; @@ -504,35 +504,6 @@ impl DeepSizeOf for InvertedIndex { } } -/// Resolve any `Pending` candidates that wand emitted via the -/// deferred-row_id path. After this returns, every entry in -/// `candidates` carries a real row_id. -async fn resolve_deferred_candidates( - docs: &LazyDocSet, - candidates: &mut [DocCandidate], -) -> Result<()> { - let pending: Vec = candidates - .iter() - .filter_map(|c| match c.addr { - CandidateAddr::Pending(d) => Some(d), - CandidateAddr::RowId(_) => None, - }) - .collect(); - if pending.is_empty() { - return Ok(()); - } - let mut iter = docs.resolve_row_ids(&pending).await?.into_iter(); - for c in candidates { - if matches!(c.addr, CandidateAddr::Pending(_)) { - let r = iter.next().ok_or_else(|| { - Error::internal("resolve_row_ids returned fewer items than requested") - })?; - c.addr = CandidateAddr::RowId(r); - } - } - Ok(()) -} - impl InvertedIndex { fn format_version(&self) -> InvertedListFormatVersion { self.format_version @@ -880,30 +851,54 @@ impl InvertedIndex { return Ok((Vec::new(), Vec::new())); } + /// Global-heap entry carrying the score plus where its row_id comes + /// from. `Pending` candidates keep (slot, doc_id) through the top-k + /// so only the final survivors pay row_id resolution. + struct MergedCandidate { + score: OrderedFloat, + /// Index into the per-query partition list for `Pending` addresses. + slot: u32, + addr: CandidateAddr, + } + + impl PartialEq for MergedCandidate { + fn eq(&self, other: &Self) -> bool { + self.score == other.score + } + } + + impl Eq for MergedCandidate {} + + impl PartialOrd for MergedCandidate { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } + } + + impl Ord for MergedCandidate { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.score.cmp(&other.score) + } + } + fn push_scored_candidate( - candidates: &mut BinaryHeap>, + candidates: &mut BinaryHeap>, limit: usize, + slot: u32, addr: CandidateAddr, score: f32, - ) -> Result<()> { - // resolve_deferred_candidates ran upstream, so every candidate - // carries a real row_id at this point. - let row_id = match addr { - CandidateAddr::RowId(r) => r, - CandidateAddr::Pending(_) => { - return Err(Error::internal( - "bm25_search post-condition: deferred candidate left unresolved", - )); - } + ) { + let candidate = MergedCandidate { + score: OrderedFloat(score), + slot, + addr, }; - if candidates.len() < limit { - candidates.push(Reverse(ScoredDoc::new(row_id, score))); + candidates.push(Reverse(candidate)); } else if candidates.peek().unwrap().0.score.0 < score { candidates.pop(); - candidates.push(Reverse(ScoredDoc::new(row_id, score))); + candidates.push(Reverse(candidate)); } - Ok(()) } let mask = prefilter.mask(); @@ -947,9 +942,9 @@ impl InvertedIndex { // No hits in this partition; its DocSet stays // unloaded, so we never pay the per-doc // row_id/num_tokens download for it. - return Result::Ok(PartitionCandidates::empty()); + return Result::Ok((part, PartitionCandidates::empty())); } - let docs_for_wand = part.docs.docs_for_wand(mask.as_ref()).await?; + let docs_for_wand = part.docs.docs_for_wand(operator, mask.as_ref()).await?; let max_position = postings .iter() .map(|posting| posting.term_index() as usize) @@ -985,23 +980,26 @@ impl InvertedIndex { std::result::Result::<_, Error>::Ok(candidates) }) .await?; - let mut partition_result = PartitionCandidates { + let partition_result = PartitionCandidates { tokens_by_position, grouped_expansions, candidates, }; - resolve_deferred_candidates(&part.docs, &mut partition_result.candidates) - .await?; - Result::Ok(partition_result) + Result::Ok((part, partition_result)) } }) .collect::>(); let mut parts = stream::iter(parts).buffer_unordered(get_num_compute_intensive_cpus()); let mut idf_cache: HashMap = HashMap::new(); - while let Some(res) = parts.try_next().await? { + // Partitions that produced candidates, indexed by the `slot` carried + // in deferred heap entries. + let mut resolving_parts: Vec> = Vec::new(); + while let Some((part, res)) = parts.try_next().await? { if res.candidates.is_empty() { continue; } + let slot = resolving_parts.len() as u32; + resolving_parts.push(part); let PartitionCandidates { tokens_by_position, grouped_expansions, @@ -1034,7 +1032,7 @@ impl InvertedIndex { score += idf_by_position[term_index as usize] * scorer.doc_weight(freq, doc_length); } - push_scored_candidate(&mut candidates, limit, addr, score)?; + push_scored_candidate(&mut candidates, limit, slot, addr, score); } } else { let grouped_positions = grouped_expansions @@ -1065,16 +1063,57 @@ impl InvertedIndex { score += term.query_weight() * scorer.doc_weight(freq, doc_length); } } - push_scored_candidate(&mut candidates, limit, addr, score)?; + push_scored_candidate(&mut candidates, limit, slot, addr, score); } } } - Ok(candidates - .into_sorted_vec() - .into_iter() - .map(|Reverse(doc)| (doc.row_id, doc.score.0)) - .unzip()) + // Resolve row_ids only for the candidates that survived the global + // top-k: group deferred survivors per partition and batch-resolve — + // at most `limit` lookups regardless of how many partitions + // contributed candidates. + /// One partition's surviving deferred candidates: positions in the + /// merged result list paired with the doc_ids to resolve. + type DeferredGroup = Vec<(usize, u32)>; + let sorted = candidates.into_sorted_vec(); + let mut row_ids = Vec::with_capacity(sorted.len()); + let mut scores = Vec::with_capacity(sorted.len()); + let mut deferred: HashMap = HashMap::new(); + for (pos, Reverse(candidate)) in sorted.into_iter().enumerate() { + scores.push(candidate.score.0); + match candidate.addr { + CandidateAddr::RowId(row_id) => row_ids.push(row_id), + CandidateAddr::Pending(doc_id) => { + deferred + .entry(candidate.slot) + .or_default() + .push((pos, doc_id)); + // Placeholder, overwritten by the batch resolution below. + row_ids.push(0); + } + } + } + if !deferred.is_empty() { + let groups = deferred + .into_iter() + .map(|(slot, entries)| (resolving_parts[slot as usize].clone(), entries)) + .collect::>(); + let resolved: Vec<(DeferredGroup, Vec)> = + stream::iter(groups.into_iter().map(|(part, entries)| async move { + let doc_ids: Vec = entries.iter().map(|&(_, doc_id)| doc_id).collect(); + let resolved = part.docs.resolve_row_ids(&doc_ids).await?; + Result::Ok((entries, resolved)) + })) + .buffer_unordered(get_num_compute_intensive_cpus()) + .try_collect() + .await?; + for (entries, resolved) in resolved { + for ((pos, _), row_id) in entries.into_iter().zip(resolved) { + row_ids[pos] = row_id; + } + } + } + Ok((row_ids, scores)) } async fn load_legacy_index( @@ -1722,6 +1761,8 @@ impl InvertedPartition { let docs = Arc::new(LazyDocSet::new( store.clone(), docs_path, + id, + WeakLanceCache::from(index_cache), num_docs, false, frag_reuse_index, @@ -2377,10 +2418,7 @@ impl InvertedPartition { self.inverted_list.block_size(), ); builder.tokens = self.tokens.into_mutable(); - // into_builder rewrites every doc, so materialize the full - // DocSet now and clone it out of the Arc. - let docs_arc = self.docs.ensure_loaded().await?; - builder.docs = (*docs_arc).clone(); + builder.docs = self.docs.owned_docset().await?; builder .posting_lists @@ -6366,11 +6404,78 @@ impl NumTokens { } } +/// `Shared` is a zero-copy view of the partition's `DocRowIdsKey` cache +/// entry; hold it only transiently, or evicting the entry can no longer +/// free the memory. +#[derive(Debug, Clone)] +enum RowIds { + Owned(Vec), + Shared(ScalarBuffer), +} + +impl Default for RowIds { + fn default() -> Self { + Self::Owned(Vec::new()) + } +} + +impl std::ops::Deref for RowIds { + type Target = [u64]; + + fn deref(&self) -> &Self::Target { + match self { + Self::Owned(values) => values, + Self::Shared(values) => values, + } + } +} + +impl DeepSizeOf for RowIds { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + match self { + Self::Owned(values) => values.deep_size_of_children(context), + // Weighed by the DocRowIdsKey cache entry. + Self::Shared(_) => 0, + } + } +} + +impl RowIds { + fn with_capacity(capacity: usize) -> Self { + Self::Owned(Vec::with_capacity(capacity)) + } + + fn into_owned(self) -> Vec { + match self { + Self::Owned(values) => values, + Self::Shared(values) => values.to_vec(), + } + } + + fn push(&mut self, value: u64) { + match self { + Self::Owned(values) => values.push(value), + Self::Shared(values) => { + let mut owned = values.to_vec(); + owned.push(value); + *self = Self::Owned(owned); + } + } + } + + fn memory_size(&self) -> usize { + match self { + Self::Owned(values) => values.capacity() * std::mem::size_of::(), + Self::Shared(_) => 0, + } + } +} + // 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)] pub struct DocSet { - row_ids: Vec, + row_ids: RowIds, num_tokens: NumTokens, // (row_id, doc_id) pairs sorted by row_id inv: Vec<(u64, u32)>, @@ -6548,7 +6653,7 @@ impl DocSet { total_tokens: u64, ) -> Self { Self { - row_ids: Vec::new(), + row_ids: RowIds::default(), num_tokens: NumTokens::Shared(num_tokens_col.values().clone()), inv: Vec::new(), total_tokens, @@ -6557,6 +6662,41 @@ impl DocSet { } } + /// Per-query view for the masked wand path: this num-tokens-only set + /// plus a transient borrow of the row-ids cache entry. + pub(crate) fn with_shared_row_ids(&self, row_ids: ScalarBuffer) -> Self { + let mut docs = self.clone(); + docs.row_ids = RowIds::Shared(row_ids); + docs + } + + /// Resident scoring set for a modern partition: shared `num_tokens` plus + /// `inv`; no row_ids — the column lives only in the cache entry. + pub(crate) fn from_cached_num_tokens_with_inv( + row_ids_col: &UInt64Array, + num_tokens_col: &arrow_array::UInt32Array, + total_tokens: u64, + ) -> Self { + let row_ids = row_ids_col.values(); + let mut inv: Vec<(u64, u32)> = row_ids + .iter() + .enumerate() + .map(|(doc_id, row_id)| (*row_id, doc_id as u32)) + .collect(); + if !row_ids.is_sorted() { + inv.sort_unstable_by_key(|entry| entry.0); + } + let mut docs = Self::from_cached_num_tokens(num_tokens_col, total_tokens); + docs.inv = inv; + docs + } + + /// True iff `doc_ids()` can answer reverse lookups; flat_search requires + /// this. + pub(crate) fn supports_reverse_lookup(&self) -> bool { + !self.inv.is_empty() || self.has_row_ids() + } + /// Build a `DocSet` from already-loaded `row_id` and `num_tokens` /// arrow columns. Lets callers that have one column already in hand /// (e.g. `LazyDocSet` after `total_tokens_num` pre-fetched @@ -6585,7 +6725,7 @@ impl DocSet { let total_tokens = num_tokens.iter().map(|&x| x as u64).sum(); return Ok(Self { - row_ids, + row_ids: RowIds::Owned(row_ids), num_tokens: NumTokens::Owned(num_tokens), inv: Vec::new(), total_tokens, @@ -6628,7 +6768,7 @@ impl DocSet { let total_tokens = num_tokens.iter().map(|&x| x as u64).sum(); return Ok(Self { - row_ids, + row_ids: RowIds::Owned(row_ids), num_tokens: NumTokens::Owned(num_tokens), inv, total_tokens, @@ -6649,7 +6789,7 @@ impl DocSet { } let total_tokens = num_tokens.iter().map(|&x| x as u64).sum(); Ok(Self { - row_ids, + row_ids: RowIds::Owned(row_ids), num_tokens: NumTokens::Owned(num_tokens), inv, total_tokens, @@ -6663,7 +6803,7 @@ impl DocSet { pub fn remap(&mut self, mapping: &RowAddrRemap) -> Vec { let mut removed = Vec::new(); let len = self.len(); - let row_ids = std::mem::replace(&mut self.row_ids, Vec::with_capacity(len)); + let row_ids = std::mem::replace(&mut self.row_ids, RowIds::with_capacity(len)).into_owned(); let num_tokens = std::mem::replace(&mut self.num_tokens, NumTokens::with_capacity(len)).into_owned(); self.invalidate_norms(); @@ -6755,7 +6895,7 @@ impl DocSet { } pub(crate) fn memory_size(&self) -> usize { - self.row_ids.capacity() * std::mem::size_of::() + self.row_ids.memory_size() + self.num_tokens.memory_size() + self.inv.capacity() * std::mem::size_of::<(u64, u32)>() } @@ -7426,11 +7566,42 @@ mod tests { use std::sync::Arc; use std::sync::atomic::{AtomicU32, Ordering}; + use crate::scalar::inverted::lazy_docset::DocRowIdsKey; use crate::scalar::inverted::tokenizer::document_tokenizer::TextTokenizer; use lance_tokenizer::{Language, SimpleTokenizer, StopWordFilter, TextAnalyzer}; use super::*; + /// Test helper: resolve any `Pending` candidates a partition-level wand + /// walk emitted, so per-partition results can be asserted on real + /// row_ids. Production resolves only global top-k survivors inside + /// `InvertedIndex::bm25_search`. + async fn resolve_deferred_candidates( + docs: &LazyDocSet, + candidates: &mut [DocCandidate], + ) -> Result<()> { + let pending: Vec = candidates + .iter() + .filter_map(|c| match c.addr { + CandidateAddr::Pending(d) => Some(d), + CandidateAddr::RowId(_) => None, + }) + .collect(); + if pending.is_empty() { + return Ok(()); + } + let mut iter = docs.resolve_row_ids(&pending).await?.into_iter(); + for c in candidates { + if matches!(c.addr, CandidateAddr::Pending(_)) { + let r = iter.next().ok_or_else(|| { + Error::internal("resolve_row_ids returned fewer items than requested") + })?; + c.addr = CandidateAddr::RowId(r); + } + } + Ok(()) + } + #[derive(Debug)] struct MetadataAccessDeniedStore { inner: Arc, @@ -8507,6 +8678,503 @@ mod tests { ); } + /// Counts docs-file ROW_ID reads so the test below can assert the + /// caching contract of deferred row_id resolution: one full-column read + /// per resolving partition on first use, zero scattered single-row reads, + /// and zero additional reads on subsequent queries. + #[derive(Debug, Default)] + struct DocsRowIdReadCounter { + full_column_reads: std::sync::atomic::AtomicUsize, + scattered_reads: std::sync::atomic::AtomicUsize, + } + + #[cfg_attr(coverage, coverage(off))] + impl DocsRowIdReadCounter { + fn full_column_reads(&self) -> usize { + self.full_column_reads + .load(std::sync::atomic::Ordering::Relaxed) + } + fn scattered_reads(&self) -> usize { + self.scattered_reads + .load(std::sync::atomic::Ordering::Relaxed) + } + } + + struct DocsRowIdCountingReader { + inner: Arc, + counter: Arc, + } + + #[cfg_attr(coverage, coverage(off))] + #[async_trait] + impl IndexReader for DocsRowIdCountingReader { + async fn read_record_batch(&self, n: u64, batch_size: u64) -> Result { + self.inner.read_record_batch(n, batch_size).await + } + async fn read_global_buffer(&self, index: u32) -> Result { + self.inner.read_global_buffer(index).await + } + async fn read_range( + &self, + range: std::ops::Range, + projection: Option<&[&str]>, + ) -> Result { + if projection + .map(|cols| cols.contains(&lance_core::ROW_ID)) + .unwrap_or(false) + { + self.counter + .full_column_reads + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + self.inner.read_range(range, projection).await + } + async fn read_ranges( + &self, + ranges: &[std::ops::Range], + projection: Option<&[&str]>, + ) -> Result { + if projection + .map(|cols| cols.contains(&lance_core::ROW_ID)) + .unwrap_or(false) + { + self.counter + .scattered_reads + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + self.inner.read_ranges(ranges, projection).await + } + async fn num_batches(&self, batch_size: u64) -> u32 { + self.inner.num_batches(batch_size).await + } + fn num_rows(&self) -> usize { + self.inner.num_rows() + } + fn schema(&self) -> &lance_core::datatypes::Schema { + self.inner.schema() + } + } + + #[derive(Debug)] + struct DocsRowIdCountingStore { + inner: Arc, + counter: Arc, + } + + #[cfg_attr(coverage, coverage(off))] + impl DeepSizeOf for DocsRowIdCountingStore { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.inner.deep_size_of_children(context) + } + } + + #[cfg_attr(coverage, coverage(off))] + #[async_trait] + impl IndexStore for DocsRowIdCountingStore { + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn clone_arc(&self) -> Arc { + Arc::new(Self { + inner: self.inner.clone(), + counter: self.counter.clone(), + }) + } + fn io_parallelism(&self) -> usize { + self.inner.io_parallelism() + } + fn with_io_priority(&self, io_priority: u64) -> Arc { + Arc::new(Self { + inner: self.inner.with_io_priority(io_priority), + counter: self.counter.clone(), + }) + } + async fn new_index_file( + &self, + name: &str, + schema: Arc, + ) -> Result> { + self.inner.new_index_file(name, schema).await + } + async fn open_index_file(&self, name: &str) -> Result> { + let reader = self.inner.open_index_file(name).await?; + if name.ends_with(DOCS_FILE) { + Ok(Arc::new(DocsRowIdCountingReader { + inner: reader, + counter: self.counter.clone(), + })) + } else { + Ok(reader) + } + } + async fn copy_index_file( + &self, + name: &str, + dest_store: &dyn IndexStore, + ) -> Result { + self.inner.copy_index_file(name, dest_store).await + } + async fn copy_index_file_to( + &self, + name: &str, + new_name: &str, + dest_store: &dyn IndexStore, + ) -> Result { + self.inner + .copy_index_file_to(name, new_name, dest_store) + .await + } + async fn rename_index_file( + &self, + name: &str, + new_name: &str, + ) -> Result { + self.inner.rename_index_file(name, new_name).await + } + async fn delete_index_file(&self, name: &str) -> Result<()> { + self.inner.delete_index_file(name).await + } + async fn list_files_with_sizes(&self) -> Result> { + self.inner.list_files_with_sizes().await + } + } + + /// Matching-partition count used by `write_many_partition_index`. + const MANY_PARTITIONS: u64 = 40; + + /// Writes `MANY_PARTITIONS` partitions whose only token is "pipeline" + /// (1-3 docs each), one extra partition whose only token does not match, + /// and the metadata file. Returns every matching doc's row_id. + async fn write_many_partition_index(store: &LanceIndexStore) -> Vec { + let mut expected_row_ids: Vec = Vec::with_capacity((MANY_PARTITIONS * 3) as usize); + for pid in 0..MANY_PARTITIONS { + let mut builder = InnerBuilder::new(pid, false, TokenSetFormat::default()); + builder.tokens.add("pipeline".to_owned()); + builder.posting_lists.push(PostingListBuilder::new(false)); + let ndocs = (pid % 3) + 1; + for d in 0..ndocs { + builder.posting_lists[0].add(d as u32, PositionRecorder::Count(1)); + let row_id = pid * 1000 + d; + builder.docs.append(row_id, 1); + expected_row_ids.push(row_id); + } + builder.write(store).await.unwrap(); + } + // A partition whose only token does NOT match the query. + let mut empty_builder = + InnerBuilder::new(MANY_PARTITIONS, false, TokenSetFormat::default()); + empty_builder.tokens.add("unrelated".to_owned()); + empty_builder + .posting_lists + .push(PostingListBuilder::new(false)); + empty_builder.posting_lists[0].add(0, PositionRecorder::Count(1)); + empty_builder.docs.append(999_999, 1); + empty_builder.write(store).await.unwrap(); + + let all_partitions: Vec = (0..=MANY_PARTITIONS).collect(); + let metadata = std::collections::HashMap::from_iter(vec![ + ( + "partitions".to_owned(), + serde_json::to_string(&all_partitions).unwrap(), + ), + ( + "params".to_owned(), + serde_json::to_string(&InvertedIndexParams::default()).unwrap(), + ), + ( + TOKEN_SET_FORMAT_KEY.to_owned(), + TokenSetFormat::default().to_string(), + ), + ]); + let mut writer = store + .new_index_file(METADATA_FILE, Arc::new(arrow_schema::Schema::empty())) + .await + .unwrap(); + writer.finish_with_metadata(metadata).await.unwrap(); + expected_row_ids + } + + #[tokio::test] + async fn test_bm25_search_many_partitions_resolves_exact_row_ids() { + // Regression test for deferred row_id resolution: with many partitions + // (including one whose only token does not match the query), a search + // must resolve every candidate's row_id exactly once and correctly. + // Also asserts the caching contract: each resolving partition reads the + // ROW_ID column exactly once (one full-column read, no scattered + // single-row reads), and subsequent queries perform no further reads. + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + let expected_row_ids = write_many_partition_index(store.as_ref()).await; + + // Load through a store that counts docs-file ROW_ID reads. + let counter = Arc::new(DocsRowIdReadCounter::default()); + let inner_store: Arc = store.clone(); + let counting_store: Arc = Arc::new(DocsRowIdCountingStore { + inner: inner_store, + counter: counter.clone(), + }); + let cache = Arc::new(LanceCache::with_capacity(64 * 1024 * 1024)); + let index = InvertedIndex::load(counting_store, None, cache.as_ref()) + .await + .unwrap(); + assert_eq!(index.partitions.len(), 41); + assert_eq!(counter.full_column_reads(), 0); + assert_eq!(counter.scattered_reads(), 0); + + // Every matching doc must come back exactly once with a correct row_id. + let tokens = Arc::new(Tokens::new(vec!["pipeline".to_owned()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(1000))); + let prefilter = Arc::new(NoFilter); + let metrics = Arc::new(NoOpMetricsCollector); + let (row_ids, scores) = index + .bm25_search( + tokens.clone(), + params, + Operator::Or, + prefilter.clone(), + metrics.clone(), + None, + ) + .await + .unwrap(); + let mut got = row_ids.clone(); + got.sort_unstable(); + let mut want = expected_row_ids.clone(); + want.sort_unstable(); + assert_eq!(got, want, "every matching doc resolved exactly once"); + assert_eq!(scores.len(), row_ids.len()); + assert!(scores.iter().all(|s| *s > 0.0)); + + // Caching contract: one full-column ROW_ID read per resolving + // partition, no scattered single-row reads. The non-matching + // partition must not read its ROW_ID column at all. + assert_eq!( + counter.full_column_reads(), + MANY_PARTITIONS as usize, + "each resolving partition reads the ROW_ID column exactly once" + ); + assert_eq!( + counter.scattered_reads(), + 0, + "resolution must not issue scattered single-row reads" + ); + + // Top-k smaller than the total: exactly k results, and the cached + // columns serve resolution with no further docs-file reads. + let params_k = Arc::new(FtsSearchParams::new().with_limit(Some(7))); + let (row_ids_k, scores_k) = index + .bm25_search(tokens, params_k, Operator::Or, prefilter, metrics, None) + .await + .unwrap(); + assert_eq!(row_ids_k.len(), 7); + assert_eq!(scores_k.len(), 7); + assert!(row_ids_k.iter().all(|id| expected_row_ids.contains(id))); + assert_eq!( + counter.full_column_reads(), + MANY_PARTITIONS as usize, + "subsequent queries must not re-read the ROW_ID column" + ); + assert_eq!(counter.scattered_reads(), 0); + } + + #[tokio::test] + async fn test_bm25_search_resolves_only_topk_survivors_and_accounts_cache() { + // Late resolution: row_ids are resolved only for candidates that + // survive the global top-k, so with k much smaller than the partition + // count, only the partitions holding survivors load their ROW_ID + // column. Each loaded column must be visible to the index cache as + // its own weighed entry. + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + let expected_row_ids = write_many_partition_index(store.as_ref()).await; + + let counter = Arc::new(DocsRowIdReadCounter::default()); + let inner_store: Arc = store.clone(); + let counting_store: Arc = Arc::new(DocsRowIdCountingStore { + inner: inner_store, + counter: counter.clone(), + }); + let cache = Arc::new(LanceCache::with_capacity(64 * 1024 * 1024)); + let index = InvertedIndex::load(counting_store, None, cache.as_ref()) + .await + .unwrap(); + + // k=3 with 40 matching partitions: survivors span at most 3 + // partitions, so at most 3 ROW_ID columns are read even though every + // matching partition produced candidates. + let tokens = Arc::new(Tokens::new(vec!["pipeline".to_owned()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(3))); + let (row_ids, scores) = index + .bm25_search( + tokens, + params, + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + assert_eq!(row_ids.len(), 3); + assert_eq!(scores.len(), 3); + assert!(row_ids.iter().all(|id| expected_row_ids.contains(id))); + let reads = counter.full_column_reads(); + assert!( + (1..=3).contains(&reads), + "only partitions with surviving candidates load their ROW_ID column, got {reads}" + ); + assert_eq!(counter.scattered_reads(), 0); + + // Accounting: every loaded column is present in the index cache as a + // DocRowIds entry (weighed at insert, evictable under pressure). + // Partition caches are scoped with a `part-{id}` prefix on load. + let mut cached_entries = 0; + for partition_id in 0..=MANY_PARTITIONS { + if cache + .with_key_prefix(format!("part-{partition_id}").as_str()) + .get_with_key(&DocRowIdsKey { partition_id }) + .await + .is_some() + { + cached_entries += 1; + } + } + assert_eq!( + cached_entries, reads, + "each loaded ROW_ID column is its own index-cache entry" + ); + } + + #[tokio::test] + async fn test_prewarm_fills_row_ids_cache_entry_without_a_second_copy() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + let expected_row_ids = write_many_partition_index(store.as_ref()).await; + + let counter = Arc::new(DocsRowIdReadCounter::default()); + let inner_store: Arc = store.clone(); + let counting_store: Arc = Arc::new(DocsRowIdCountingStore { + inner: inner_store, + counter: counter.clone(), + }); + let cache = Arc::new(LanceCache::with_capacity(64 * 1024 * 1024)); + let index = InvertedIndex::load(counting_store, None, cache.as_ref()) + .await + .unwrap(); + + index + .prewarm_with_options(&FtsPrewarmOptions::default()) + .await + .unwrap(); + let total_partitions = (MANY_PARTITIONS + 1) as usize; + assert_eq!(counter.full_column_reads(), total_partitions); + assert_eq!(counter.scattered_reads(), 0); + + for partition_id in 0..=MANY_PARTITIONS { + assert!( + cache + .with_key_prefix(format!("part-{partition_id}").as_str()) + .get_with_key(&DocRowIdsKey { partition_id }) + .await + .is_some(), + "prewarm must fill the DocRowIds entry for partition {partition_id}" + ); + } + for part in &index.partitions { + let resident = part.docs.ensure_loaded().await.unwrap(); + assert!(!resident.has_row_ids()); + assert!(resident.supports_reverse_lookup()); + } + + let tokens = Arc::new(Tokens::new(vec!["pipeline".to_owned()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(1000))); + let (row_ids, _) = index + .bm25_search( + tokens, + params, + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + let mut got = row_ids; + got.sort_unstable(); + let mut want = expected_row_ids; + want.sort_unstable(); + assert_eq!(got, want); + assert_eq!(counter.full_column_reads(), total_partitions); + assert_eq!(counter.scattered_reads(), 0); + } + + #[tokio::test] + async fn test_row_ids_resolution_reloads_after_eviction() { + // no_cache retains nothing — the always-evicted worst case: every + // query must reload the column and stay correct. + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + let expected_row_ids = write_many_partition_index(store.as_ref()).await; + + let counter = Arc::new(DocsRowIdReadCounter::default()); + let inner_store: Arc = store.clone(); + let counting_store: Arc = Arc::new(DocsRowIdCountingStore { + inner: inner_store, + counter: counter.clone(), + }); + let cache = LanceCache::no_cache(); + let index = InvertedIndex::load(counting_store, None, &cache) + .await + .unwrap(); + + let tokens = Arc::new(Tokens::new(vec!["pipeline".to_owned()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(1000))); + let mut want = expected_row_ids; + want.sort_unstable(); + let mut reads_after_first = 0; + for round in 0..2 { + let (row_ids, _) = index + .bm25_search( + tokens.clone(), + params.clone(), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + let mut got = row_ids; + got.sort_unstable(); + assert_eq!(got, want, "round {round} must stay correct"); + if round == 0 { + reads_after_first = counter.full_column_reads(); + assert!(reads_after_first > 0); + } + } + assert!( + counter.full_column_reads() > reads_after_first, + "with nothing retained, the next query reloads the column" + ); + assert_eq!(counter.scattered_reads(), 0); + } + #[tokio::test] async fn test_modern_prewarm_packs_group_with_shared_posting_buffer() { let tmpdir = TempObjDir::default(); @@ -9384,14 +10052,35 @@ mod tests { assert_eq!(first.total_tokens_num(), 100); let all_rows = RowAddrMask::all_rows(); - let wand_view = partition.docs.docs_for_wand(&all_rows).await.unwrap(); + let wand_view = partition + .docs + .docs_for_wand(Operator::Or, &all_rows) + .await + .unwrap(); assert!(Arc::ptr_eq(first, &wand_view)); + // Flat-shaped mask (OR + tiny allow-list) → resident set with `inv`. let filtered = RowAddrMask::allow_nothing(); - let full = partition.docs.docs_for_wand(&filtered).await.unwrap(); - assert!(full.has_row_ids()); - assert!(matches!(&full.num_tokens, NumTokens::Owned(_))); - assert_eq!(full.total_tokens_num(), 100); + let resident = partition + .docs + .docs_for_wand(Operator::Or, &filtered) + .await + .unwrap(); + assert!(!resident.has_row_ids()); + assert!(resident.supports_reverse_lookup()); + assert!(matches!(&resident.num_tokens, NumTokens::Shared(_))); + assert_eq!(resident.total_tokens_num(), 100); + + // Masked non-flat walk (AND is never flat) → per-query view. + let masked_view = partition + .docs + .docs_for_wand(Operator::And, &filtered) + .await + .unwrap(); + assert!(masked_view.has_row_ids()); + assert!(matches!(&masked_view.row_ids, RowIds::Shared(_))); + assert_eq!(masked_view.total_tokens_num(), 100); + assert_eq!( partition.docs.resolve_row_ids(&[0, 99]).await.unwrap(), [0, 99] @@ -10189,7 +10878,11 @@ mod tests { assert!(grouped_expansions.is_empty()); let mask = NoFilter.mask(); - let docs_for_wand = partition.docs.docs_for_wand(mask.as_ref()).await.unwrap(); + let docs_for_wand = partition + .docs + .docs_for_wand(Operator::Or, mask.as_ref()) + .await + .unwrap(); let mut candidates = partition .bm25_search( docs_for_wand.as_ref(), diff --git a/rust/lance-index/src/scalar/inverted/lazy_docset.rs b/rust/lance-index/src/scalar/inverted/lazy_docset.rs index 025c91cde21..faf87dd6bac 100644 --- a/rust/lance-index/src/scalar/inverted/lazy_docset.rs +++ b/rust/lance-index/src/scalar/inverted/lazy_docset.rs @@ -14,7 +14,14 @@ //! demand and cache. Wand scoring still needs per-doc num_tokens, but //! only partitions that actually contribute hits pay //! `ensure_num_tokens_loaded`/`ensure_loaded`. +//! +//! For modern partitions the `doc_id -> row_id` column has exactly one +//! home: the [`DocRowIdsKey`] index-cache entry. No `DocSet` keeps a +//! long-lived copy or reference, so evicting the entry really frees the +//! memory; borrowers reload it single-flight. Legacy and frag-reuse +//! partitions rewrite the mapping at load and keep a private owned copy. +use std::borrow::Cow; use std::sync::Arc; use arrow::array::AsArray; @@ -22,11 +29,14 @@ use arrow::datatypes::{UInt32Type, UInt64Type}; use arrow_array::{Array, UInt32Array, UInt64Array}; use lance_core::ROW_ID; use lance_core::Result; +use lance_core::cache::{CacheKey, WeakLanceCache}; use lance_core::deepsize::DeepSizeOf; use tokio::sync::OnceCell; use crate::scalar::RowIdRemapper; use crate::scalar::inverted::index::{DocSet, NUM_TOKEN_COL}; +use crate::scalar::inverted::query::Operator; +use crate::scalar::inverted::wand::should_flat_search; use crate::scalar::{IndexReader, IndexStore}; use lance_select::mask::RowAddrMask; @@ -60,19 +70,58 @@ struct NumTokensSnapshot { docs: Arc, } +/// A partition's full `doc_id -> row_id` column, stored as its own +/// index-cache entry so the cache weighs it at insert time and can +/// evict it independently of the partition object. Keeping it out of +/// the partition's lazily-populated fields keeps the memory visible +/// to capacity accounting. +#[derive(Debug)] +pub struct CachedDocRowIds { + pub row_ids: Arc, +} + +impl DeepSizeOf for CachedDocRowIds { + fn deep_size_of_children(&self, _ctx: &mut lance_core::deepsize::Context) -> usize { + self.row_ids.len() * std::mem::size_of::() + } +} + +/// Cache key for [`CachedDocRowIds`], scoped per partition. +#[derive(Debug, Clone)] +pub struct DocRowIdsKey { + pub partition_id: u64, +} + +impl CacheKey for DocRowIdsKey { + type ValueType = CachedDocRowIds; + + fn key(&self) -> Cow<'_, str> { + format!("doc-row-ids-{}", self.partition_id).into() + } + + fn type_name() -> &'static str { + "DocRowIds" + } +} + /// Store-backed DocSet view that loads on demand and caches. /// /// Holds the [`IndexStore`] and docs-file path rather than an open /// [`IndexReader`], so a cached partition does not pin a docs-file /// handle for its whole lifetime. The reader is re-opened on demand /// inside each column accessor and dropped when that read completes; -/// because the resulting buffers are cached in the `OnceCell`s below, -/// a contributing partition re-opens only on a cold miss, and a -/// partition that never scores never opens the docs file at all after -/// construction. +/// because the resulting buffers are cached (num_tokens and the resident +/// scoring set in the `OnceCell`s below, the row_id column as its own +/// index-cache entry), a contributing partition re-opens only on a +/// cold miss, and a partition that never scores never opens the docs +/// file at all after construction. pub struct DeferredDocSet { store: Arc, docs_path: String, + /// Scopes this partition's [`DocRowIdsKey`] in the index cache. + partition_id: u64, + /// The index cache holding the row_id column entry. + index_cache: WeakLanceCache, is_legacy: bool, frag_reuse_index: Option>, /// 256-document-block partitions score with quantized document lengths; the @@ -83,10 +132,9 @@ pub struct DeferredDocSet { /// `NUM_TOKEN_COL` and its zero-copy scoring view carrying the cached sum, /// published together on first read. num_tokens: OnceCell, - /// `ROW_ID` arrow buffer cached on first read. - row_ids_col: OnceCell>, - /// Full DocSet, materialized on first `ensure_loaded`. - full: OnceCell>, + /// Modern partitions: num_tokens + `inv`, no row_ids (the column stays + /// in the cache entry). Legacy / frag reuse: the full rewritten DocSet. + resident: OnceCell>, } impl std::fmt::Debug for LazyDocSet { @@ -102,10 +150,10 @@ impl std::fmt::Debug for LazyDocSet { .field("num_rows", &d.num_rows) .field( "total_tokens_loaded", - &(d.num_tokens.initialized() || d.full.initialized()), + &(d.num_tokens.initialized() || d.resident.initialized()), ) .field("num_tokens_loaded", &d.num_tokens.initialized()) - .field("full_loaded", &d.full.initialized()) + .field("resident_loaded", &d.resident.initialized()) .finish(), } } @@ -116,7 +164,7 @@ impl DeepSizeOf for LazyDocSet { match self { Self::Loaded(l) => l.docs.deep_size_of_children(ctx), Self::Deferred(d) => { - d.full + d.resident .get() .map(|d| d.deep_size_of_children(ctx)) .unwrap_or(0) @@ -128,13 +176,8 @@ impl DeepSizeOf for LazyDocSet { + arr.deep_size_of_children(ctx) }) .unwrap_or(0) - + d.row_ids_col - .get() - .map(|arr| { - let arr: &dyn Array = arr.as_ref(); - arr.deep_size_of_children(ctx) - }) - .unwrap_or(0) + // The row_id column is not a field here: it lives in the + // index cache as its own weighed entry (`DocRowIdsKey`). } } } @@ -145,6 +188,8 @@ impl LazyDocSet { pub fn new( store: Arc, docs_path: String, + partition_id: u64, + index_cache: WeakLanceCache, num_rows: usize, is_legacy: bool, frag_reuse_index: Option>, @@ -153,13 +198,14 @@ impl LazyDocSet { Self::Deferred(Box::new(DeferredDocSet { store, docs_path, + partition_id, + index_cache, is_legacy, frag_reuse_index, quantized_scoring, num_rows, num_tokens: OnceCell::new(), - row_ids_col: OnceCell::new(), - full: OnceCell::new(), + resident: OnceCell::new(), })) } @@ -191,7 +237,7 @@ impl LazyDocSet { match self { Self::Loaded(l) => Some(l.total_tokens), Self::Deferred(d) => d - .full + .resident .get() .map(|docs| docs.total_tokens_num()) .or_else(|| { @@ -221,7 +267,10 @@ impl LazyDocSet { } } - /// Materialize the full DocSet, including row_ids. + /// Make the scoring state resident: num_tokens, `inv`, and the row-ids + /// column loaded into its [`DocRowIdsKey`] cache entry. The returned + /// DocSet carries no row_ids for modern partitions, so nothing pins the + /// entry. Prewarm calls this. pub async fn ensure_loaded(&self) -> Result> { match self { Self::Loaded(l) => Ok(l.docs.clone()), @@ -229,6 +278,25 @@ impl LazyDocSet { } } + /// A fully-owned DocSet, row_ids included, for rebuild paths that + /// mutate it. Never stashed, so it does not pin the cache entry. + pub async fn owned_docset(&self) -> Result { + match self { + Self::Loaded(l) => Ok((*l.docs).clone()), + Self::Deferred(d) => { + if d.is_legacy || d.frag_reuse_index.is_some() { + return Ok((*d.ensure_loaded().await?).clone()); + } + let (snapshot, row_ids) = + futures::try_join!(d.num_tokens_snapshot(), d.row_ids_column())?; + let mut docs = + DocSet::from_columns(row_ids.as_ref(), snapshot.column.as_ref(), false, None)?; + docs.set_quantized_scoring(d.quantized_scoring); + Ok(docs) + } + } + } + /// Materialize a DocSet that carries num_tokens but no row_ids. /// Used by the deferred-row_id scoring path; the per-partition /// caller resolves surviving doc_ids -> row_ids post-wand via @@ -241,16 +309,36 @@ impl LazyDocSet { } } - /// Pick the right DocSet shape for a wand walk under `mask`: - /// the num_tokens-only deferred form when the mask is trivial - /// AND no FragReuseIndex needs to filter row_ids; otherwise the - /// full DocSet. Encapsulates the policy so callers don't have to - /// rederive the conditions for the targeted-read fast path. - pub async fn docs_for_wand(&self, mask: &RowAddrMask) -> Result> { - if mask.is_select_all() && !self.has_frag_reuse_remap() { - self.ensure_num_tokens_loaded().await - } else { - self.ensure_loaded().await + /// Pick the DocSet shape for a wand walk: trivial mask → num_tokens + /// only; legacy / frag reuse → owned rewritten DocSet; flat-shaped mask + /// (same [`should_flat_search`] predicate wand uses) → resident set with + /// `inv`; any other mask → per-query view borrowing the cache entry. + pub async fn docs_for_wand( + &self, + operator: Operator, + mask: &RowAddrMask, + ) -> Result> { + match self { + Self::Loaded(l) => Ok(l.docs.clone()), + Self::Deferred(d) => { + if mask.is_select_all() && !self.has_frag_reuse_remap() { + return self.ensure_num_tokens_loaded().await; + } + if d.is_legacy + || d.frag_reuse_index.is_some() + || should_flat_search(operator, mask, d.num_rows as u64) + { + return self.ensure_loaded().await; + } + let (snapshot_docs, row_ids) = { + let (snapshot, row_ids) = + futures::try_join!(d.num_tokens_snapshot(), d.row_ids_column())?; + (snapshot.docs.clone(), row_ids) + }; + Ok(Arc::new( + snapshot_docs.with_shared_row_ids(row_ids.values().clone()), + )) + } } } @@ -277,8 +365,8 @@ impl DeferredDocSet { } async fn total_tokens_num(&self) -> Result { - if let Some(full) = self.full.get() { - return Ok(full.total_tokens_num()); + if let Some(resident) = self.resident.get() { + return Ok(resident.total_tokens_num()); } Ok(self.num_tokens_snapshot().await?.docs.total_tokens_num()) } @@ -302,38 +390,52 @@ impl DeferredDocSet { .await } + /// The full `ROW_ID` column, loaded through the index cache as its own + /// weighed entry ([`DocRowIdsKey`]) rather than a field on this struct, + /// so the memory is accounted for at insert time and evictable under + /// pressure. Concurrent loads are deduped by the cache backend. async fn row_ids_column(&self) -> Result> { - self.row_ids_col - .get_or_try_init(|| async { - let reader = self.reader().await?; - let batch = reader.read_range(0..self.num_rows, Some(&[ROW_ID])).await?; - Result::Ok(Arc::new(batch[ROW_ID].as_primitive::().clone())) - }) - .await - .cloned() + let store = self.store.clone(); + let docs_path = self.docs_path.clone(); + let num_rows = self.num_rows; + let cached = self + .index_cache + .get_or_insert_with_key( + DocRowIdsKey { + partition_id: self.partition_id, + }, + || async move { + let reader = store.open_index_file(&docs_path).await?; + let batch = reader.read_range(0..num_rows, Some(&[ROW_ID])).await?; + Result::Ok(CachedDocRowIds { + row_ids: Arc::new(batch[ROW_ID].as_primitive::().clone()), + }) + }, + ) + .await?; + Ok(cached.row_ids.clone()) } async fn ensure_loaded(&self) -> Result> { let docs = self - .full + .resident .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 mut docs = if let Some(num_tokens) = self.num_tokens.get() { - let row_ids = self.row_ids_column().await?; - DocSet::from_columns( - row_ids.as_ref(), - num_tokens.column.as_ref(), - self.is_legacy, - self.frag_reuse_index.clone(), - )? - } else { + let mut docs = if self.is_legacy || self.frag_reuse_index.is_some() { + // Both rewrite the mapping, so keep a private owned copy. DocSet::load( self.reader().await?, self.is_legacy, self.frag_reuse_index.clone(), ) .await? + } else { + let (snapshot, row_ids) = + futures::try_join!(self.num_tokens_snapshot(), self.row_ids_column())?; + DocSet::from_cached_num_tokens_with_inv( + row_ids.as_ref(), + snapshot.column.as_ref(), + snapshot.docs.total_tokens_num(), + ) }; docs.set_quantized_scoring(self.quantized_scoring); Result::Ok(Arc::new(docs)) @@ -344,29 +446,22 @@ impl DeferredDocSet { } async fn ensure_num_tokens_loaded(&self) -> Result> { - if let Some(full) = self.full.get() { - return Ok(full.clone()); + if let Some(resident) = self.resident.get() { + return Ok(resident.clone()); } Ok(self.num_tokens_snapshot().await?.docs.clone()) } async fn resolve_row_ids(&self, doc_ids: &[u32]) -> Result> { - if let Some(full) = self.full.get() - && full.has_row_ids() + // Only legacy / frag-reuse residents carry (rewritten) row_ids; + // modern partitions read the cache entry, keeping it LRU-hot. + if let Some(resident) = self.resident.get() + && resident.has_row_ids() { - return Ok(doc_ids.iter().map(|&d| full.row_id(d)).collect()); + return Ok(doc_ids.iter().map(|&d| resident.row_id(d)).collect()); } - if let Some(arr) = self.row_ids_col.get() { - return Ok(doc_ids.iter().map(|&d| arr.value(d as usize)).collect()); - } - let ranges: Vec> = doc_ids - .iter() - .map(|&d| d as usize..d as usize + 1) - .collect(); - let reader = self.reader().await?; - let batch = reader.read_ranges(&ranges, Some(&[ROW_ID])).await?; - let arr = batch[ROW_ID].as_primitive::(); - Ok((0..arr.len()).map(|i| arr.value(i)).collect()) + let arr = self.row_ids_column().await?; + Ok(doc_ids.iter().map(|&d| arr.value(d as usize)).collect()) } } @@ -379,14 +474,24 @@ mod tests { use lance_io::object_store::ObjectStore; #[tokio::test] - async fn test_full_docset_is_a_complete_cached_snapshot() { + async fn test_resident_docset_serves_num_tokens_reads() { let temp_dir = TempObjDir::default(); + let cache = Arc::new(LanceCache::no_cache()); let store = Arc::new(LanceIndexStore::new( ObjectStore::local().into(), temp_dir.clone(), - Arc::new(LanceCache::no_cache()), + cache.clone(), )); - let docs = LazyDocSet::new(store, "unused".to_owned(), 3, false, None, false); + let docs = LazyDocSet::new( + store, + "unused".to_owned(), + 0, + WeakLanceCache::from(&cache), + 3, + false, + None, + false, + ); assert_eq!(docs.total_tokens_cached(), None); let row_ids = UInt64Array::from(vec![10, 20, 30]); @@ -395,7 +500,7 @@ mod tests { let LazyDocSet::Deferred(deferred) = &docs else { panic!("expected a deferred DocSet"); }; - deferred.full.set(full.clone()).unwrap(); + deferred.resident.set(full.clone()).unwrap(); let wand_docs = docs.ensure_num_tokens_loaded().await.unwrap(); assert!(Arc::ptr_eq(&wand_docs, &full)); diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index b6e2972bdbe..f730d6f9d4b 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::ops::Deref; use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::{Arc, LazyLock}; use std::{ @@ -203,6 +202,17 @@ pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { .parse::() .unwrap_or(10) }); + +/// Single shared definition of the flat-search predicate: [`Wand::search`] +/// and `LazyDocSet::docs_for_wand` MUST agree — a masked query scored +/// without row_ids silently skips the `mask.selected` filter. +pub(super) fn should_flat_search(operator: Operator, mask: &RowAddrMask, num_docs: u64) -> bool { + operator == Operator::Or + && mask.iter_addrs().is_some() + && mask.max_len().is_some_and(|num_rows_matched| { + num_rows_matched * 100 <= *FLAT_SEARCH_PERCENT_THRESHOLD * num_docs + }) +} // Bulk MAXSCORE path for top-k disjunctions (Lucene MaxScoreBulkScorer // 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 @@ -1804,15 +1814,11 @@ impl<'a, S: Scorer> Wand<'a, S> { return Ok(vec![]); } - match (mask.max_len(), mask.iter_addrs()) { - (Some(num_rows_matched), Some(row_ids)) - if self.operator == Operator::Or - && num_rows_matched * 100 - <= FLAT_SEARCH_PERCENT_THRESHOLD.deref() * self.docs.len() as u64 => - { - return self.flat_search(params, row_ids, metrics); - } - _ => {} + if should_flat_search(self.operator, &mask, self.docs.len() as u64) { + let row_ids = mask + .iter_addrs() + .expect("should_flat_search guarantees an iterable mask"); + return self.flat_search(params, row_ids, metrics); } // Top-k disjunctions over compressed lists can opt into the bulk @@ -1993,6 +1999,11 @@ impl<'a, S: Scorer> Wand<'a, S> { if limit == 0 { return Ok(vec![]); } + debug_assert!( + self.docs.is_empty() || self.docs.supports_reverse_lookup(), + "flat_search needs reverse lookups (inv, or sorted legacy row_ids); \ + the caller picked the wrong DocSet shape for this mask" + ); // we need to map the row ids to doc ids, and sort them, // because WAND PostingIterator can't go back to the previous doc id.