From 7e2acc690baffa0f0769a836741e37c5cb2a6ee6 Mon Sep 17 00:00:00 2001 From: Lu Qiu Date: Tue, 21 Jul 2026 15:38:03 -0700 Subject: [PATCH 1/7] perf(fts): resolve deferred row_ids via the cached whole ROW_ID column resolve_row_ids' slow path opened a fresh docs-file reader and issued scattered single-row read_ranges on every call, caching nothing - so any partition whose DocSet is not materialized pays a file open plus random single-row reads per query, forever. Under concurrency these per-row reads funnel through the shared scheduler/cache locks and serialize the whole search path. Use the existing row_ids_column() loader instead: one sequential read of the ROW_ID column, cached in the row_ids_col OnceCell, after which every resolve for the partition is an in-memory lookup. Costs ~8 bytes per doc resident (~2.5MB per 312k-doc partition), the same lazily-populated accounting semantics as num_tokens_col. Measured on a 320-core node, 100M-doc 42-language dataset, tier-100-200 5-term match_any, k=100, 960GiB index cache, single-node bench: - c1: resolve stage 237ms -> 0.1ms; qps 2.5 -> 3.2; mean 406 -> 316ms - c16, 300s: qps 4.1 -> 107.6 (26x), mean 3.9s -> 149ms Adds a many-partition regression test asserting exact row_id resolution. --- rust/lance-index/src/scalar/inverted/index.rs | 101 ++++++++++++++++++ .../src/scalar/inverted/lazy_docset.rs | 10 +- 2 files changed, 103 insertions(+), 8 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 9b1496f7140..10ed31085d4 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -8507,6 +8507,107 @@ mod tests { ); } + #[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. + // Exercises DeferredDocSet::resolve_row_ids' cached whole-column path. + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + const NUM_MATCHING_PARTITIONS: u64 = 40; + let mut expected_row_ids: Vec = Vec::new(); + for pid in 0..NUM_MATCHING_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.as_ref()).await.unwrap(); + } + // A partition whose only token does NOT match the query. + let mut empty_builder = + InnerBuilder::new(NUM_MATCHING_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.as_ref()).await.unwrap(); + + let all_partitions: Vec = (0..=NUM_MATCHING_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(); + + let cache = Arc::new(LanceCache::with_capacity(64 * 1024 * 1024)); + let index = InvertedIndex::load(store.clone(), None, cache.as_ref()) + .await + .unwrap(); + assert_eq!(index.partitions.len(), 41); + + // 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)); + + // Top-k smaller than the total: exactly k results. + 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))); + } + #[tokio::test] async fn test_modern_prewarm_packs_group_with_shared_posting_buffer() { let tmpdir = TempObjDir::default(); diff --git a/rust/lance-index/src/scalar/inverted/lazy_docset.rs b/rust/lance-index/src/scalar/inverted/lazy_docset.rs index 025c91cde21..4387a1dc02a 100644 --- a/rust/lance-index/src/scalar/inverted/lazy_docset.rs +++ b/rust/lance-index/src/scalar/inverted/lazy_docset.rs @@ -359,14 +359,8 @@ impl DeferredDocSet { 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()) } } From 408350f7b91755bfd1775c1ec1b19a2ceb00d72b Mon Sep 17 00:00:00 2001 From: Lu Qiu Date: Tue, 21 Jul 2026 16:06:00 -0700 Subject: [PATCH 2/7] test: assert the cached ROW_ID read contract in the resolve regression test Address review: the exact-row-ids test would also pass with the previous scattered-read implementation. Wrap the index store with a docs-file ROW_ID read counter and assert: - one full-column ROW_ID read per resolving partition on the first query - zero scattered single-row reads - zero additional docs-file ROW_ID reads on a subsequent query - the non-matching partition never reads its ROW_ID column Also preallocate the expected_row_ids vector to its known bound. --- rust/lance-index/src/scalar/inverted/index.rs | 198 +++++++++++++++++- 1 file changed, 194 insertions(+), 4 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 10ed31085d4..049be701f03 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -8507,12 +8507,171 @@ 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, + } + + 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, + } + + #[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, + } + + impl DeepSizeOf for DocsRowIdCountingStore { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.inner.deep_size_of_children(context) + } + } + + #[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 + } + } + #[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. - // Exercises DeferredDocSet::resolve_row_ids' cached whole-column path. + // 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(), @@ -8521,7 +8680,8 @@ mod tests { )); const NUM_MATCHING_PARTITIONS: u64 = 40; - let mut expected_row_ids: Vec = Vec::new(); + let mut expected_row_ids: Vec = + Vec::with_capacity((NUM_MATCHING_PARTITIONS * 3) as usize); for pid in 0..NUM_MATCHING_PARTITIONS { let mut builder = InnerBuilder::new(pid, false, TokenSetFormat::default()); builder.tokens.add("pipeline".to_owned()); @@ -8567,11 +8727,20 @@ mod tests { .unwrap(); writer.finish_with_metadata(metadata).await.unwrap(); + // 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(store.clone(), None, cache.as_ref()) + 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)); @@ -8597,7 +8766,22 @@ mod tests { assert_eq!(scores.len(), row_ids.len()); assert!(scores.iter().all(|s| *s > 0.0)); - // Top-k smaller than the total: exactly k results. + // 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(), + NUM_MATCHING_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) @@ -8606,6 +8790,12 @@ mod tests { 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(), + NUM_MATCHING_PARTITIONS as usize, + "subsequent queries must not re-read the ROW_ID column" + ); + assert_eq!(counter.scattered_reads(), 0); } #[tokio::test] From 15f2e44ca47eeb4be9a75d2a2d9aae72a335b9a1 Mon Sep 17 00:00:00 2001 From: Lu Qiu Date: Tue, 21 Jul 2026 16:41:50 -0700 Subject: [PATCH 3/7] test: exclude test-only I/O wrappers from coverage --- rust/lance-index/src/scalar/inverted/index.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 049be701f03..16517d4c4f3 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -8517,6 +8517,7 @@ mod tests { scattered_reads: std::sync::atomic::AtomicUsize, } + #[cfg_attr(coverage, coverage(off))] impl DocsRowIdReadCounter { fn full_column_reads(&self) -> usize { self.full_column_reads @@ -8533,6 +8534,7 @@ mod tests { 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 { @@ -8588,12 +8590,14 @@ mod tests { 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 { From 861bc4383f98a6e24a875f6500eec8d3243b65a8 Mon Sep 17 00:00:00 2001 From: Lu Qiu Date: Tue, 21 Jul 2026 17:07:19 -0700 Subject: [PATCH 4/7] test: gate coverage_attribute feature for lance-index coverage builds --- rust/lance-index/src/lib.rs | 1 + 1 file changed, 1 insertion(+) 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 //! From 48fa667568fe055a9c03567c06ae1c34ac975b01 Mon Sep 17 00:00:00 2001 From: Lu Qiu Date: Wed, 22 Jul 2026 11:02:34 -0700 Subject: [PATCH 5/7] perf(index): resolve fts row ids after the global top-k and cache row-id columns as weighed entries - Carry (partition slot, doc_id) through the global top-k heap and resolve row ids only for the final survivors: at most limit lookups per query instead of up to partitions * limit, and only partitions holding survivors load their ROW_ID column. - Store each partition's ROW_ID column as its own index-cache entry (DocRowIdsKey -> CachedDocRowIds) loaded through WeakLanceCache::get_or_insert_with_key: weighed at insert time so the cache accounts for the memory, independently evictable, and single-flight deduplicated on concurrent loads. The lazily populated OnceCell copy is removed. --- rust/lance-index/src/scalar/inverted/index.rs | 320 +++++++++++++----- .../src/scalar/inverted/lazy_docset.rs | 113 +++++-- 2 files changed, 322 insertions(+), 111 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 16517d4c4f3..065eecb30e7 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,7 +942,7 @@ 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 max_position = postings @@ -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, @@ -7426,11 +7467,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, @@ -8668,25 +8740,15 @@ mod tests { } } - #[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()), - )); + /// Matching-partition count used by `write_many_partition_index`. + const MANY_PARTITIONS: u64 = 40; - const NUM_MATCHING_PARTITIONS: u64 = 40; - let mut expected_row_ids: Vec = - Vec::with_capacity((NUM_MATCHING_PARTITIONS * 3) as usize); - for pid in 0..NUM_MATCHING_PARTITIONS { + /// 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)); @@ -8697,20 +8759,20 @@ mod tests { builder.docs.append(row_id, 1); expected_row_ids.push(row_id); } - builder.write(store.as_ref()).await.unwrap(); + builder.write(store).await.unwrap(); } // A partition whose only token does NOT match the query. let mut empty_builder = - InnerBuilder::new(NUM_MATCHING_PARTITIONS, false, TokenSetFormat::default()); + 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.as_ref()).await.unwrap(); + empty_builder.write(store).await.unwrap(); - let all_partitions: Vec = (0..=NUM_MATCHING_PARTITIONS).collect(); + let all_partitions: Vec = (0..=MANY_PARTITIONS).collect(); let metadata = std::collections::HashMap::from_iter(vec![ ( "partitions".to_owned(), @@ -8730,6 +8792,24 @@ mod tests { .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()); @@ -8775,7 +8855,7 @@ mod tests { // partition must not read its ROW_ID column at all. assert_eq!( counter.full_column_reads(), - NUM_MATCHING_PARTITIONS as usize, + MANY_PARTITIONS as usize, "each resolving partition reads the ROW_ID column exactly once" ); assert_eq!( @@ -8796,12 +8876,84 @@ mod tests { assert!(row_ids_k.iter().all(|id| expected_row_ids.contains(id))); assert_eq!( counter.full_column_reads(), - NUM_MATCHING_PARTITIONS as usize, + 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_modern_prewarm_packs_group_with_shared_posting_buffer() { let tmpdir = TempObjDir::default(); diff --git a/rust/lance-index/src/scalar/inverted/lazy_docset.rs b/rust/lance-index/src/scalar/inverted/lazy_docset.rs index 4387a1dc02a..4c88d57daac 100644 --- a/rust/lance-index/src/scalar/inverted/lazy_docset.rs +++ b/rust/lance-index/src/scalar/inverted/lazy_docset.rs @@ -15,6 +15,7 @@ //! only partitions that actually contribute hits pay //! `ensure_num_tokens_loaded`/`ensure_loaded`. +use std::borrow::Cow; use std::sync::Arc; use arrow::array::AsArray; @@ -22,6 +23,7 @@ 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; @@ -60,19 +62,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 full +/// DocSet 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,8 +124,6 @@ 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>, } @@ -128,13 +167,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 +179,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,12 +189,13 @@ 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(), })) } @@ -302,15 +339,30 @@ 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> { @@ -356,9 +408,6 @@ impl DeferredDocSet { { return Ok(doc_ids.iter().map(|&d| full.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 arr = self.row_ids_column().await?; Ok(doc_ids.iter().map(|&d| arr.value(d as usize)).collect()) } @@ -375,12 +424,22 @@ mod tests { #[tokio::test] async fn test_full_docset_is_a_complete_cached_snapshot() { 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]); From 985d7d323db3aabd2b21a8ef8df21db2969db767 Mon Sep 17 00:00:00 2001 From: Lu Qiu Date: Wed, 22 Jul 2026 12:43:13 -0700 Subject: [PATCH 6/7] perf(index): make the row-id cache entry the fts column's only home The DocRowIdsKey index-cache entry is now the sole owner of a modern partition's doc_id->row_id column. ensure_loaded (and therefore prewarm) loads the column into the entry and keeps a resident DocSet of num_tokens + inv only; the masked non-flat wand path borrows the entry per query via a zero-copy view; top-k resolution always reads through the entry so the LRU sees real usage. Evicting the entry now actually frees the memory, and reloads are single-flight. Flat-path shape selection shares wand's own should_flat_search predicate so the two decisions cannot drift, and flat_search asserts it got a reverse-lookup-capable DocSet. Legacy and frag-reuse partitions keep their rewritten owned copies and do not populate the raw-column entry. into_builder copies an owned DocSet out of the entry instead of materializing one on the partition. --- rust/lance-index/src/scalar/inverted/index.rs | 296 ++++++++++++++++-- .../src/scalar/inverted/lazy_docset.rs | 148 ++++++--- rust/lance-index/src/scalar/inverted/wand.rs | 33 +- 3 files changed, 408 insertions(+), 69 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 065eecb30e7..bdf7c406ca1 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -944,7 +944,7 @@ impl InvertedIndex { // row_id/num_tokens download for it. 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) @@ -2418,10 +2418,9 @@ 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(); + // into_builder rewrites every doc, so it needs a fully-owned DocSet + // (row_ids included) it can mutate; the copy is not stashed anywhere. + builder.docs = self.docs.owned_docset().await?; builder .posting_lists @@ -6407,11 +6406,80 @@ impl NumTokens { } } +/// Per-doc `row_id` storage for a [`DocSet`]. `Shared` is a zero-copy view +/// of the partition's `DocRowIdsKey` index-cache entry and must only be held +/// transiently (per-query views) — a long-lived `Shared` would pin the +/// allocation so evicting the entry could no longer free it. +#[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; counting it here too + // would double-bill the cache for the same allocation. + 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)>, @@ -6589,7 +6657,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, @@ -6598,6 +6666,45 @@ impl DocSet { } } + /// A per-query scoring view: this num-tokens-only set plus a zero-copy + /// `row_ids` borrow of the partition's `DocRowIdsKey` cache entry, for + /// the masked non-flat wand path. Dropped with the query, so the borrow + /// pins the entry's allocation only while the query is in flight. + 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 + } + + /// The resident scoring set for a modern partition: shared `num_tokens` + /// plus `inv` derived from a transient borrow of the row-ids column. + /// No `row_ids` — that column lives solely in the `DocRowIdsKey` cache + /// entry so the cache can weigh and evict it. + 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 (`inv` built, or the + /// sorted legacy `row_ids` present). The flat-search path 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 @@ -6626,7 +6733,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, @@ -6669,7 +6776,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, @@ -6690,7 +6797,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, @@ -6704,7 +6811,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(); @@ -6796,7 +6903,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)>() } @@ -8954,6 +9061,132 @@ mod tests { ); } + #[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); + + // Prewarm loads every column into its cache entry; the resident + // DocSet keeps reverse lookups but no row_ids copy. + 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()); + } + + // A warm query resolves entirely from the cache entries. + 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 reloads the column single-flight and must stay correct; + // nothing may depend on a pinned copy surviving eviction. + 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(); @@ -9831,14 +10064,37 @@ 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)); + // A flat-shaped mask (OR + tiny allow-list) gets the resident set: + // reverse lookups via `inv`, still no row_ids copy. 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); + + // A masked non-flat walk (AND never takes the flat path) gets a + // per-query view borrowing the row-ids cache entry. + 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] @@ -10636,7 +10892,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 4c88d57daac..4473c3a8167 100644 --- a/rust/lance-index/src/scalar/inverted/lazy_docset.rs +++ b/rust/lance-index/src/scalar/inverted/lazy_docset.rs @@ -14,6 +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 and the next borrower reloads it single-flight (index files are +//! immutable, so a reload is always equivalent). Legacy and frag-reuse +//! partitions rewrite the mapping at load (sort / tombstone) and keep a +//! private owned copy instead, leaving the raw-column entry unpopulated. use std::borrow::Cow; use std::sync::Arc; @@ -29,6 +37,8 @@ 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; @@ -102,8 +112,8 @@ impl CacheKey for DocRowIdsKey { /// [`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 (num_tokens and the full -/// DocSet in the `OnceCell`s below, the row_id column as its own +/// 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. @@ -124,8 +134,10 @@ pub struct DeferredDocSet { /// `NUM_TOKEN_COL` and its zero-copy scoring view carrying the cached sum, /// published together on first read. num_tokens: OnceCell, - /// Full DocSet, materialized on first `ensure_loaded`. - full: OnceCell>, + /// Resident scoring set from first `ensure_loaded`: modern partitions + /// get num_tokens + `inv` without row_ids (the column stays in the + /// cache entry); legacy / frag reuse get their full rewritten DocSet. + resident: OnceCell>, } impl std::fmt::Debug for LazyDocSet { @@ -141,10 +153,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(), } } @@ -155,7 +167,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) @@ -196,7 +208,7 @@ impl LazyDocSet { quantized_scoring, num_rows, num_tokens: OnceCell::new(), - full: OnceCell::new(), + resident: OnceCell::new(), })) } @@ -228,7 +240,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(|| { @@ -258,7 +270,13 @@ impl LazyDocSet { } } - /// Materialize the full DocSet, including row_ids. + /// Make this partition's scoring state fully resident: num_tokens, the + /// reverse-lookup `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 and the cache stays + /// free to evict it. This is what prewarm calls: afterwards queries do + /// no IO as long as the cache retains the prewarmed entries (the same + /// contract posting lists have). pub async fn ensure_loaded(&self) -> Result> { match self { Self::Loaded(l) => Ok(l.docs.clone()), @@ -266,6 +284,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 @@ -278,16 +315,39 @@ 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 of `operator` under `mask`: + /// trivial mask → num_tokens-only (survivors resolve post-merge); + /// legacy / frag reuse → the owned rewritten DocSet; flat-shaped mask + /// (same [`should_flat_search`] predicate wand uses, so the two cannot + /// disagree) → the resident set with `inv`; any other mask → a + /// per-query view borrowing the row-ids cache entry, dropped when the + /// query finishes. + 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()), + )) + } } } @@ -314,8 +374,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()) } @@ -367,25 +427,28 @@ impl DeferredDocSet { 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 (sort / tombstone), so keep a + // private owned copy; caching the raw column next to it + // would just double-store. DocSet::load( self.reader().await?, self.is_legacy, self.frag_reuse_index.clone(), ) .await? + } else { + // Borrow the cache entry once to derive `inv`; the + // column itself stays only in the entry. + 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)) @@ -396,17 +459,20 @@ 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 always read through the cache entry, which also + // keeps a hot prewarmed column recently-used in the LRU. + 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()); } let arr = self.row_ids_column().await?; Ok(doc_ids.iter().map(|&d| arr.value(d as usize)).collect()) @@ -422,7 +488,7 @@ 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( @@ -448,7 +514,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..77ca64e65c8 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,19 @@ pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { .parse::() .unwrap_or(10) }); + +/// Whether [`Wand::search`] will take the flat-search path. +/// `LazyDocSet::docs_for_wand` uses the same predicate to pick the DocSet +/// shape beforehand; the two MUST agree — a masked query scored without +/// row_ids silently skips the `mask.selected` filter — so this is the +/// single shared definition. +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 +1816,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 +2001,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. From e35c1332025e48cdecf0d414007e2b329cd36a02 Mon Sep 17 00:00:00 2001 From: Lu Qiu Date: Wed, 22 Jul 2026 13:22:45 -0700 Subject: [PATCH 7/7] refactor(index): trim comments to essential invariants --- rust/lance-index/src/scalar/inverted/index.rs | 42 ++++++----------- .../src/scalar/inverted/lazy_docset.rs | 46 +++++++------------ rust/lance-index/src/scalar/inverted/wand.rs | 8 ++-- 3 files changed, 33 insertions(+), 63 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index bdf7c406ca1..d00d270c1a2 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -2418,8 +2418,6 @@ impl InvertedPartition { self.inverted_list.block_size(), ); builder.tokens = self.tokens.into_mutable(); - // into_builder rewrites every doc, so it needs a fully-owned DocSet - // (row_ids included) it can mutate; the copy is not stashed anywhere. builder.docs = self.docs.owned_docset().await?; builder @@ -6406,10 +6404,9 @@ impl NumTokens { } } -/// Per-doc `row_id` storage for a [`DocSet`]. `Shared` is a zero-copy view -/// of the partition's `DocRowIdsKey` index-cache entry and must only be held -/// transiently (per-query views) — a long-lived `Shared` would pin the -/// allocation so evicting the entry could no longer free it. +/// `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), @@ -6437,8 +6434,7 @@ 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; counting it here too - // would double-bill the cache for the same allocation. + // Weighed by the DocRowIdsKey cache entry. Self::Shared(_) => 0, } } @@ -6666,20 +6662,16 @@ impl DocSet { } } - /// A per-query scoring view: this num-tokens-only set plus a zero-copy - /// `row_ids` borrow of the partition's `DocRowIdsKey` cache entry, for - /// the masked non-flat wand path. Dropped with the query, so the borrow - /// pins the entry's allocation only while the query is in flight. + /// 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 } - /// The resident scoring set for a modern partition: shared `num_tokens` - /// plus `inv` derived from a transient borrow of the row-ids column. - /// No `row_ids` — that column lives solely in the `DocRowIdsKey` cache - /// entry so the cache can weigh and evict it. + /// 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, @@ -6699,8 +6691,8 @@ impl DocSet { docs } - /// True iff `doc_ids()` can answer reverse lookups (`inv` built, or the - /// sorted legacy `row_ids` present). The flat-search path requires this. + /// 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() } @@ -9090,8 +9082,6 @@ mod tests { assert_eq!(counter.full_column_reads(), total_partitions); assert_eq!(counter.scattered_reads(), 0); - // Prewarm loads every column into its cache entry; the resident - // DocSet keeps reverse lookups but no row_ids copy. for partition_id in 0..=MANY_PARTITIONS { assert!( cache @@ -9108,7 +9098,6 @@ mod tests { assert!(resident.supports_reverse_lookup()); } - // A warm query resolves entirely from the cache entries. 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 @@ -9133,9 +9122,8 @@ mod tests { #[tokio::test] async fn test_row_ids_resolution_reloads_after_eviction() { - // no_cache retains nothing — the always-evicted worst case. Every - // query reloads the column single-flight and must stay correct; - // nothing may depend on a pinned copy surviving 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(), @@ -10071,8 +10059,7 @@ mod tests { .unwrap(); assert!(Arc::ptr_eq(first, &wand_view)); - // A flat-shaped mask (OR + tiny allow-list) gets the resident set: - // reverse lookups via `inv`, still no row_ids copy. + // Flat-shaped mask (OR + tiny allow-list) → resident set with `inv`. let filtered = RowAddrMask::allow_nothing(); let resident = partition .docs @@ -10084,8 +10071,7 @@ mod tests { assert!(matches!(&resident.num_tokens, NumTokens::Shared(_))); assert_eq!(resident.total_tokens_num(), 100); - // A masked non-flat walk (AND never takes the flat path) gets a - // per-query view borrowing the row-ids cache entry. + // Masked non-flat walk (AND is never flat) → per-query view. let masked_view = partition .docs .docs_for_wand(Operator::And, &filtered) diff --git a/rust/lance-index/src/scalar/inverted/lazy_docset.rs b/rust/lance-index/src/scalar/inverted/lazy_docset.rs index 4473c3a8167..faf87dd6bac 100644 --- a/rust/lance-index/src/scalar/inverted/lazy_docset.rs +++ b/rust/lance-index/src/scalar/inverted/lazy_docset.rs @@ -15,13 +15,11 @@ //! 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 +//! 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 and the next borrower reloads it single-flight (index files are -//! immutable, so a reload is always equivalent). Legacy and frag-reuse -//! partitions rewrite the mapping at load (sort / tombstone) and keep a -//! private owned copy instead, leaving the raw-column entry unpopulated. +//! 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; @@ -134,9 +132,8 @@ pub struct DeferredDocSet { /// `NUM_TOKEN_COL` and its zero-copy scoring view carrying the cached sum, /// published together on first read. num_tokens: OnceCell, - /// Resident scoring set from first `ensure_loaded`: modern partitions - /// get num_tokens + `inv` without row_ids (the column stays in the - /// cache entry); legacy / frag reuse get their full rewritten DocSet. + /// Modern partitions: num_tokens + `inv`, no row_ids (the column stays + /// in the cache entry). Legacy / frag reuse: the full rewritten DocSet. resident: OnceCell>, } @@ -270,13 +267,10 @@ impl LazyDocSet { } } - /// Make this partition's scoring state fully resident: num_tokens, the - /// reverse-lookup `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 and the cache stays - /// free to evict it. This is what prewarm calls: afterwards queries do - /// no IO as long as the cache retains the prewarmed entries (the same - /// contract posting lists have). + /// 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()), @@ -315,13 +309,10 @@ impl LazyDocSet { } } - /// Pick the DocSet shape for a wand walk of `operator` under `mask`: - /// trivial mask → num_tokens-only (survivors resolve post-merge); - /// legacy / frag reuse → the owned rewritten DocSet; flat-shaped mask - /// (same [`should_flat_search`] predicate wand uses, so the two cannot - /// disagree) → the resident set with `inv`; any other mask → a - /// per-query view borrowing the row-ids cache entry, dropped when the - /// query finishes. + /// 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, @@ -430,9 +421,7 @@ impl DeferredDocSet { .resident .get_or_try_init(|| async { let mut docs = if self.is_legacy || self.frag_reuse_index.is_some() { - // Both rewrite the mapping (sort / tombstone), so keep a - // private owned copy; caching the raw column next to it - // would just double-store. + // Both rewrite the mapping, so keep a private owned copy. DocSet::load( self.reader().await?, self.is_legacy, @@ -440,8 +429,6 @@ impl DeferredDocSet { ) .await? } else { - // Borrow the cache entry once to derive `inv`; the - // column itself stays only in the entry. let (snapshot, row_ids) = futures::try_join!(self.num_tokens_snapshot(), self.row_ids_column())?; DocSet::from_cached_num_tokens_with_inv( @@ -466,9 +453,8 @@ impl DeferredDocSet { } async fn resolve_row_ids(&self, doc_ids: &[u32]) -> Result> { - // Only legacy / frag-reuse residents carry (rewritten) row_ids. - // Modern partitions always read through the cache entry, which also - // keeps a hot prewarmed column recently-used in the LRU. + // 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() { diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index 77ca64e65c8..f730d6f9d4b 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -203,11 +203,9 @@ pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { .unwrap_or(10) }); -/// Whether [`Wand::search`] will take the flat-search path. -/// `LazyDocSet::docs_for_wand` uses the same predicate to pick the DocSet -/// shape beforehand; the two MUST agree — a masked query scored without -/// row_ids silently skips the `mask.selected` filter — so this is the -/// single shared definition. +/// 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()