From 63da3b9801ff31035bb7e0114b0fe14c358ef938 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 20 Jul 2026 16:52:27 +0800 Subject: [PATCH 01/10] refactor(fts): separate document identity from row addresses --- rust/lance-index/src/scalar/inverted.rs | 2 +- .../src/scalar/inverted/builder.rs | 7 +- .../src/scalar/inverted/documents.rs | 1312 +++++++++++++++++ rust/lance-index/src/scalar/inverted/index.rs | 767 ++++++---- .../src/scalar/inverted/lazy_docset.rs | 405 ----- .../lance-index/src/scalar/inverted/scorer.rs | 22 +- rust/lance-index/src/scalar/inverted/wand.rs | 677 +++++---- 7 files changed, 2217 insertions(+), 975 deletions(-) create mode 100644 rust/lance-index/src/scalar/inverted/documents.rs delete mode 100644 rust/lance-index/src/scalar/inverted/lazy_docset.rs diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index ca27367febc..dc0e3fc131c 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -3,12 +3,12 @@ pub mod builder; mod cache_codec; +mod documents; mod encoding; mod impact; mod index; mod iter; pub mod json; -mod lazy_docset; pub mod parser; pub mod query; mod scorer; diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index 5cf754dfa3a..ec3efc4b3dd 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -1200,7 +1200,12 @@ impl InnerBuilder { let batch = docs.to_batch()?; let mut writer = store.new_index_file(path, batch.schema()).await?; writer.write_record_batch(batch).await?; - writer.finish().await + writer + .finish_with_metadata(HashMap::from([( + super::documents::TOTAL_TOKENS_KEY.to_owned(), + docs.total_tokens_num().to_string(), + )])) + .await } } diff --git a/rust/lance-index/src/scalar/inverted/documents.rs b/rust/lance-index/src/scalar/inverted/documents.rs new file mode 100644 index 00000000000..3d8b66ad6ae --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/documents.rs @@ -0,0 +1,1312 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Typed document-side state for partitioned FTS indices. +//! +//! Posting lists use dense, partition-local document identifiers. This module +//! keeps that identity separate from dataset-version row addresses so scoring +//! never has to infer which value a numeric slot represents. + +use std::ops::Range; +use std::sync::{Arc, OnceLock}; + +use arrow::buffer::ScalarBuffer; +use arrow_array::{Array, RecordBatch, UInt32Array, UInt64Array}; +use lance_core::deepsize::DeepSizeOf; +use lance_core::{Error, ROW_ID, Result}; +use lance_select::RowAddrMask; +use object_store::path::Path; +use roaring::RoaringBitmap; +use tokio::sync::OnceCell; + +use crate::scalar::{IndexReader, IndexStore, RowIdRemapper}; + +use super::index::{DocSet, NUM_TOKEN_COL, dequantize_doc_length, quantize_doc_length}; + +/// Schema metadata key persisted in every modern `docs.lance` partition. +pub(super) const TOTAL_TOKENS_KEY: &str = "total_tokens"; + +/// Maximum number of disjoint point ranges used to resolve final addresses. +const MAX_POINT_ADDRESS_RANGES: usize = 256; +/// Maximum projected address bytes used by the point-read plan. +const MAX_POINT_ADDRESS_BYTES: usize = 64 * 1024; + +/// Dense, immutable document identity inside one FTS partition. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub(super) struct DocId(u32); + +impl DocId { + pub(crate) fn new(value: u32) -> Self { + Self(value) + } + + pub(crate) fn get(self) -> u32 { + self.0 + } + + fn as_usize(self) -> usize { + self.0 as usize + } +} + +/// Immutable corpus statistics for one partition. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct PartitionStats { + pub(crate) num_docs: usize, + pub(crate) total_tokens: u64, +} + +/// Exact document lengths plus the optional quantized scoring representation. +#[derive(Debug)] +pub(super) struct DocLengths { + values: ScalarBuffer, + total_tokens: u64, + quantized_scoring: bool, + norms: OnceLock>, +} + +impl DeepSizeOf for DocLengths { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.values.deep_size_of_children(context) + + self + .norms + .get() + .map(|norms| std::mem::size_of_val(norms.as_ref())) + .unwrap_or(0) + } +} + +impl DocLengths { + fn try_new( + values: ScalarBuffer, + expected_num_docs: usize, + persisted_total_tokens: Option, + quantized_scoring: bool, + path: &str, + ) -> Result { + if values.len() != expected_num_docs { + return Err(corrupt_docs( + path, + format!( + "{NUM_TOKEN_COL} has {} rows but the file footer reports {expected_num_docs}", + values.len() + ), + )); + } + let total_tokens = values.iter().try_fold(0_u64, |total, &value| { + total + .checked_add(u64::from(value)) + .ok_or_else(|| corrupt_docs(path, format!("{NUM_TOKEN_COL} sum overflows u64"))) + })?; + if let Some(expected) = persisted_total_tokens + && expected != total_tokens + { + return Err(corrupt_docs( + path, + format!( + "{TOTAL_TOKENS_KEY} metadata is {expected}, but {NUM_TOKEN_COL} sums to {total_tokens}" + ), + )); + } + Ok(Self { + values, + total_tokens, + quantized_scoring, + norms: OnceLock::new(), + }) + } + + pub(crate) fn len(&self) -> usize { + self.values.len() + } + + pub(crate) fn total_tokens(&self) -> u64 { + self.total_tokens + } + + #[inline] + pub(crate) fn exact(&self, doc_id: DocId) -> u32 { + self.values[doc_id.as_usize()] + } + + pub(crate) fn scoring_norms(&self) -> Option<&[u8]> { + if !self.quantized_scoring { + return None; + } + Some( + self.norms + .get_or_init(|| { + self.values + .iter() + .map(|&length| quantize_doc_length(length)) + .collect() + }) + .as_ref(), + ) + } + + #[inline] + pub(crate) fn scoring(&self, doc_id: DocId) -> u32 { + match self.scoring_norms() { + Some(norms) => dequantize_doc_length(norms[doc_id.as_usize()]), + None => self.exact(doc_id), + } + } +} + +#[derive(Debug)] +enum AddressValues { + Shared(ScalarBuffer), + Owned(Vec), +} + +impl AddressValues { + fn len(&self) -> usize { + match self { + Self::Shared(values) => values.len(), + Self::Owned(values) => values.len(), + } + } + + fn get(&self, index: usize) -> u64 { + match self { + Self::Shared(values) => values[index], + Self::Owned(values) => values[index], + } + } +} + +impl DeepSizeOf for AddressValues { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + match self { + Self::Shared(values) => values.deep_size_of_children(context), + Self::Owned(values) => values.deep_size_of_children(context), + } + } +} + +/// Addresses projected into the dataset version that opened the index. +#[derive(Debug)] +pub(super) struct VersionAddressProjection { + addresses: AddressValues, + /// `None` means every slot is live. Deleted documents retain their DocId + /// slot but are absent from this bitmap. + live_docs: Option, +} + +impl DeepSizeOf for VersionAddressProjection { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.addresses.deep_size_of_children(context) + + self + .live_docs + .as_ref() + .map(|live| live.serialized_size()) + .unwrap_or(0) + } +} + +impl VersionAddressProjection { + fn try_new( + raw: &UInt64Array, + expected_num_docs: usize, + remapper: Option<&dyn RowIdRemapper>, + path: &str, + ) -> Result { + if raw.len() != expected_num_docs { + return Err(corrupt_docs( + path, + format!( + "{ROW_ID} has {} rows but the file footer reports {expected_num_docs}", + raw.len() + ), + )); + } + if raw.null_count() != 0 { + return Err(corrupt_docs(path, format!("{ROW_ID} contains null values"))); + } + + let Some(remapper) = remapper else { + return Ok(Self { + addresses: AddressValues::Shared(raw.values().clone()), + live_docs: None, + }); + }; + + let mut addresses = Vec::with_capacity(raw.len()); + let mut live_docs = RoaringBitmap::new(); + for (doc_id, &address) in raw.values().iter().enumerate() { + match remapper.remap_row_id(address) { + Some(current) => { + addresses.push(current); + live_docs.insert(doc_id as u32); + } + None => { + // The value in a dead slot is intentionally meaningless; + // callers must consult `live_docs` before reading it. + addresses.push(0); + } + } + } + Ok(Self { + addresses: AddressValues::Owned(addresses), + live_docs: Some(live_docs), + }) + } + + fn address(&self, doc_id: DocId) -> Option { + if self + .live_docs + .as_ref() + .is_some_and(|live| !live.contains(doc_id.get())) + { + None + } else { + Some(self.addresses.get(doc_id.as_usize())) + } + } + + fn visibility(&self, mask: &RowAddrMask) -> DocVisibility { + if mask.is_select_all() && self.live_docs.is_none() { + return DocVisibility::All; + } + let selected = (0..self.addresses.len()) + .filter_map(|doc_id| { + let doc_id = DocId::new(doc_id as u32); + self.address(doc_id) + .filter(|address| mask.selected(*address)) + .map(|_| doc_id.get()) + }) + .collect(); + DocVisibility::Selected(selected) + } +} + +/// Query-local selection in the partition-local DocId domain. +#[derive(Debug, Clone)] +pub(super) enum DocVisibility { + All, + Selected(RoaringBitmap), +} + +impl DocVisibility { + #[inline] + pub(crate) fn selected(&self, doc_id: DocId) -> bool { + match self { + Self::All => true, + Self::Selected(selected) => selected.contains(doc_id.get()), + } + } + + pub(crate) fn len(&self, total_docs: usize) -> usize { + match self { + Self::All => total_docs, + Self::Selected(selected) => selected.len() as usize, + } + } + + pub(crate) fn is_empty(&self) -> bool { + matches!(self, Self::Selected(selected) if selected.is_empty()) + } + + pub(crate) fn iter(&self) -> Option + '_> { + match self { + Self::All => None, + Self::Selected(selected) => Some(selected.iter().map(DocId::new)), + } + } +} + +/// Modern query-side document state for one partition. +pub(super) struct PartitionDocuments { + store: Arc, + path: String, + num_docs: usize, + persisted_total_tokens: Option, + quantized_scoring: bool, + remapper: Option>, + lengths: OnceCell>, + projection: OnceCell>, +} + +/// Load-boundary discriminator between the read-only legacy representation and +/// the typed partitioned representation. Query code dispatches on this enum +/// once; modern scoring never receives a partial legacy [`DocSet`]. +#[derive(Debug, Clone)] +pub(super) enum PartitionDocumentStore { + Legacy(Arc), + Modern(Arc), +} + +impl DeepSizeOf for PartitionDocumentStore { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + match self { + Self::Legacy(docs) => docs.deep_size_of_children(context), + Self::Modern(docs) => docs.deep_size_of_children(context), + } + } +} + +impl PartitionDocumentStore { + pub(crate) fn len(&self) -> usize { + match self { + Self::Legacy(docs) => docs.len(), + Self::Modern(docs) => docs.len(), + } + } + + pub(crate) fn legacy(&self) -> Option<&Arc> { + match self { + Self::Legacy(docs) => Some(docs), + Self::Modern(_) => None, + } + } + + pub(crate) fn modern(&self) -> Option<&Arc> { + match self { + Self::Legacy(_) => None, + Self::Modern(docs) => Some(docs), + } + } + + pub(crate) async fn stats(&self) -> Result { + match self { + Self::Legacy(docs) => Ok(PartitionStats { + num_docs: docs.len(), + total_tokens: docs.total_tokens_num(), + }), + Self::Modern(docs) => docs.stats().await, + } + } + + pub(crate) fn cached_stats(&self) -> Option { + match self { + Self::Legacy(docs) => Some(PartitionStats { + num_docs: docs.len(), + total_tokens: docs.total_tokens_num(), + }), + Self::Modern(docs) => docs.cached_stats(), + } + } + + pub(crate) async fn prewarm(&self) -> Result<()> { + match self { + Self::Legacy(_) => Ok(()), + Self::Modern(docs) => docs.prewarm().await, + } + } + + pub(crate) async fn load_build_docset(&self) -> Result { + match self { + Self::Legacy(docs) => Ok((**docs).clone()), + Self::Modern(docs) => docs.load_build_docset().await, + } + } +} + +impl std::fmt::Debug for PartitionDocuments { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PartitionDocuments") + .field("path", &self.path) + .field("num_docs", &self.num_docs) + .field("persisted_total_tokens", &self.persisted_total_tokens) + .field("lengths_loaded", &self.lengths.initialized()) + .field("projection_loaded", &self.projection.initialized()) + .finish() + } +} + +impl DeepSizeOf for PartitionDocuments { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.lengths + .get() + .map(|lengths| lengths.deep_size_of_children(context)) + .unwrap_or(0) + + self + .projection + .get() + .map(|projection| projection.deep_size_of_children(context)) + .unwrap_or(0) + } +} + +impl PartitionDocuments { + pub(crate) fn try_new( + store: Arc, + path: String, + reader: &dyn IndexReader, + remapper: Option>, + quantized_scoring: bool, + ) -> Result { + let num_docs = reader.num_rows(); + if num_docs > u32::MAX as usize { + return Err(corrupt_docs( + &path, + format!("document count {num_docs} exceeds dense DocId capacity"), + )); + } + let persisted_total_tokens = reader + .schema() + .metadata + .get(TOTAL_TOKENS_KEY) + .map(|value| { + value.parse::().map_err(|error| { + corrupt_docs( + &path, + format!("invalid {TOTAL_TOKENS_KEY} metadata value {value:?}: {error}"), + ) + }) + }) + .transpose()?; + Ok(Self { + store, + path, + num_docs, + persisted_total_tokens, + quantized_scoring, + remapper, + lengths: OnceCell::new(), + projection: OnceCell::new(), + }) + } + + pub(crate) fn len(&self) -> usize { + self.num_docs + } + + #[cfg(test)] + pub(crate) fn lengths_loaded(&self) -> bool { + self.lengths.initialized() + } + + #[cfg(test)] + pub(crate) fn projection_loaded(&self) -> bool { + self.projection.initialized() + } + + async fn reader(&self) -> Result> { + self.store.open_index_file(&self.path).await + } + + pub(crate) async fn stats(&self) -> Result { + let total_tokens = match self.persisted_total_tokens { + Some(total_tokens) => total_tokens, + None => self.lengths().await?.total_tokens(), + }; + Ok(PartitionStats { + num_docs: self.num_docs, + total_tokens, + }) + } + + pub(crate) fn cached_stats(&self) -> Option { + self.persisted_total_tokens + .or_else(|| self.lengths.get().map(|lengths| lengths.total_tokens())) + .map(|total_tokens| PartitionStats { + num_docs: self.num_docs, + total_tokens, + }) + } + + pub(crate) async fn lengths(&self) -> Result> { + self.lengths + .get_or_try_init(|| async { + let reader = self.reader().await?; + let batch = reader + .read_range(0..self.num_docs, Some(&[NUM_TOKEN_COL])) + .await?; + let column = required_u32_column(&batch, NUM_TOKEN_COL, &self.path)?; + if column.null_count() != 0 { + return Err(corrupt_docs( + &self.path, + format!("{NUM_TOKEN_COL} contains null values"), + )); + } + Ok(Arc::new(DocLengths::try_new( + column.values().clone(), + self.num_docs, + self.persisted_total_tokens, + self.quantized_scoring, + &self.path, + )?)) + }) + .await + .cloned() + } + + pub(crate) async fn projection(&self) -> Result> { + self.projection + .get_or_try_init(|| async { + let reader = self.reader().await?; + let batch = reader.read_range(0..self.num_docs, Some(&[ROW_ID])).await?; + let column = required_u64_column(&batch, ROW_ID, &self.path)?; + Ok(Arc::new(VersionAddressProjection::try_new( + column, + self.num_docs, + self.remapper.as_deref(), + &self.path, + )?)) + }) + .await + .cloned() + } + + pub(crate) async fn visibility(&self, mask: &RowAddrMask) -> Result { + if mask.is_select_all() && self.remapper.is_none() { + Ok(DocVisibility::All) + } else { + Ok(self.projection().await?.visibility(mask)) + } + } + + /// Resolve final global top-k DocIds to current row addresses. + pub(crate) async fn resolve_addresses(&self, doc_ids: &[DocId]) -> Result> { + if doc_ids.is_empty() { + return Ok(Vec::new()); + } + for doc_id in doc_ids { + if doc_id.as_usize() >= self.num_docs { + return Err(corrupt_docs( + &self.path, + format!( + "candidate DocId {} is outside [0, {})", + doc_id.get(), + self.num_docs + ), + )); + } + } + if let Some(projection) = self.projection.get() { + return doc_ids + .iter() + .map(|&doc_id| { + projection.address(doc_id).ok_or_else(|| { + corrupt_docs( + &self.path, + format!("candidate DocId {} is not live", doc_id.get()), + ) + }) + }) + .collect(); + } + if self.remapper.is_some() { + let projection = self.projection().await?; + return doc_ids + .iter() + .map(|&doc_id| { + projection.address(doc_id).ok_or_else(|| { + corrupt_docs( + &self.path, + format!("candidate DocId {} is not live", doc_id.get()), + ) + }) + }) + .collect(); + } + + let mut unique = doc_ids + .iter() + .map(|doc_id| doc_id.get()) + .collect::>(); + unique.sort_unstable(); + unique.dedup(); + let ranges = coalesce_doc_ids(&unique); + let point_bytes = unique.len().saturating_mul(std::mem::size_of::()); + let use_bulk = ranges.len() > MAX_POINT_ADDRESS_RANGES + || point_bytes > MAX_POINT_ADDRESS_BYTES + || unique.len() == self.num_docs; + let reader = self.reader().await?; + let address_column = if use_bulk { + let batch = reader.read_range(0..self.num_docs, Some(&[ROW_ID])).await?; + required_u64_column(&batch, ROW_ID, &self.path)?.clone() + } else { + let batch = reader.read_ranges(&ranges, Some(&[ROW_ID])).await?; + required_u64_column(&batch, ROW_ID, &self.path)?.clone() + }; + if address_column.null_count() != 0 { + return Err(corrupt_docs( + &self.path, + format!("{ROW_ID} contains null values"), + )); + } + let addresses = address_column.values().to_vec(); + if use_bulk && addresses.len() != self.num_docs { + return Err(corrupt_docs( + &self.path, + format!( + "{ROW_ID} bulk read returned {} rows, expected {}", + addresses.len(), + self.num_docs + ), + )); + } + if !use_bulk && addresses.len() != unique.len() { + return Err(corrupt_docs( + &self.path, + format!( + "{ROW_ID} point read returned {} rows, expected {}", + addresses.len(), + unique.len() + ), + )); + } + + if use_bulk { + Ok(doc_ids + .iter() + .map(|doc_id| addresses[doc_id.as_usize()]) + .collect()) + } else { + let by_doc_id = unique + .into_iter() + .zip(addresses) + .collect::>(); + doc_ids + .iter() + .map(|doc_id| { + by_doc_id.get(&doc_id.get()).copied().ok_or_else(|| { + Error::internal(format!( + "resolved address missing for DocId {}", + doc_id.get() + )) + }) + }) + .collect() + } + } + + /// Materialize the build-side table for rewrite/update operations. + pub(crate) async fn load_build_docset(&self) -> Result { + DocSet::load(self.reader().await?, false, self.remapper.clone()).await + } + + pub(crate) async fn prewarm(&self) -> Result<()> { + self.lengths().await?; + self.projection().await?; + Ok(()) + } +} + +fn required_u32_column<'a>( + batch: &'a RecordBatch, + name: &str, + path: &str, +) -> Result<&'a UInt32Array> { + let column = batch + .column_by_name(name) + .ok_or_else(|| corrupt_docs(path, format!("required column {name} is missing")))?; + column + .as_any() + .downcast_ref::() + .ok_or_else(|| { + corrupt_docs( + path, + format!( + "column {name} has type {}, expected UInt32", + column.data_type() + ), + ) + }) +} + +fn required_u64_column<'a>( + batch: &'a RecordBatch, + name: &str, + path: &str, +) -> Result<&'a UInt64Array> { + let column = batch + .column_by_name(name) + .ok_or_else(|| corrupt_docs(path, format!("required column {name} is missing")))?; + column + .as_any() + .downcast_ref::() + .ok_or_else(|| { + corrupt_docs( + path, + format!( + "column {name} has type {}, expected UInt64", + column.data_type() + ), + ) + }) +} + +fn coalesce_doc_ids(doc_ids: &[u32]) -> Vec> { + let mut ranges = Vec::new(); + let Some((&first, rest)) = doc_ids.split_first() else { + return ranges; + }; + let mut start = first; + let mut end = first + 1; + for &doc_id in rest { + if doc_id == end { + end += 1; + } else { + ranges.push(start as usize..end as usize); + start = doc_id; + end = doc_id + 1; + } + } + ranges.push(start as usize..end as usize); + ranges +} + +fn corrupt_docs(path: &str, message: impl Into) -> Error { + Error::corrupt_file(Path::from(path), message) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use arrow_array::{ArrayRef, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use async_trait::async_trait; + use lance_core::cache::LanceCache; + use lance_core::utils::tempfile::TempObjDir; + use lance_io::object_store::ObjectStore; + use lance_select::RowAddrTreeMap; + use roaring::RoaringTreemap; + + use crate::scalar::lance_format::LanceIndexStore; + use crate::scalar::{IndexFile, IndexWriter}; + + use super::*; + + #[derive(Debug, Default)] + struct DocumentReadCounts { + range_calls: AtomicUsize, + ranges_calls: AtomicUsize, + rows: AtomicUsize, + length_rows: AtomicUsize, + address_rows: AtomicUsize, + } + + impl DocumentReadCounts { + fn record(&self, rows: usize, projection: Option<&[&str]>) { + self.rows.fetch_add(rows, Ordering::Relaxed); + if projection.is_none_or(|columns| columns.contains(&NUM_TOKEN_COL)) { + self.length_rows.fetch_add(rows, Ordering::Relaxed); + } + if projection.is_none_or(|columns| columns.contains(&ROW_ID)) { + self.address_rows.fetch_add(rows, Ordering::Relaxed); + } + } + } + + struct CountingReader { + inner: Arc, + counts: Arc, + } + + #[async_trait] + impl IndexReader for CountingReader { + 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: Range, + projection: Option<&[&str]>, + ) -> Result { + self.counts.range_calls.fetch_add(1, Ordering::Relaxed); + self.counts.record(range.len(), projection); + self.inner.read_range(range, projection).await + } + + async fn read_ranges( + &self, + ranges: &[Range], + projection: Option<&[&str]>, + ) -> Result { + self.counts.ranges_calls.fetch_add(1, Ordering::Relaxed); + self.counts + .record(ranges.iter().map(Range::len).sum(), projection); + 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() + } + + fn file_size_bytes(&self) -> Option { + self.inner.file_size_bytes() + } + } + + #[derive(Debug)] + struct CountingStore { + inner: Arc, + target: String, + counts: Arc, + } + + impl DeepSizeOf for CountingStore { + 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 CountingStore { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn clone_arc(&self) -> Arc { + Arc::new(Self { + inner: self.inner.clone(), + target: self.target.clone(), + counts: self.counts.clone(), + }) + } + + fn io_parallelism(&self) -> usize { + self.inner.io_parallelism() + } + + 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 == self.target { + Ok(Arc::new(CountingReader { + inner: reader, + counts: self.counts.clone(), + })) + } else { + Ok(reader) + } + } + + fn with_io_priority(&self, io_priority: u64) -> Arc { + Arc::new(Self { + inner: self.inner.with_io_priority(io_priority), + target: self.target.clone(), + counts: self.counts.clone(), + }) + } + + 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 + } + } + + fn test_store() -> (TempObjDir, Arc) { + let directory = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + directory.clone(), + Arc::new(LanceCache::no_cache()), + )); + (directory, store) + } + + async fn write_documents( + store: &dyn IndexStore, + path: &str, + addresses: UInt64Array, + lengths: UInt32Array, + total_tokens: Option<&str>, + ) { + let schema = Arc::new(Schema::new(vec![ + Field::new(ROW_ID, DataType::UInt64, addresses.null_count() != 0), + Field::new(NUM_TOKEN_COL, DataType::UInt32, lengths.null_count() != 0), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(addresses) as ArrayRef, + Arc::new(lengths) as ArrayRef, + ], + ) + .unwrap(); + let mut writer = store.new_index_file(path, schema).await.unwrap(); + writer.write_record_batch(batch).await.unwrap(); + if let Some(total_tokens) = total_tokens { + writer + .finish_with_metadata(HashMap::from([( + TOTAL_TOKENS_KEY.to_owned(), + total_tokens.to_owned(), + )])) + .await + .unwrap(); + } else { + writer.finish().await.unwrap(); + } + } + + async fn open_documents( + store: Arc, + path: &str, + remapper: Option>, + ) -> Result { + let reader = store.open_index_file(path).await?; + PartitionDocuments::try_new(store, path.to_owned(), reader.as_ref(), remapper, false) + } + + fn counted_store( + inner: Arc, + target: &str, + ) -> (Arc, Arc) { + let counts = Arc::new(DocumentReadCounts::default()); + ( + Arc::new(CountingStore { + inner, + target: target.to_owned(), + counts: counts.clone(), + }), + counts, + ) + } + + #[derive(Debug)] + struct TestRemapper { + mapping: HashMap>, + } + + impl RowIdRemapper for TestRemapper { + fn remap_row_id(&self, row_id: u64) -> Option { + self.mapping.get(&row_id).copied().unwrap_or(Some(row_id)) + } + + fn remap_row_addrs_tree_map(&self, _: &RowAddrTreeMap) -> RowAddrTreeMap { + unreachable!("not used by document projection tests") + } + + fn remap_row_ids_roaring_tree_map(&self, _: &RoaringTreemap) -> RoaringTreemap { + unreachable!("not used by document projection tests") + } + + fn remap_row_ids_record_batch(&self, _: RecordBatch, _: usize) -> Result { + unreachable!("not used by document projection tests") + } + } + + #[test] + fn coalesces_sorted_doc_ids() { + assert_eq!( + coalesce_doc_ids(&[1, 2, 3, 7, 9, 10]), + vec![1..4, 7..8, 9..11] + ); + } + + #[test] + fn required_document_columns_validate_name_and_type() { + let schema = Arc::new(Schema::new(vec![Field::new( + NUM_TOKEN_COL, + DataType::UInt64, + false, + )])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(UInt64Array::from(vec![1])) as ArrayRef], + ) + .unwrap(); + + let missing = required_u64_column(&batch, ROW_ID, "docs.lance").unwrap_err(); + assert!( + missing + .to_string() + .contains("required column _rowid is missing") + ); + let wrong_type = required_u32_column(&batch, NUM_TOKEN_COL, "docs.lance").unwrap_err(); + assert!(wrong_type.to_string().contains("expected UInt32")); + } + + #[test] + fn visibility_uses_doc_ids() { + let projection = VersionAddressProjection { + addresses: AddressValues::Owned(vec![10, 20, 30]), + live_docs: Some(RoaringBitmap::from_iter([0, 2])), + }; + let mask = RowAddrMask::all_rows(); + let DocVisibility::Selected(selected) = projection.visibility(&mask) else { + panic!("remapped projection must preserve its live-doc bitmap") + }; + assert_eq!(selected.iter().collect::>(), vec![0, 2]); + } + + #[test] + fn remap_preserves_doc_id_slots_and_filters_in_current_address_domain() { + let raw = UInt64Array::from(vec![10, 20, 30, 40]); + let remapper = TestRemapper { + mapping: HashMap::from([(10, Some(100)), (20, None), (30, Some(300))]), + }; + let projection = VersionAddressProjection::try_new(&raw, 4, Some(&remapper), "docs") + .expect("valid projection"); + + assert_eq!(projection.address(DocId::new(0)), Some(100)); + assert_eq!(projection.address(DocId::new(1)), None); + assert_eq!(projection.address(DocId::new(2)), Some(300)); + assert_eq!(projection.address(DocId::new(3)), Some(40)); + + let DocVisibility::Selected(all_live) = projection.visibility(&RowAddrMask::all_rows()) + else { + panic!("a remapped projection must retain a live-doc bitmap") + }; + assert_eq!(all_live.iter().collect::>(), vec![0, 2, 3]); + + let allowed = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([100, 40])); + let DocVisibility::Selected(selected) = projection.visibility(&allowed) else { + panic!("allow-list must compile to DocIds") + }; + assert_eq!(selected.iter().collect::>(), vec![0, 3]); + + let blocked = RowAddrMask::from_block(RowAddrTreeMap::from_iter([300])); + let DocVisibility::Selected(selected) = projection.visibility(&blocked) else { + panic!("block-list must compile to DocIds") + }; + assert_eq!(selected.iter().collect::>(), vec![0, 3]); + } + + #[test] + fn doc_lengths_validate_shape_total_and_memory() { + let mismatch = DocLengths::try_new(ScalarBuffer::from(vec![2, 3]), 3, None, false, "docs") + .unwrap_err(); + assert!(mismatch.to_string().contains("2 rows")); + + let mismatch = DocLengths::try_new( + ScalarBuffer::from(vec![2, 3, 5]), + 3, + Some(11), + false, + "docs", + ) + .unwrap_err(); + assert!(mismatch.to_string().contains("sums to 10")); + + let lengths = + DocLengths::try_new(ScalarBuffer::from(vec![2, 3, 5]), 3, Some(10), true, "docs") + .unwrap(); + let before_norms = lengths.deep_size_of(); + assert_eq!(lengths.total_tokens(), 10); + assert_eq!(lengths.scoring_norms().unwrap().len(), 3); + assert_eq!(lengths.deep_size_of() - before_norms, 3); + } + + #[tokio::test] + async fn persisted_stats_are_footer_only_and_compatible_with_full_docset_reader() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + Some("10"), + ) + .await; + + let reader = store.open_index_file(path).await.unwrap(); + assert_eq!( + reader.schema().metadata.get(TOTAL_TOKENS_KEY), + Some(&"10".to_owned()) + ); + let complete = DocSet::load(reader, false, None).await.unwrap(); + assert_eq!(complete.len(), 3); + assert_eq!(complete.row_id(1), 20); + assert_eq!(complete.total_tokens_num(), 10); + + let (counting, counts) = counted_store(store, path); + let documents = open_documents(counting, path, None).await.unwrap(); + assert_eq!( + documents.stats().await.unwrap(), + PartitionStats { + num_docs: 3, + total_tokens: 10, + } + ); + assert_eq!(counts.rows.load(Ordering::Relaxed), 0); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); + } + + #[tokio::test] + async fn missing_stats_fall_back_once_to_lengths() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + None, + ) + .await; + let (counting, counts) = counted_store(store, path); + let documents = open_documents(counting, path, None).await.unwrap(); + + assert_eq!(documents.stats().await.unwrap().total_tokens, 10); + assert_eq!(documents.stats().await.unwrap().total_tokens, 10); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 1); + assert_eq!(counts.length_rows.load(Ordering::Relaxed), 3); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 0); + assert!(documents.lengths_loaded()); + assert!(!documents.projection_loaded()); + } + + #[tokio::test] + async fn invalid_or_mismatched_stats_are_corruption_not_fallback() { + let (_directory, store) = test_store(); + write_documents( + store.as_ref(), + "invalid.lance", + UInt64Array::from(vec![10]), + UInt32Array::from(vec![2]), + Some("not-a-u64"), + ) + .await; + let error = open_documents(store.clone(), "invalid.lance", None) + .await + .unwrap_err(); + assert!(error.to_string().contains("invalid total_tokens")); + + write_documents( + store.as_ref(), + "mismatch.lance", + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + Some("11"), + ) + .await; + let (counting, counts) = counted_store(store, "mismatch.lance"); + let documents = open_documents(counting, "mismatch.lance", None) + .await + .unwrap(); + assert_eq!(documents.stats().await.unwrap().total_tokens, 11); + let error = documents.lengths().await.unwrap_err(); + assert!(error.to_string().contains("sums to 10")); + assert_eq!(counts.length_rows.load(Ordering::Relaxed), 3); + assert!(!documents.lengths_loaded()); + } + + #[tokio::test] + async fn final_address_resolution_obeys_point_and_bulk_budgets() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + let num_docs = 600_u64; + write_documents( + store.as_ref(), + path, + UInt64Array::from_iter_values((0..num_docs).map(|id| id + 1_000)), + UInt32Array::from_iter_values((0..num_docs).map(|_| 1)), + Some("600"), + ) + .await; + let (counting, counts) = counted_store(store, path); + let documents = open_documents(counting, path, None).await.unwrap(); + + assert!(documents.resolve_addresses(&[]).await.unwrap().is_empty()); + assert_eq!(counts.rows.load(Ordering::Relaxed), 0); + + let point_ids = [DocId::new(5), DocId::new(6), DocId::new(10), DocId::new(5)]; + assert_eq!( + documents.resolve_addresses(&point_ids).await.unwrap(), + vec![1005, 1006, 1010, 1005] + ); + assert_eq!(counts.ranges_calls.load(Ordering::Relaxed), 1); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 0); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 3); + + let bulk_ids = (0..=512).step_by(2).map(DocId::new).collect::>(); + assert_eq!(bulk_ids.len(), MAX_POINT_ADDRESS_RANGES + 1); + let resolved = documents.resolve_addresses(&bulk_ids).await.unwrap(); + assert_eq!(resolved.first(), Some(&1000)); + assert_eq!(resolved.last(), Some(&1512)); + assert_eq!(counts.ranges_calls.load(Ordering::Relaxed), 1); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 1); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 603); + assert!(!documents.projection_loaded()); + } + + #[tokio::test] + async fn document_column_nulls_are_reported_as_corruption() { + let (_directory, store) = test_store(); + write_documents( + store.as_ref(), + "null-length.lance", + UInt64Array::from(vec![Some(10), Some(20)]), + UInt32Array::from(vec![Some(2), None]), + Some("2"), + ) + .await; + let documents = open_documents(store.clone(), "null-length.lance", None) + .await + .unwrap(); + assert!( + documents + .lengths() + .await + .unwrap_err() + .to_string() + .contains("_num_tokens contains null") + ); + + write_documents( + store.as_ref(), + "null-address.lance", + UInt64Array::from(vec![Some(10), None]), + UInt32Array::from(vec![2, 3]), + Some("5"), + ) + .await; + let documents = open_documents(store, "null-address.lance", None) + .await + .unwrap(); + assert!( + documents + .resolve_addresses(&[DocId::new(1)]) + .await + .unwrap_err() + .to_string() + .contains("_rowid contains null") + ); + } +} diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index e04d6d1924a..422c0128194 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -53,10 +53,12 @@ use std::sync::LazyLock; use tokio::{sync::OnceCell, task::spawn_blocking}; use tracing::{info, instrument, warn}; +use super::documents::{ + DocId, DocLengths, DocVisibility, PartitionDocumentStore, PartitionDocuments, +}; use super::encoding::{MAX_POSTING_BLOCK_SIZE, PositionBlockBuilder}; use super::impact::{IMPACT_LEVEL1_BLOCKS, ImpactSkipData, ImpactSkipDataBuilder}; use super::iter::PostingListIterator; -use super::lazy_docset::LazyDocSet; use super::tokenizer::{LEGACY_BLOCK_SIZE, validate_block_size}; use super::{InvertedIndexBuilder, InvertedIndexParams, wand::*}; use super::{ @@ -272,13 +274,13 @@ pub fn validate_format_version_block_size( } #[derive(Debug)] -struct PartitionCandidates { +struct PartitionCandidates { tokens_by_position: Vec, grouped_expansions: Vec, - candidates: Vec, + candidates: Vec>, } -impl PartitionCandidates { +impl PartitionCandidates { fn empty() -> Self { Self { tokens_by_position: Vec::new(), @@ -288,6 +290,78 @@ impl PartitionCandidates { } } +fn push_scored_key( + candidates: &mut BinaryHeap>, + limit: usize, + key: u64, + score: f32, +) { + if candidates.len() < limit { + candidates.push(Reverse(ScoredDoc::new(key, score))); + } else if candidates + .peek() + .is_some_and(|candidate| candidate.0.score.0 < score) + { + candidates.pop(); + candidates.push(Reverse(ScoredDoc::new(key, score))); + } +} + +fn rescore_partition_candidates( + partition: PartitionCandidates, + scorer: &MemBM25Scorer, + idf_cache: &mut HashMap, +) -> Vec<(C, f32)> { + let PartitionCandidates { + tokens_by_position, + grouped_expansions, + candidates, + } = partition; + let idf_by_position = tokens_by_position + .iter() + .map(|token| { + *idf_cache + .entry(token.clone()) + .or_insert_with(|| scorer.query_weight(token)) + }) + .collect::>(); + let grouped_positions = grouped_expansions + .iter() + .map(|group| group.position) + .collect::>(); + + candidates + .into_iter() + .map( + |DocCandidate { + document, + posting_doc_id, + freqs, + doc_length, + }| { + let mut score = 0.0; + for (term_index, freq) in freqs { + if grouped_positions.contains(&term_index) { + continue; + } + debug_assert!((term_index as usize) < idf_by_position.len()); + score += + idf_by_position[term_index as usize] * scorer.doc_weight(freq, doc_length); + } + for group in &grouped_expansions { + for term in group.terms.iter() { + let Some(freq) = term.frequency(posting_doc_id) else { + continue; + }; + score += term.query_weight() * scorer.doc_weight(freq, doc_length); + } + } + (document, score) + }, + ) + .collect() +} + #[derive(Debug)] struct LoadedPostings { postings: Vec, @@ -296,6 +370,27 @@ struct LoadedPostings { exact_scoring_required: bool, } +enum LoadedDocLengths { + Legacy(Arc), + Modern(Arc), +} + +impl LoadedDocLengths { + fn scoring_num_tokens(&self, doc_id: u32) -> u32 { + match self { + Self::Legacy(docs) => docs.scoring_num_tokens(doc_id), + Self::Modern(lengths) => lengths.scoring(DocId::new(doc_id)), + } + } + + fn num_tokens_by_row_id(&self, row_id: u64) -> u32 { + match self { + Self::Legacy(docs) => docs.num_tokens_by_row_id(row_id), + Self::Modern(_) => unreachable!("modern posting lists use dense DocIds"), + } + } +} + impl LoadedPostings { fn empty() -> Self { Self { @@ -508,35 +603,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 @@ -725,29 +791,36 @@ impl InvertedIndex { Ok((total_tokens, num_docs, token_docs)) } - /// Aggregate per-partition `total_tokens` and `num_docs` across the - /// index. `len` is cheap (no IO); `total_tokens_num` reads only the - /// num_tokens column the first time per partition and caches it on - /// `LazyDocSet`. Avoids materializing the full DocSet just to get - /// these two scalars. + /// Aggregate immutable per-partition corpus statistics. New modern files + /// read both values from the already-opened docs footer; older partitioned + /// files scan `_num_tokens` once as a compatibility fallback. async fn aggregate_corpus_stats(&self) -> Result<(u64, usize)> { self.corpus_stats .get_or_try_init(|| async { let io_parallelism = self.store.io_parallelism(); - let num_docs: usize = self.partitions.iter().map(|p| p.docs.len()).sum(); let futures = self .partitions .iter() .map(|p| { - let docs = p.docs.clone(); - async move { docs.total_tokens_num().await } + let part = p.clone(); + async move { part.docs.stats().await } }) .collect::>(); - let totals: Vec = stream::iter(futures) + let stats = stream::iter(futures) .buffer_unordered(io_parallelism) - .try_collect() + .try_collect::>() .await?; - Ok((totals.into_iter().sum(), num_docs)) + let mut total_tokens = 0_u64; + let mut num_docs = 0_usize; + for stat in stats { + total_tokens = total_tokens + .checked_add(stat.total_tokens) + .ok_or_else(|| Error::index("FTS corpus token count overflows u64"))?; + num_docs = num_docs + .checked_add(stat.num_docs) + .ok_or_else(|| Error::index("FTS corpus document count overflows usize"))?; + } + Ok((total_tokens, num_docs)) }) .await .copied() @@ -883,42 +956,48 @@ impl InvertedIndex { if limit == 0 { return Ok((Vec::new(), Vec::new())); } - - fn push_scored_candidate( - candidates: &mut BinaryHeap>, - limit: usize, - 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", - )); - } - }; - - if candidates.len() < limit { - candidates.push(Reverse(ScoredDoc::new(row_id, score))); - } else if candidates.peek().unwrap().0.score.0 < score { - candidates.pop(); - candidates.push(Reverse(ScoredDoc::new(row_id, score))); - } - Ok(()) - } - let mask = prefilter.mask(); + if self.is_legacy() { + self.bm25_search_legacy( + tokens, + params, + operator, + mask, + metrics, + scorer, + impact_scorer, + limit, + ) + .await + } else { + self.bm25_search_modern( + tokens, + params, + operator, + mask, + metrics, + scorer, + impact_scorer, + limit, + ) + .await + } + } - let mut candidates = BinaryHeap::new(); - // Shared top-k floor across this query's partitions. Seeded to -inf so - // the first real score wins; each partition publishes its local k-th - // and prunes against the running global k-th (a lower bound on the true - // global k-th - see `Wand::shared_threshold`). + #[allow(clippy::too_many_arguments)] + async fn bm25_search_legacy( + &self, + tokens: Arc, + params: Arc, + operator: Operator, + mask: Arc, + metrics: Arc, + scorer: &MemBM25Scorer, + impact_scorer: Arc, + limit: usize, + ) -> Result<(Vec, Vec)> { let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); - let legacy_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + let local_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); let parts = self .partitions .iter() @@ -930,9 +1009,14 @@ impl InvertedIndex { let metrics = metrics.clone(); let impact_scorer = impact_scorer.clone(); let impact_shared_threshold = impact_shared_threshold.clone(); - let legacy_shared_threshold = legacy_shared_threshold.clone(); + let local_shared_threshold = local_shared_threshold.clone(); async move { - let loaded_postings = part + let LoadedPostings { + postings, + grouped_expansions, + impact_safe, + exact_scoring_required, + } = part .load_posting_lists( tokens.as_ref(), params.as_ref(), @@ -941,19 +1025,9 @@ impl InvertedIndex { metrics.as_ref(), ) .await?; - let LoadedPostings { - postings, - grouped_expansions, - impact_safe, - exact_scoring_required, - } = loaded_postings; if postings.is_empty() { - // 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()); } - let docs_for_wand = part.docs.docs_for_wand(mask.as_ref()).await?; let max_position = postings .iter() .map(|posting| posting.term_index() as usize) @@ -961,124 +1035,203 @@ impl InvertedIndex { .unwrap_or_default(); let mut tokens_by_position = vec![String::new(); max_position + 1]; for posting in &postings { - let idx = posting.term_index() as usize; - tokens_by_position[idx] = posting.token().to_owned(); + tokens_by_position[posting.term_index() as usize] = + posting.token().to_owned(); } - let params = params.clone(); - let mask = mask.clone(); - let metrics = metrics.clone(); - let part_for_wand = part.clone(); + let docs = part.docs.legacy().cloned().ok_or_else(|| { + Error::internal("legacy index contains modern partition documents") + })?; let use_global_scorer = impact_safe || exact_scoring_required; - let partition_threshold = if use_global_scorer { + let threshold = if use_global_scorer { impact_shared_threshold } else { - legacy_shared_threshold + local_shared_threshold }; let wand_scorer = use_global_scorer.then(|| impact_scorer.clone()); + let part_for_wand = part.clone(); let candidates = spawn_cpu(move || { - let candidates = part_for_wand.bm25_search( - docs_for_wand.as_ref(), + part_for_wand.bm25_search_legacy( + docs.as_ref(), params.as_ref(), operator, - mask, + mask.as_ref(), postings, wand_scorer, metrics.as_ref(), - partition_threshold, - )?; - std::result::Result::<_, Error>::Ok(candidates) + threshold, + ) }) .await?; - let mut partition_result = PartitionCandidates { + Result::Ok(PartitionCandidates { tokens_by_position, grouped_expansions, candidates, - }; - resolve_deferred_candidates(&part.docs, &mut partition_result.candidates) - .await?; - Result::Ok(partition_result) + }) } }) .collect::>(); + + let mut ranked = BinaryHeap::new(); + let mut idf_cache = HashMap::new(); 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? { - if res.candidates.is_empty() { - continue; - } - let PartitionCandidates { - tokens_by_position, - grouped_expansions, - candidates: part_candidates, - } = res; - let mut idf_by_position = Vec::with_capacity(tokens_by_position.len()); - for token in &tokens_by_position { - let idf_weight = match idf_cache.get(token) { - Some(weight) => *weight, - None => { - let weight = scorer.query_weight(token); - idf_cache.insert(token.clone(), weight); - weight - } - }; - idf_by_position.push(idf_weight); + while let Some(partition) = parts.try_next().await? { + for (row_id, score) in rescore_partition_candidates(partition, scorer, &mut idf_cache) { + push_scored_key(&mut ranked, limit, row_id, score); } + } + Ok(ranked + .into_sorted_vec() + .into_iter() + .map(|Reverse(doc)| (doc.row_id, doc.score.0)) + .unzip()) + } - if grouped_expansions.is_empty() { - for DocCandidate { - addr, - freqs, - doc_length, - .. - } in part_candidates - { - let mut score = 0.0; - for (term_index, freq) in freqs.into_iter() { - debug_assert!((term_index as usize) < idf_by_position.len()); - score += idf_by_position[term_index as usize] - * scorer.doc_weight(freq, doc_length); + #[allow(clippy::too_many_arguments)] + async fn bm25_search_modern( + &self, + tokens: Arc, + params: Arc, + operator: Operator, + mask: Arc, + metrics: Arc, + scorer: &MemBM25Scorer, + impact_scorer: Arc, + limit: usize, + ) -> Result<(Vec, Vec)> { + if self.partitions.len() > u32::MAX as usize { + return Err(Error::index(format!( + "FTS partition count {} exceeds candidate identity capacity", + self.partitions.len() + ))); + } + // Ensures old partitioned files without persisted stats populate their + // fallback stats before any synchronous local scorer is constructed. + self.aggregate_corpus_stats().await?; + + let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + let local_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + let parts = self + .partitions + .iter() + .enumerate() + .map(|(partition_ordinal, part)| { + let part = part.clone(); + let tokens = tokens.clone(); + let params = params.clone(); + let mask = mask.clone(); + let metrics = metrics.clone(); + let impact_scorer = impact_scorer.clone(); + let impact_shared_threshold = impact_shared_threshold.clone(); + let local_shared_threshold = local_shared_threshold.clone(); + async move { + let LoadedPostings { + postings, + grouped_expansions, + impact_safe, + exact_scoring_required, + } = part + .load_posting_lists( + tokens.as_ref(), + params.as_ref(), + operator, + impact_scorer.as_ref(), + metrics.as_ref(), + ) + .await?; + if postings.is_empty() { + return Result::Ok((partition_ordinal, PartitionCandidates::empty())); } - push_scored_candidate(&mut candidates, limit, addr, score)?; - } - } else { - let grouped_positions = grouped_expansions - .iter() - .map(|group| group.position) - .collect::>(); - for DocCandidate { - addr, - posting_doc_id, - freqs, - doc_length, - } in part_candidates - { - let mut score = 0.0; - for (term_index, freq) in freqs.into_iter() { - if grouped_positions.contains(&term_index) { - continue; - } - debug_assert!((term_index as usize) < idf_by_position.len()); - score += idf_by_position[term_index as usize] - * scorer.doc_weight(freq, doc_length); + let documents = part.docs.modern().cloned().ok_or_else(|| { + Error::internal("modern index contains legacy partition documents") + })?; + let visibility = Arc::new(documents.visibility(mask.as_ref()).await?); + if visibility.is_empty() { + return Result::Ok((partition_ordinal, PartitionCandidates::empty())); } - for group in &grouped_expansions { - for term in group.terms.iter() { - let Some(freq) = term.frequency(posting_doc_id) else { - continue; - }; - score += term.query_weight() * scorer.doc_weight(freq, doc_length); - } + let lengths = documents.lengths().await?; + let max_position = postings + .iter() + .map(|posting| posting.term_index() as usize) + .max() + .unwrap_or_default(); + let mut tokens_by_position = vec![String::new(); max_position + 1]; + for posting in &postings { + tokens_by_position[posting.term_index() as usize] = + posting.token().to_owned(); } - push_scored_candidate(&mut candidates, limit, addr, score)?; + let use_global_scorer = impact_safe || exact_scoring_required; + let threshold = if use_global_scorer { + impact_shared_threshold + } else { + local_shared_threshold + }; + let wand_scorer = use_global_scorer.then(|| impact_scorer.clone()); + let part_for_wand = part.clone(); + let candidates = spawn_cpu(move || { + part_for_wand.bm25_search_modern( + lengths.as_ref(), + visibility.as_ref(), + params.as_ref(), + operator, + postings, + wand_scorer, + metrics.as_ref(), + threshold, + ) + }) + .await?; + Result::Ok(( + partition_ordinal, + PartitionCandidates { + tokens_by_position, + grouped_expansions, + candidates, + }, + )) } + }) + .collect::>(); + + let mut ranked = BinaryHeap::new(); + let mut idf_cache = HashMap::new(); + let mut parts = stream::iter(parts).buffer_unordered(get_num_compute_intensive_cpus()); + while let Some((partition_ordinal, partition)) = parts.try_next().await? { + for (doc_id, score) in rescore_partition_candidates(partition, scorer, &mut idf_cache) { + let key = ((partition_ordinal as u64) << 32) | u64::from(doc_id.get()); + push_scored_key(&mut ranked, limit, key, score); + } + } + + let ranked = ranked.into_sorted_vec(); + let mut addresses = vec![0_u64; ranked.len()]; + let mut by_partition = BTreeMap::>::new(); + for (rank, Reverse(candidate)) in ranked.iter().enumerate() { + let partition_ordinal = (candidate.row_id >> 32) as usize; + let doc_id = DocId::new(candidate.row_id as u32); + by_partition + .entry(partition_ordinal) + .or_default() + .push((rank, doc_id)); + } + for (partition_ordinal, entries) in by_partition { + let documents = self.partitions[partition_ordinal] + .docs + .modern() + .expect("modern search only groups modern partitions"); + let doc_ids = entries + .iter() + .map(|(_, doc_id)| *doc_id) + .collect::>(); + let resolved = documents.resolve_addresses(&doc_ids).await?; + for ((rank, _), address) in entries.into_iter().zip(resolved) { + addresses[rank] = address; } } - - Ok(candidates - .into_sorted_vec() + let scores = ranked .into_iter() - .map(|Reverse(doc)| (doc.row_id, doc.score.0)) - .unzip()) + .map(|Reverse(candidate)| candidate.score.0) + .collect(); + Ok((addresses, scores)) } async fn load_legacy_index( @@ -1138,7 +1291,7 @@ impl InvertedIndex { store, tokens, inverted_list, - docs: Arc::new(LazyDocSet::from_loaded(docs)), + docs: PartitionDocumentStore::Legacy(Arc::new(docs)), token_set_format: TokenSetFormat::Arrow, })], corpus_stats: Arc::new(OnceCell::new()), @@ -1147,7 +1300,7 @@ impl InvertedIndex { } pub fn is_legacy(&self) -> bool { - self.partitions.len() == 1 && self.partitions[0].is_legacy() + self.partitions.len() == 1 && self.partitions[0].docs.legacy().is_some() } /// Read only the index's [`InvertedIndexParams`], @@ -1515,12 +1668,8 @@ impl InvertedIndex { elapsed_ms = partition_started.elapsed().as_millis() as u64, "fts partition posting lists prewarmed" ); - // Materialize the deferred DocSet too: prewarm's contract is - // that subsequent queries do no IO, so the per-doc row_ids / - // num_tokens must be resident, not lazily faulted in at query - // time. `ensure_loaded` opens, reads, and drops the reader. let docs_started = Instant::now(); - if let Err(err) = part.docs.ensure_loaded().await { + if let Err(err) = part.docs.prewarm().await { warn!( partition_id = part.id(), error = %err, @@ -1670,10 +1819,9 @@ pub struct InvertedPartition { store: Arc, pub(crate) tokens: TokenSet, pub(crate) inverted_list: Arc, - /// Per-doc row_id + num_tokens. Wrapped in `LazyDocSet` so partitions - /// that don't contribute hits to a query never pay the full-array - /// download. Scoring paths call `ensure_loaded` before walking wand. - pub(crate) docs: Arc, + /// Legacy documents stay in their original complete `DocSet`; modern + /// documents use typed, independently-loaded lengths and addresses. + pub(super) docs: PartitionDocumentStore, token_set_format: TokenSetFormat, } @@ -1715,30 +1863,22 @@ impl InvertedPartition { let tokens = TokenSet::load(token_file, token_set_format).await?; let invert_list_file = store.open_index_file(&posting_file_path(id)).await?; let inverted_list = PostingListReader::try_new(invert_list_file, index_cache).await?; - // Defer the per-doc row_id/num_tokens read. Construction reads only - // the doc count (one footer read) and then drops the reader; the bulk - // load happens on first scoring use, re-opening the docs file on - // demand, and partitions that never score skip it entirely. Storing - // the store + path instead of an open reader keeps a cached partition - // from pinning a docs-file handle for its whole lifetime. let docs_path = doc_file_path(id); - let num_docs = store.open_index_file(&docs_path).await?.num_rows(); - let docs = Arc::new(LazyDocSet::new( + let docs_reader = store.open_index_file(&docs_path).await?; + let docs = PartitionDocuments::try_new( store.clone(), docs_path, - num_docs, - false, + docs_reader.as_ref(), frag_reuse_index, - // V3 (256-doc block) partitions score with quantized doc lengths. inverted_list.block_size() == MAX_POSTING_BLOCK_SIZE, - )); + )?; Ok(Self { id, store, tokens, inverted_list: Arc::new(inverted_list), - docs, + docs: PartitionDocumentStore::Modern(Arc::new(docs)), token_set_format, }) } @@ -1842,7 +1982,7 @@ impl InvertedPartition { doc_ids: &[u32], frequencies: &[u32], block_size: usize, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Vec { @@ -1868,7 +2008,7 @@ impl InvertedPartition { fn union_plain_posting_lists( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Result { @@ -1904,7 +2044,7 @@ impl InvertedPartition { fn union_plain_posting_lists_with_positions( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Result { @@ -1960,7 +2100,7 @@ impl InvertedPartition { fn union_compressed_posting_lists( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Result { @@ -2019,7 +2159,7 @@ impl InvertedPartition { fn union_compressed_posting_lists_with_positions( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Result { @@ -2083,7 +2223,7 @@ impl InvertedPartition { fn union_posting_lists( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, with_positions: bool, query_weight: f32, scorer: &MemBM25Scorer, @@ -2190,6 +2330,12 @@ impl InvertedPartition { .try_collect::>() .await?; + if self.docs.modern().is_some() { + for (_, token, _, posting) in &loaded_postings { + validate_modern_posting_doc_ids(posting, token, num_docs)?; + } + } + let needs_union = loaded_postings .windows(2) .any(|window| window[0].2 == window[1].2); @@ -2239,7 +2385,12 @@ impl InvertedPartition { } let docs_for_union = if needs_union { - Some(self.docs.ensure_num_tokens_loaded().await?) + Some(match &self.docs { + PartitionDocumentStore::Legacy(docs) => LoadedDocLengths::Legacy(docs.clone()), + PartitionDocumentStore::Modern(documents) => { + LoadedDocLengths::Modern(documents.lengths().await?) + } + }) } else { None }; @@ -2278,7 +2429,7 @@ impl InvertedPartition { .into_iter() .map(|(_, _, posting)| posting) .collect::>(); - let docs = docs_for_union.as_deref().ok_or_else(|| { + let docs = docs_for_union.as_ref().ok_or_else(|| { Error::index("union docs were not loaded for grouped query terms".to_string()) })?; let posting = Self::union_posting_lists( @@ -2336,38 +2487,79 @@ impl InvertedPartition { }) } - #[instrument(level = "debug", skip_all)] - // Deferred-DocSet adds the `docs` param (caller materializes it) on top of - // the cross-partition `shared_threshold`, tipping this hot-path search fn - // one over the limit. Bundling args isn't worth the churn here. #[allow(clippy::too_many_arguments)] - pub fn bm25_search( + fn bm25_search_legacy( &self, docs: &DocSet, params: &FtsSearchParams, operator: Operator, - mask: Arc, + mask: &RowAddrMask, + postings: Vec, + impact_scorer: Option>, + metrics: &dyn MetricsCollector, + shared_threshold: Arc, + ) -> Result>> { + let documents = LegacyWandDocuments::new(docs, mask); + self.bm25_search_with_documents( + &documents, + params, + operator, + postings, + impact_scorer, + metrics, + shared_threshold, + ) + } + + #[allow(clippy::too_many_arguments)] + fn bm25_search_modern( + &self, + lengths: &DocLengths, + visibility: &DocVisibility, + params: &FtsSearchParams, + operator: Operator, postings: Vec, impact_scorer: Option>, metrics: &dyn MetricsCollector, shared_threshold: Arc, - ) -> Result> { + ) -> Result>> { + let documents = ModernWandDocuments::new(lengths, visibility); + self.bm25_search_with_documents( + &documents, + params, + operator, + postings, + impact_scorer, + metrics, + shared_threshold, + ) + } + + #[instrument(level = "debug", skip_all)] + #[allow(clippy::too_many_arguments)] + fn bm25_search_with_documents( + &self, + documents: &D, + params: &FtsSearchParams, + operator: Operator, + postings: Vec, + impact_scorer: Option>, + metrics: &dyn MetricsCollector, + shared_threshold: Arc, + ) -> Result>> { if postings.is_empty() { return Ok(Vec::new()); } - // Caller selects the DocSet shape via `LazyDocSet::docs_for_wand` - // and passes it in here; wand uses `docs.has_row_ids()` to - // handle the num_tokens-only case. let hits = if let Some(scorer) = impact_scorer { - let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) + let mut wand = Wand::new(operator, postings.into_iter(), documents, scorer) .with_shared_threshold(shared_threshold); - wand.search(params, mask, metrics)? + wand.search(params, metrics)? } else { let scorer = IndexBM25Scorer::new(std::iter::once(self)); - let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) + let mut wand = Wand::new(operator, postings.into_iter(), documents, scorer) .with_shared_threshold(shared_threshold); - wand.search(params, mask, metrics)? + wand.search(params, metrics)? }; Ok(hits) } @@ -2381,10 +2573,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.load_build_docset().await?; builder .posting_lists @@ -6562,9 +6751,7 @@ impl DocSet { } /// 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 - /// `num_tokens`) skip re-reading that column. + /// Arrow columns without re-reading either column. pub fn from_columns( row_id_col: &UInt64Array, num_tokens_col: &arrow_array::UInt32Array, @@ -9342,56 +9529,88 @@ mod tests { } #[tokio::test] - async fn test_stats_then_num_tokens_view_reuses_shared_storage() { + async fn test_persisted_stats_do_not_load_document_columns() { let (index, _counter, _tmpdir) = load_counted_v2_index(100, LanceCache::no_cache()).await; + assert!(!index.is_legacy()); let partition = index.partitions[0].clone(); + let documents = partition.docs.modern().unwrap(); assert_eq!(index.aggregate_corpus_stats().await.unwrap(), (100, 100)); - assert_eq!(partition.docs.total_tokens_cached(), Some(100)); + assert_eq!(documents.cached_stats().unwrap().total_tokens, 100); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); - let views = - futures::future::join_all((0..8).map(|_| partition.docs.ensure_num_tokens_loaded())) - .await - .into_iter() - .collect::>>() - .unwrap(); + let views = futures::future::join_all((0..8).map(|_| documents.lengths())) + .await + .into_iter() + .collect::>>() + .unwrap(); let first = &views[0]; assert!(views.iter().all(|view| Arc::ptr_eq(first, view))); - assert!(!first.has_row_ids()); - assert!(matches!(&first.num_tokens, NumTokens::Shared(_))); - assert_eq!(first.total_tokens_num(), 100); + assert_eq!(first.total_tokens(), 100); let all_rows = RowAddrMask::all_rows(); - let wand_view = partition.docs.docs_for_wand(&all_rows).await.unwrap(); - assert!(Arc::ptr_eq(first, &wand_view)); + assert!(matches!( + documents.visibility(&all_rows).await.unwrap(), + DocVisibility::All + )); + assert!(!documents.projection_loaded()); 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 visibility = documents.visibility(&filtered).await.unwrap(); + assert!(visibility.is_empty()); + assert!(documents.projection_loaded()); assert_eq!( - partition.docs.resolve_row_ids(&[0, 99]).await.unwrap(), + documents + .resolve_addresses(&[DocId::new(0), DocId::new(99)]) + .await + .unwrap(), [0, 99] ); } #[tokio::test] - async fn test_concurrent_total_and_num_tokens_view_initialization() { + async fn test_no_hit_partition_does_not_load_document_columns() { let (index, _counter, _tmpdir) = load_counted_v2_index(100, LanceCache::no_cache()).await; - let docs = index.partitions[0].docs.clone(); + let documents = index.partitions[0].docs.modern().unwrap(); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); - let totals = futures::future::join_all((0..8).map(|_| docs.total_tokens_num())); - let views = futures::future::join_all((0..8).map(|_| docs.ensure_num_tokens_loaded())); - let (totals, views) = tokio::join!(totals, views); + let tokens = Arc::new(Tokens::new(vec!["missing-token".to_owned()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(10))); + let (row_ids, scores) = index + .bm25_search( + tokens, + params, + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); - let totals = totals.into_iter().collect::>>().unwrap(); - assert_eq!(totals, vec![100; 8]); + assert!(row_ids.is_empty()); + assert!(scores.is_empty()); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); + } + + #[tokio::test] + async fn test_concurrent_stats_and_lengths_initialization() { + let (index, _counter, _tmpdir) = load_counted_v2_index(100, LanceCache::no_cache()).await; + let docs = index.partitions[0].docs.modern().unwrap().clone(); + + let stats = futures::future::join_all((0..8).map(|_| docs.stats())); + let views = futures::future::join_all((0..8).map(|_| docs.lengths())); + let (stats, views) = tokio::join!(stats, views); + + let stats = stats.into_iter().collect::>>().unwrap(); + assert!(stats.iter().all(|stats| stats.total_tokens == 100)); let views = views.into_iter().collect::>>().unwrap(); let first = &views[0]; assert!(views.iter().all(|view| Arc::ptr_eq(first, view))); - assert!(matches!(&first.num_tokens, NumTokens::Shared(_))); - assert_eq!(docs.total_tokens_cached(), Some(100)); + assert_eq!(docs.cached_stats().unwrap().total_tokens, 100); } #[tokio::test] @@ -10146,7 +10365,7 @@ mod tests { params: &FtsSearchParams, scorer: Arc, shared_threshold: Arc, - ) -> Vec { + ) -> Vec> { let LoadedPostings { postings, grouped_expansions, @@ -10166,24 +10385,24 @@ mod tests { assert!(!exact_scoring_required); assert!(grouped_expansions.is_empty()); - let mask = NoFilter.mask(); - let docs_for_wand = partition.docs.docs_for_wand(mask.as_ref()).await.unwrap(); - let mut candidates = partition - .bm25_search( - docs_for_wand.as_ref(), + let documents = partition.docs.modern().unwrap(); + let lengths = documents.lengths().await.unwrap(); + let visibility = documents + .visibility(NoFilter.mask().as_ref()) + .await + .unwrap(); + partition + .bm25_search_modern( + lengths.as_ref(), + &visibility, params, Operator::Or, - mask, postings, Some(scorer), &NoOpMetricsCollector, shared_threshold, ) - .unwrap(); - resolve_deferred_candidates(&partition.docs, &mut candidates) - .await - .unwrap(); - candidates + .unwrap() } #[tokio::test] @@ -10241,10 +10460,7 @@ mod tests { ) .await; assert_eq!(first_candidates.len(), 1); - assert!(matches!( - first_candidates[0].addr, - CandidateAddr::RowId(100) - )); + assert_eq!(first_candidates[0].document, DocId::new(0)); let first_score = scorer.query_weight("alpha") * scorer.doc_weight(1, first_candidates[0].doc_length); let published_threshold = f32::from_bits(shared_threshold.load(Ordering::Relaxed)); @@ -10262,10 +10478,7 @@ mod tests { ) .await; assert_eq!(second_candidates.len(), 1); - assert!(matches!( - second_candidates[0].addr, - CandidateAddr::RowId(200) - )); + assert_eq!(second_candidates[0].document, DocId::new(0)); let second_score = scorer.query_weight("alpha") * scorer.doc_weight(1, second_candidates[0].doc_length); assert!( diff --git a/rust/lance-index/src/scalar/inverted/lazy_docset.rs b/rust/lance-index/src/scalar/inverted/lazy_docset.rs deleted file mode 100644 index d062cca4f63..00000000000 --- a/rust/lance-index/src/scalar/inverted/lazy_docset.rs +++ /dev/null @@ -1,405 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors - -//! Deferred-load wrapper around [`DocSet`]. -//! -//! The inverted-index `DocSet` holds the per-doc `row_id` and `num_tokens` -//! arrays for a partition. Eager loading on partition open pulls roughly -//! 12 bytes × num_docs per partition; across thousands of partitions on -//! cold object storage that's tens of GiB of IO before a query has even -//! checked whether a partition contains the term it's looking for. -//! -//! [`LazyDocSet`] defers the load. Cheap sync getters (`len`, -//! `total_tokens_cached`) work without IO; async getters fetch on -//! 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`. - -use std::sync::Arc; - -use arrow::array::AsArray; -use arrow::datatypes::{UInt32Type, UInt64Type}; -use arrow_array::{Array, UInt32Array, UInt64Array}; -use lance_core::ROW_ID; -use lance_core::Result; -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::{IndexReader, IndexStore}; -use lance_select::mask::RowAddrMask; - -/// Lazy view over an inverted-index partition's `DocSet`. -/// -/// Two variants: -/// - `Loaded`: a pre-materialized DocSet (legacy paths, tests). -/// Sync accessors return cached values; async accessors return -/// the same DocSet. -/// - `Deferred`: backed by an [`IndexReader`]; columns are read and -/// cached on first request. -pub enum LazyDocSet { - Loaded(LoadedDocSet), - Deferred(Box), -} - -/// Pre-materialized DocSet view -- no reader, no IO. -pub struct LoadedDocSet { - docs: Arc, - num_rows: usize, - total_tokens: u64, -} - -/// Atomically published num-tokens state for deferred scoring. -/// -/// Keeping the Arrow column and the zero-copy `DocSet` view that carries its -/// total in one `OnceCell` prevents cancellation from exposing a partially -/// initialized scoring cache. -struct NumTokensSnapshot { - column: Arc, - docs: Arc, -} - -/// 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. -pub struct DeferredDocSet { - store: Arc, - docs_path: String, - is_legacy: bool, - frag_reuse_index: Option>, - /// V3 (256-doc block) partitions score with quantized doc lengths; the - /// flag is applied to every `DocSet` this deferred set materializes. - quantized_scoring: bool, - /// Doc count cached at construction so `len()` stays sync + IO-free. - num_rows: usize, - /// `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>, -} - -impl std::fmt::Debug for LazyDocSet { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Loaded(l) => f - .debug_struct("LazyDocSet::Loaded") - .field("num_rows", &l.num_rows) - .field("total_tokens", &l.total_tokens) - .finish(), - Self::Deferred(d) => f - .debug_struct("LazyDocSet::Deferred") - .field("num_rows", &d.num_rows) - .field( - "total_tokens_loaded", - &(d.num_tokens.initialized() || d.full.initialized()), - ) - .field("num_tokens_loaded", &d.num_tokens.initialized()) - .field("full_loaded", &d.full.initialized()) - .finish(), - } - } -} - -impl DeepSizeOf for LazyDocSet { - fn deep_size_of_children(&self, ctx: &mut lance_core::deepsize::Context) -> usize { - match self { - Self::Loaded(l) => l.docs.deep_size_of_children(ctx), - Self::Deferred(d) => { - d.full - .get() - .map(|d| d.deep_size_of_children(ctx)) - .unwrap_or(0) - + d.num_tokens - .get() - .map(|snapshot| { - let arr: &dyn Array = snapshot.column.as_ref(); - snapshot.docs.deep_size_of_children(ctx) - + 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) - } - } - } -} - -impl LazyDocSet { - #[allow(clippy::too_many_arguments)] - pub fn new( - store: Arc, - docs_path: String, - num_rows: usize, - is_legacy: bool, - frag_reuse_index: Option>, - quantized_scoring: bool, - ) -> Self { - Self::Deferred(Box::new(DeferredDocSet { - store, - docs_path, - is_legacy, - frag_reuse_index, - quantized_scoring, - num_rows, - num_tokens: OnceCell::new(), - row_ids_col: OnceCell::new(), - full: OnceCell::new(), - })) - } - - /// Wrap an already-materialized [`DocSet`]. Used by legacy paths - /// and tests that need to seed a partition without a reader. - pub fn from_loaded(docs: DocSet) -> Self { - let num_rows = docs.len(); - let total_tokens = docs.total_tokens_num(); - Self::Loaded(LoadedDocSet { - docs: Arc::new(docs), - num_rows, - total_tokens, - }) - } - - pub fn len(&self) -> usize { - match self { - Self::Loaded(l) => l.num_rows, - Self::Deferred(d) => d.num_rows, - } - } - - /// Sync read of cached `total_tokens`. Returns `None` for a - /// `Deferred` LazyDocSet that hasn't yet had any of - /// `total_tokens_num` / `ensure_num_tokens_loaded` / `ensure_loaded` - /// run. Used by sync scoring code that has already paid for one - /// of those async calls. - pub fn total_tokens_cached(&self) -> Option { - match self { - Self::Loaded(l) => Some(l.total_tokens), - Self::Deferred(d) => d - .full - .get() - .map(|docs| docs.total_tokens_num()) - .or_else(|| { - d.num_tokens - .get() - .map(|snapshot| snapshot.docs.total_tokens_num()) - }), - } - } - - /// True if this DocSet carries a FragReuseIndex. Callers MUST - /// avoid the deferred-row_id path when this is set: targeted - /// row_id reads return raw stored ids, bypassing the per-id - /// `remap_row_id` filter that `DocSet::from_columns` applies. - pub fn has_frag_reuse_remap(&self) -> bool { - match self { - Self::Loaded(_) => false, - Self::Deferred(d) => d.frag_reuse_index.is_some(), - } - } - - /// Sum of `num_tokens` across all docs. - pub async fn total_tokens_num(&self) -> Result { - match self { - Self::Loaded(l) => Ok(l.total_tokens), - Self::Deferred(d) => d.total_tokens_num().await, - } - } - - /// Materialize the full DocSet, including row_ids. - pub async fn ensure_loaded(&self) -> Result> { - match self { - Self::Loaded(l) => Ok(l.docs.clone()), - Self::Deferred(d) => d.ensure_loaded().await, - } - } - - /// 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 - /// [`Self::resolve_row_ids`]. The tokens-only result is cached separately; - /// a later `ensure_loaded` must still produce a full DocSet. - pub async fn ensure_num_tokens_loaded(&self) -> Result> { - match self { - Self::Loaded(l) => Ok(l.docs.clone()), - Self::Deferred(d) => d.ensure_num_tokens_loaded().await, - } - } - - /// 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 - } - } - - /// Resolve a batch of `doc_id`s to their `row_id`s. Used by the - /// deferred-row_id scoring path to map post-wand top-K candidates - /// without going through a full DocSet build. - /// - /// Not safe with a FragReuseIndex (see - /// [`Self::has_frag_reuse_remap`]): the targeted reads return - /// raw stored ids without applying the remap/skip. - pub async fn resolve_row_ids(&self, doc_ids: &[u32]) -> Result> { - match self { - Self::Loaded(l) => Ok(doc_ids.iter().map(|&d| l.docs.row_id(d)).collect()), - Self::Deferred(d) => d.resolve_row_ids(doc_ids).await, - } - } -} - -impl DeferredDocSet { - /// Open a fresh docs-file reader. Dropped by the caller once its read - /// completes, so no handle is pinned across the partition's lifetime. - async fn reader(&self) -> Result> { - self.store.open_index_file(&self.docs_path).await - } - - async fn total_tokens_num(&self) -> Result { - if let Some(full) = self.full.get() { - return Ok(full.total_tokens_num()); - } - Ok(self.num_tokens_snapshot().await?.docs.total_tokens_num()) - } - - async fn num_tokens_snapshot(&self) -> Result<&NumTokensSnapshot> { - self.num_tokens - .get_or_try_init(|| async { - let reader = self.reader().await?; - let batch = reader - .read_range(0..self.num_rows, Some(&[NUM_TOKEN_COL])) - .await?; - let column = Arc::new(batch[NUM_TOKEN_COL].as_primitive::().clone()); - let total_tokens = column.values().iter().map(|&n| n as u64).sum(); - let mut docs = DocSet::from_cached_num_tokens(column.as_ref(), total_tokens); - docs.set_quantized_scoring(self.quantized_scoring); - Result::Ok(NumTokensSnapshot { - column, - docs: Arc::new(docs), - }) - }) - .await - } - - 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() - } - - async fn ensure_loaded(&self) -> Result> { - let docs = self - .full - .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 { - DocSet::load( - self.reader().await?, - self.is_legacy, - self.frag_reuse_index.clone(), - ) - .await? - }; - docs.set_quantized_scoring(self.quantized_scoring); - Result::Ok(Arc::new(docs)) - }) - .await? - .clone(); - Ok(docs) - } - - async fn ensure_num_tokens_loaded(&self) -> Result> { - if let Some(full) = self.full.get() { - return Ok(full.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() - { - 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 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()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::scalar::lance_format::LanceIndexStore; - use lance_core::cache::LanceCache; - use lance_core::utils::tempfile::TempObjDir; - use lance_io::object_store::ObjectStore; - - #[tokio::test] - async fn test_full_docset_is_a_complete_cached_snapshot() { - let temp_dir = TempObjDir::default(); - let store = Arc::new(LanceIndexStore::new( - ObjectStore::local().into(), - temp_dir.clone(), - Arc::new(LanceCache::no_cache()), - )); - let docs = LazyDocSet::new(store, "unused".to_owned(), 3, false, None, false); - assert_eq!(docs.total_tokens_cached(), None); - - let row_ids = UInt64Array::from(vec![10, 20, 30]); - let num_tokens = UInt32Array::from(vec![3, 5, 8]); - let full = Arc::new(DocSet::from_columns(&row_ids, &num_tokens, false, None).unwrap()); - let LazyDocSet::Deferred(deferred) = &docs else { - panic!("expected a deferred DocSet"); - }; - deferred.full.set(full.clone()).unwrap(); - - let wand_docs = docs.ensure_num_tokens_loaded().await.unwrap(); - assert!(Arc::ptr_eq(&wand_docs, &full)); - assert_eq!(wand_docs.total_tokens_num(), 16); - assert_eq!(docs.total_tokens_cached(), Some(16)); - } -} diff --git a/rust/lance-index/src/scalar/inverted/scorer.rs b/rust/lance-index/src/scalar/inverted/scorer.rs index 3a33a67ff7a..37b3dc5381c 100644 --- a/rust/lance-index/src/scalar/inverted/scorer.rs +++ b/rust/lance-index/src/scalar/inverted/scorer.rs @@ -164,24 +164,20 @@ pub struct IndexBM25Scorer<'a> { } impl<'a> IndexBM25Scorer<'a> { - /// Sync constructor. Reads each partition's cached `total_tokens` via - /// `LazyDocSet::total_tokens_cached()`; callers must have already - /// populated it (via `ensure_loaded`, `ensure_num_tokens_loaded`, or - /// `total_tokens_num`). Panics with a clear message otherwise — this - /// is the wand-scoring path where the contract is statically known. + /// Sync constructor. Query setup populates immutable partition stats + /// before entering the CPU-only WAND executor. pub fn new(partitions: impl Iterator) -> Self { let partitions = partitions.collect::>(); - let num_docs = partitions.iter().map(|p| p.docs.len()).sum(); - let total_tokens: u64 = partitions + let stats = partitions .iter() - .map(|p| { - p.docs.total_tokens_cached().expect( - "IndexBM25Scorer::new requires each partition's total_tokens to be \ - cached; call `ensure_loaded` / `ensure_num_tokens_loaded` / \ - `total_tokens_num` first", + .map(|partition| { + partition.docs.cached_stats().expect( + "IndexBM25Scorer::new requires partition stats to be loaded before WAND", ) }) - .sum(); + .collect::>(); + let num_docs = stats.iter().map(|stats| stats.num_docs).sum(); + let total_tokens: u64 = stats.iter().map(|stats| stats.total_tokens).sum(); let avgdl = total_tokens as f32 / num_docs as f32; Self { partitions, diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index e2eb352175b..18fefecc2e6 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::{ @@ -22,6 +21,7 @@ use crate::metrics::MetricsCollector; use super::{ CompressedPositionStorage, + documents::{DocId, DocLengths, DocVisibility}, impact::{IMPACT_LEVEL1_BLOCKS, ImpactScoreCache, ImpactSkipData}, index::{PositionStreamCodec, dequantize_doc_length}, query::Operator, @@ -168,10 +168,10 @@ impl TopKCollector { } } - fn into_candidates( + fn into_candidates( self, - mut to_addr: impl FnMut(u64) -> CandidateAddr, - ) -> Result> { + mut to_candidate: impl FnMut(u64) -> C, + ) -> Result>> { let Self { heap, mut frequency_slots, @@ -181,7 +181,7 @@ impl TopKCollector { .map( |Reverse((doc, doc_length, posting_doc_id, frequency_slot))| { Ok(DocCandidate { - addr: to_addr(doc.row_id), + document: to_candidate(doc.row_id), posting_doc_id, freqs: frequency_slots.take(frequency_slot)?, doc_length, @@ -1318,12 +1318,12 @@ impl PostingIterator { // The norm cache arg tips this hot-path fn over the limit; bundling the // scoring inputs isn't worth the churn here. #[allow(clippy::too_many_arguments)] - fn collect_window_scores( + fn collect_window_scores( &mut self, window_min: u64, up_to: u64, clause_idx: usize, - docs: &DocSet, + docs: &D, scorer: &S, norm_k: Option<(&[u8], &[f32; 256])>, acc: &mut WindowAccumulator, @@ -1392,10 +1392,7 @@ impl PostingIterator { if doc_id > up_to { break; } - let doc_length = match &doc { - DocInfo::Raw(raw) => docs.scoring_num_tokens(raw.doc_id), - DocInfo::Located(located) => docs.num_tokens_by_row_id(located.row_id), - }; + let doc_length = docs.doc_length(&doc); let score = self.score(scorer, doc.frequency(), doc_length); let slot = (doc_id - window_min) as usize; acc.add(clause_idx, slot, score, doc.frequency()); @@ -1417,6 +1414,55 @@ impl PostingIterator { PostingList::Plain(ref plain) => plain.row_ids.get(self.index + 1).cloned(), } } + + #[cfg(test)] + fn validate_modern_doc_ids(&self, num_docs: usize) -> Result<()> { + validate_modern_posting_doc_ids(&self.list, &self.token, num_docs) + } +} + +/// Validate the dense-DocId boundary before any document-length lookup. +/// Modern postings are sorted, so decoding only the final physical block is +/// sufficient to validate their maximum DocId. +pub(super) fn validate_modern_posting_doc_ids( + posting: &PostingList, + token: &str, + num_docs: usize, +) -> Result<()> { + let PostingList::Compressed(list) = posting else { + return Err(Error::index(format!( + "modern FTS posting for token {token:?} uses a legacy row-address layout" + ))); + }; + if list.length == 0 { + return Ok(()); + } + let last_block_idx = list.blocks.len().checked_sub(1).ok_or_else(|| { + Error::index(format!( + "modern FTS posting for token {token:?} has length {} but no blocks", + list.length + )) + })?; + let mut state = CompressedState::new(list.block_size); + state.decompress( + list.blocks.value(last_block_idx), + last_block_idx, + list.blocks.len(), + list.length, + list.posting_tail_codec, + list.block_size, + ); + let max_doc_id = state.doc_ids.last().copied().ok_or_else(|| { + Error::index(format!( + "modern FTS posting for token {token:?} has an empty final block" + )) + })?; + if max_doc_id as usize >= num_docs { + return Err(Error::index(format!( + "modern FTS posting for token {token:?} contains DocId {max_doc_id}, outside [0, {num_docs})" + ))); + } + Ok(()) } /// Inner window span (in doc ids) of the bulk MAXSCORE path. Same as Lucene's @@ -1486,19 +1532,9 @@ impl WindowAccumulator { } } -/// How wand identified a candidate: either it already had the real -/// row_id (DocSet carried row_ids), or only the partition-local -/// doc_id (deferred-row_id path; the caller must resolve via -/// [`super::lazy_docset::LazyDocSet::resolve_row_ids`]). -#[derive(Debug, Clone, Copy)] -pub enum CandidateAddr { - RowId(u64), - Pending(u32), -} - #[derive(Debug)] -pub struct DocCandidate { - pub addr: CandidateAddr, +pub struct DocCandidate { + pub document: C, /// The document key used by the posting lists: doc_id for compressed /// postings, row_id for legacy plain postings. pub posting_doc_id: u64, @@ -1507,6 +1543,226 @@ pub struct DocCandidate { pub doc_length: u32, } +/// Document-side contract consumed by WAND. Implementations fix candidate +/// identity and visibility before the CPU executor starts. +type FlatDocuments<'a> = (usize, Box + 'a>); + +pub(super) trait WandDocuments { + type Candidate: Copy + Debug; + + fn len(&self) -> usize; + fn scoring_norms(&self) -> Option<&[u8]>; + fn scoring_num_tokens(&self, doc_id: u32) -> u32; + fn doc_length(&self, doc: &DocInfo) -> u32; + fn document_key(&self, doc: &DocInfo) -> Option; + fn document_key_for_doc_id(&self, doc_id: u32) -> Option; + fn candidate_from_key(&self, key: u64) -> Self::Candidate; + fn flat_documents(&self) -> Option>; + fn flat_doc_length(&self, doc_id: u64, document_key: u64, compressed: bool) -> u32; +} + +pub(super) struct ModernWandDocuments<'a> { + lengths: &'a DocLengths, + visibility: &'a DocVisibility, +} + +impl<'a> ModernWandDocuments<'a> { + pub(crate) fn new(lengths: &'a DocLengths, visibility: &'a DocVisibility) -> Self { + Self { + lengths, + visibility, + } + } +} + +impl WandDocuments for ModernWandDocuments<'_> { + type Candidate = DocId; + + fn len(&self) -> usize { + self.lengths.len() + } + + fn scoring_norms(&self) -> Option<&[u8]> { + self.lengths.scoring_norms() + } + + fn scoring_num_tokens(&self, doc_id: u32) -> u32 { + self.lengths.scoring(DocId::new(doc_id)) + } + + fn doc_length(&self, doc: &DocInfo) -> u32 { + match doc { + DocInfo::Raw(doc) => self.scoring_num_tokens(doc.doc_id), + DocInfo::Located(_) => unreachable!("modern posting lists contain dense DocIds"), + } + } + + fn document_key(&self, doc: &DocInfo) -> Option { + match doc { + DocInfo::Raw(doc) if self.visibility.selected(DocId::new(doc.doc_id)) => { + Some(u64::from(doc.doc_id)) + } + DocInfo::Raw(_) => None, + DocInfo::Located(_) => unreachable!("modern posting lists contain dense DocIds"), + } + } + + fn document_key_for_doc_id(&self, doc_id: u32) -> Option { + self.visibility + .selected(DocId::new(doc_id)) + .then_some(u64::from(doc_id)) + } + + fn candidate_from_key(&self, key: u64) -> Self::Candidate { + DocId::new(key as u32) + } + + fn flat_documents(&self) -> Option<(usize, Box + '_>)> { + self.visibility.iter().map(|doc_ids| { + let len = self.visibility.len(self.lengths.len()); + let docs = doc_ids.map(|doc_id| { + let value = u64::from(doc_id.get()); + (value, value) + }); + (len, Box::new(docs) as Box>) + }) + } + + fn flat_doc_length(&self, doc_id: u64, _document_key: u64, _compressed: bool) -> u32 { + self.scoring_num_tokens(doc_id as u32) + } +} + +pub(super) struct LegacyWandDocuments<'a> { + docs: &'a DocSet, + mask: &'a RowAddrMask, +} + +impl<'a> LegacyWandDocuments<'a> { + pub(crate) fn new(docs: &'a DocSet, mask: &'a RowAddrMask) -> Self { + Self { docs, mask } + } +} + +impl WandDocuments for LegacyWandDocuments<'_> { + type Candidate = u64; + + fn len(&self) -> usize { + self.docs.len() + } + + fn scoring_norms(&self) -> Option<&[u8]> { + self.docs.scoring_norms() + } + + fn scoring_num_tokens(&self, doc_id: u32) -> u32 { + self.docs.scoring_num_tokens(doc_id) + } + + fn doc_length(&self, doc: &DocInfo) -> u32 { + match doc { + DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), + DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), + } + } + + fn document_key(&self, doc: &DocInfo) -> Option { + let row_id = match doc { + DocInfo::Raw(doc) => self.docs.row_id(doc.doc_id), + DocInfo::Located(doc) => doc.row_id, + }; + (row_id != RowAddress::TOMBSTONE_ROW && self.mask.selected(row_id)).then_some(row_id) + } + + fn document_key_for_doc_id(&self, doc_id: u32) -> Option { + let row_id = self.docs.row_id(doc_id); + (row_id != RowAddress::TOMBSTONE_ROW && self.mask.selected(row_id)).then_some(row_id) + } + + fn candidate_from_key(&self, key: u64) -> Self::Candidate { + key + } + + fn flat_documents(&self) -> Option<(usize, Box + '_>)> { + let count = self.mask.max_len()? as usize; + let row_ids = self.mask.iter_addrs()?; + let docs = row_ids.flat_map(|row_addr| { + let row_id: u64 = row_addr.into(); + self.docs + .doc_ids(row_id) + .map(move |doc_id| (doc_id, row_id)) + }); + Some((count, Box::new(docs))) + } + + fn flat_doc_length(&self, doc_id: u64, document_key: u64, compressed: bool) -> u32 { + if compressed { + self.docs.scoring_num_tokens(doc_id as u32) + } else { + self.docs.num_tokens_by_row_id(document_key) + } + } +} + +// Most unit tests exercise WAND in isolation with an in-memory complete +// DocSet. Production code must select one of the explicit modern/legacy +// adapters above, so this convenience implementation is test-only. +#[cfg(test)] +impl WandDocuments for DocSet { + type Candidate = u64; + + fn len(&self) -> usize { + self.len() + } + + fn scoring_norms(&self) -> Option<&[u8]> { + self.scoring_norms() + } + + fn scoring_num_tokens(&self, doc_id: u32) -> u32 { + self.scoring_num_tokens(doc_id) + } + + fn doc_length(&self, doc: &DocInfo) -> u32 { + match doc { + DocInfo::Raw(doc) => self.scoring_num_tokens(doc.doc_id), + DocInfo::Located(doc) => self.num_tokens_by_row_id(doc.row_id), + } + } + + fn document_key(&self, doc: &DocInfo) -> Option { + Some(match doc { + DocInfo::Raw(doc) if self.has_row_ids() => self.row_id(doc.doc_id), + DocInfo::Raw(doc) => u64::from(doc.doc_id), + DocInfo::Located(doc) => doc.row_id, + }) + } + + fn document_key_for_doc_id(&self, doc_id: u32) -> Option { + Some(if self.has_row_ids() { + self.row_id(doc_id) + } else { + u64::from(doc_id) + }) + } + + fn candidate_from_key(&self, key: u64) -> Self::Candidate { + key + } + + fn flat_documents(&self) -> Option<(usize, Box + '_>)> { + None + } + + fn flat_doc_length(&self, doc_id: u64, document_key: u64, compressed: bool) -> u32 { + if compressed { + self.scoring_num_tokens(doc_id as u32) + } else { + self.num_tokens_by_row_id(document_key) + } + } +} + struct HeadPosting { // Iterators that are already positioned on or after the next candidate doc. // The heap is ordered by smallest doc id so the top element determines @@ -1626,7 +1882,7 @@ impl Ord for TailPosting { } } -pub struct Wand<'a, S: Scorer> { +pub struct Wand<'a, S: Scorer, D: WandDocuments> { threshold: f32, // multiple of factor and the minimum score of the top-k documents operator: Operator, num_terms: usize, @@ -1663,7 +1919,7 @@ pub struct Wand<'a, S: Scorer> { bulk_and_mode_override: Option, #[cfg(test)] bulk_and_searches: usize, - docs: &'a DocSet, + documents: &'a D, scorer: S, // Shared cross-partition top-k floor. Each partition publishes its local // k-th score (`atomic_store_max_f32`) and prunes against the running value @@ -1686,11 +1942,11 @@ fn atomic_store_max_f32(slot: &AtomicU32, val: f32) { // we were using row id as doc id in the past, which is u64, // but now we are using the index as doc id, which is u32. // so here WAND is a generic struct that can be used for both u32 and u64 doc ids. -impl<'a, S: Scorer> Wand<'a, S> { +impl<'a, S: Scorer, D: WandDocuments> Wand<'a, S, D> { pub(crate) fn new( operator: Operator, postings: impl Iterator, - docs: &'a DocSet, + documents: &'a D, scorer: S, ) -> Self { let mut head = BinaryHeap::new(); @@ -1730,7 +1986,7 @@ impl<'a, S: Scorer> Wand<'a, S> { bulk_and_mode_override: None, #[cfg(test)] bulk_and_searches: 0, - docs, + documents, scorer, shared_threshold: None, } @@ -1757,8 +2013,7 @@ impl<'a, S: Scorer> Wand<'a, S> { /// the cache is bit-identical to `scorer.doc_weight`, because both /// evaluate the same expressions on the same quantized lengths. fn norm_k_cache(&self) -> Option<(&'a [u8], Box<[f32; 256]>)> { - let docs: &'a DocSet = self.docs; - let norms = docs.scoring_norms()?; + let norms = self.documents.scoring_norms()?; let mut cache = Box::new([0f32; 256]); for (code, slot) in cache.iter_mut().enumerate() { *slot = self.scorer.doc_norm(dequantize_doc_length(code as u8))?; @@ -1796,23 +2051,19 @@ impl<'a, S: Scorer> Wand<'a, S> { pub(crate) fn search( &mut self, params: &FtsSearchParams, - mask: Arc, metrics: &dyn MetricsCollector, - ) -> Result> { + ) -> Result>> { let limit = params.limit.unwrap_or(usize::MAX); if limit == 0 { 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 self.operator == Operator::Or + && let Some((num_docs_selected, documents)) = self.documents.flat_documents() + && num_docs_selected.saturating_mul(100) + <= (*FLAT_SEARCH_PERCENT_THRESHOLD as usize).saturating_mul(self.documents.len()) + { + return self.flat_search(params, documents, metrics); } // Top-k disjunctions over compressed lists can opt into the bulk @@ -1827,7 +2078,7 @@ impl<'a, S: Scorer> Wand<'a, S> { posting.posting.is_compressed() && !posting.posting.has_grouped_terms() }) { - return self.maxscore_search(params, mask, metrics); + return self.maxscore_search(params, metrics); } // Top-k conjunctions (AND and phrase) over compressed lists use the @@ -1849,15 +2100,9 @@ impl<'a, S: Scorer> Wand<'a, S> { { self.bulk_and_searches += 1; } - return self.and_bulk_search(params, mask, metrics); + return self.and_bulk_search(params, metrics); } - // Deferred-row_id path: when the DocSet was built without - // row_ids, wand emits candidates carrying just the - // partition-local doc_id; the outer caller resolves them to - // row_ids post-wand. - let docs_has_row_ids = self.docs.has_row_ids(); - let mut candidates = TopKCollector::new(limit, std::cmp::min(limit, BLOCK_SIZE * 10)); let mut num_comparisons = 0; let mut and_search_stats = (self.operator == Operator::And).then_some(AndSearchStats { @@ -1874,41 +2119,16 @@ impl<'a, S: Scorer> Wand<'a, S> { and_stats.candidates_seen += 1; } - // Either a real row_id (so we can run the mask check - // inline) or the doc_id widened to u64 (deferred path; - // the outer caller will resolve it post-wand). let posting_doc_id = doc.doc_id(); - let row_id = match &doc { - DocInfo::Raw(doc) => { - if docs_has_row_ids { - self.docs.row_id(doc.doc_id) - } else { - doc.doc_id as u64 - } - } - DocInfo::Located(doc) => doc.row_id, - }; - // Skip docs the fragment-reuse remap deleted. They are tombstoned - // in the DocSet (slot kept so posting-list doc_ids stay aligned) - // and must not surface in results. - if docs_has_row_ids && row_id == RowAddress::TOMBSTONE_ROW { - if self.operator == Operator::Or { - self.push_back_leads(doc.doc_id() + 1); - } - continue; - } - if docs_has_row_ids && !mask.selected(row_id) { + let Some(document_key) = self.documents.document_key(&doc) else { if self.operator == Operator::Or { self.push_back_leads(doc.doc_id() + 1); } continue; - } - - let doc_length = match &doc { - DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), - DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), }; + let doc_length = self.documents.doc_length(&doc); + let score = if self.operator == Operator::Or { self.advance_all_tail(doc.doc_id(), Some(doc_length), Some(&mut score)); if params.phrase_slop.is_some() @@ -1932,7 +2152,7 @@ impl<'a, S: Scorer> Wand<'a, S> { }; if candidates.insert( - ScoredDoc::new(row_id, score), + ScoredDoc::new(document_key, score), doc_length, posting_doc_id, self.iter_term_freqs(), @@ -1969,45 +2189,23 @@ impl<'a, S: Scorer> Wand<'a, S> { metrics.record_freqs_collected(and_stats.freqs_collected); } - // The heap entry's `row_id` slot is either a real row_id - // (DocSet had row_ids) or the doc_id widened to u64 - // (deferred). Tag it accordingly so the caller can match - // rather than guess. - let to_addr = |row_id_slot: u64| { - if docs_has_row_ids { - CandidateAddr::RowId(row_id_slot) - } else { - CandidateAddr::Pending(row_id_slot as u32) - } - }; - candidates.into_candidates(to_addr) + candidates.into_candidates(|key| self.documents.candidate_from_key(key)) } fn flat_search( &mut self, params: &FtsSearchParams, - row_ids: Box + '_>, + documents: Box + '_>, metrics: &dyn MetricsCollector, - ) -> Result> { + ) -> Result>> { let limit = params.limit.unwrap_or(usize::MAX); if limit == 0 { return Ok(vec![]); } - // 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. - // A list column maps one row id to several doc ids, so expand every - // document the row owns — keying on a single doc id would drop matches - // at non-last list positions (lancedb#3352). - let doc_ids = row_ids - .flat_map(|row_addr| { - let row_id: u64 = row_addr.into(); - self.docs - .doc_ids(row_id) - .map(move |doc_id| (doc_id, row_id)) - }) - .sorted_unstable() - .collect::>(); + // Posting iterators are forward-only, so selected DocIds are sorted + // before driving the sparse executor. + let documents = documents.sorted_unstable().collect::>(); let is_compressed = self .head .peek() @@ -2021,7 +2219,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let mut num_comparisons = 0; let mut candidates = TopKCollector::new(limit, 0); - for (doc_id, row_id) in doc_ids { + for (doc_id, document_key) in documents { num_comparisons += 1; self.move_head_before_target_to_tail(doc_id); self.move_head_doc_to_lead(doc_id); @@ -2051,10 +2249,9 @@ impl<'a, S: Scorer> Wand<'a, S> { } // score the doc - let doc_length = match is_compressed { - true => self.docs.scoring_num_tokens(doc_id as u32), - false => self.docs.num_tokens_by_row_id(row_id), - }; + let doc_length = self + .documents + .flat_doc_length(doc_id, document_key, is_compressed); if self.operator == Operator::Or && !self.refine_or_candidate(doc_id, doc_length) { // `flat_search` evaluates an explicit allow-list of doc ids. Unlike the // regular WAND path, skipping to the next block boundary is unsafe here @@ -2068,7 +2265,7 @@ impl<'a, S: Scorer> Wand<'a, S> { let score = self.score(doc_length); if candidates.insert( - ScoredDoc::new(row_id, score), + ScoredDoc::new(document_key, score), doc_length, doc_id, self.iter_term_freqs(), @@ -2081,9 +2278,7 @@ impl<'a, S: Scorer> Wand<'a, S> { } metrics.record_comparisons(num_comparisons); - // flat_search is driven by an explicit row_ids iterator, so - // every candidate already has a real row_id. - candidates.into_candidates(CandidateAddr::RowId) + candidates.into_candidates(|key| self.documents.candidate_from_key(key)) } /// Bulk MAXSCORE top-k disjunction, mirroring Lucene's MaxScoreBulkScorer. @@ -2098,9 +2293,8 @@ impl<'a, S: Scorer> Wand<'a, S> { fn maxscore_search( &mut self, params: &FtsSearchParams, - mask: Arc, metrics: &dyn MetricsCollector, - ) -> Result> { + ) -> Result>> { struct MaxScoreClause { posting: Box, bound: f32, @@ -2108,7 +2302,6 @@ impl<'a, S: Scorer> Wand<'a, S> { } let limit = params.limit.unwrap_or(usize::MAX); - let docs_has_row_ids = self.docs.has_row_ids(); let mut clauses = std::mem::take(&mut self.head) .into_vec() .into_iter() @@ -2284,9 +2477,10 @@ impl<'a, S: Scorer> Wand<'a, S> { } None => { essential_weight - * self - .scorer - .doc_weight(freq, self.docs.scoring_num_tokens(doc as u32)) + * self.scorer.doc_weight( + freq, + self.documents.scoring_num_tokens(doc as u32), + ) } }; if !(self.threshold > 0.0 @@ -2297,14 +2491,9 @@ impl<'a, S: Scorer> Wand<'a, S> { total_sum_upper_bound_factor, )) { - let row_id = if docs_has_row_ids { - self.docs.row_id(doc as u32) - } else { - doc - }; - let masked_out = docs_has_row_ids - && (row_id == RowAddress::TOMBSTONE_ROW || !mask.selected(row_id)); - if !masked_out { + if let Some(document_key) = + self.documents.document_key_for_doc_id(doc as u32) + { let mut total = score; let mut rejected = false; for i in (0..non_essential.len()).rev() { @@ -2337,7 +2526,7 @@ impl<'a, S: Scorer> Wand<'a, S> { None => probe.score( &self.scorer, d.frequency(), - self.docs.scoring_num_tokens(doc as u32), + self.documents.scoring_num_tokens(doc as u32), ), }; } @@ -2348,9 +2537,9 @@ impl<'a, S: Scorer> Wand<'a, S> { // which drops zero-score matches (e.g. terms // with idf 0) exactly like Wand::next does. if !rejected && total > self.threshold { - let doc_length = self.docs.scoring_num_tokens(doc as u32); + let doc_length = self.documents.scoring_num_tokens(doc as u32); if candidates.insert( - ScoredDoc::new(row_id, total), + ScoredDoc::new(document_key, total), doc_length, doc, std::iter::once((essential_term, freq)).chain( @@ -2456,7 +2645,7 @@ impl<'a, S: Scorer> Wand<'a, S> { inner_min, inner_max, clause_idx, - self.docs, + self.documents, &self.scorer, norm_k_ref, &mut acc, @@ -2491,17 +2680,11 @@ impl<'a, S: Scorer> Wand<'a, S> { continue; } - let row_id = if docs_has_row_ids { - self.docs.row_id(doc as u32) - } else { - doc - }; - if docs_has_row_ids - && (row_id == RowAddress::TOMBSTONE_ROW || !mask.selected(row_id)) - { + let Some(document_key) = self.documents.document_key_for_doc_id(doc as u32) + else { acc.clear_slot(slot); continue; - } + }; // Doc length is only needed at heap-insert time; the // non-essential completion scores go through the norm @@ -2537,7 +2720,7 @@ impl<'a, S: Scorer> Wand<'a, S> { None => { let doc_length = *doc_length_cell.get_or_insert_with(|| { - self.docs.scoring_num_tokens(doc as u32) + self.documents.scoring_num_tokens(doc as u32) }); posting.score(&self.scorer, d.frequency(), doc_length) } @@ -2547,9 +2730,9 @@ impl<'a, S: Scorer> Wand<'a, S> { if !rejected && score > self.threshold { let doc_length = doc_length_cell - .unwrap_or_else(|| self.docs.scoring_num_tokens(doc as u32)); + .unwrap_or_else(|| self.documents.scoring_num_tokens(doc as u32)); if candidates.insert( - ScoredDoc::new(row_id, score), + ScoredDoc::new(document_key, score), doc_length, doc, clauses.iter().enumerate().filter_map(|(i, clause)| { @@ -2586,14 +2769,7 @@ impl<'a, S: Scorer> Wand<'a, S> { metrics.record_comparisons(num_comparisons); - let to_addr = |row_id_slot: u64| { - if docs_has_row_ids { - CandidateAddr::RowId(row_id_slot) - } else { - CandidateAddr::Pending(row_id_slot as u32) - } - }; - candidates.into_candidates(to_addr) + candidates.into_candidates(|key| self.documents.candidate_from_key(key)) } // calculate the score of the current document @@ -2692,10 +2868,7 @@ impl<'a, S: Scorer> Wand<'a, S> { self.push_back_leads(target + 1); continue; }; - let doc_length = match &first_doc { - DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), - DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), - }; + let doc_length = self.documents.doc_length(&first_doc); let mut lead_score = 0.0; if let Some(first_posting) = self.lead.first() { lead_score += first_posting.score(&self.scorer, first_doc.frequency(), doc_length); @@ -2774,10 +2947,7 @@ impl<'a, S: Scorer> Wand<'a, S> { } let lead_doc = self.lead.first().and_then(|posting| posting.doc())?; - let doc_length = match &lead_doc { - DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), - DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), - }; + let doc_length = self.documents.doc_length(&lead_doc); if self.and_candidate_cannot_beat_threshold(doc_length) { self.and_candidates_pruned_before_return += 1; let next_target = self.and_advance_target(doc.saturating_add(1)); @@ -2811,14 +2981,12 @@ impl<'a, S: Scorer> Wand<'a, S> { fn and_bulk_search( &mut self, params: &FtsSearchParams, - mask: Arc, metrics: &dyn MetricsCollector, - ) -> Result> { + ) -> Result>> { let limit = params.limit.unwrap_or(usize::MAX); if limit == 0 { return Ok(vec![]); } - let docs_has_row_ids = self.docs.has_row_ids(); let num_lists = self.lead.len(); let phrase_slop = params.phrase_slop; @@ -3228,7 +3396,7 @@ impl<'a, S: Scorer> Wand<'a, S> { batch_norms.push(norms[doc as usize]); } } - None => match self.docs.scoring_norms() { + None => match self.documents.scoring_norms() { Some(norms) => { for &doc in batch_docs.iter() { batch_lens.push(dequantize_doc_length(norms[doc as usize])); @@ -3236,7 +3404,7 @@ impl<'a, S: Scorer> Wand<'a, S> { } None => { for &doc in batch_docs.iter() { - batch_lens.push(self.docs.scoring_num_tokens(doc)); + batch_lens.push(self.documents.scoring_num_tokens(doc)); } } }, @@ -3271,16 +3439,9 @@ impl<'a, S: Scorer> Wand<'a, S> { self.and_window_stats.candidates_returned += 1; num_comparisons += 1; - let row_id = if docs_has_row_ids { - self.docs.row_id(doc) - } else { - u64::from(doc) - }; - if docs_has_row_ids - && (row_id == RowAddress::TOMBSTONE_ROW || !mask.selected(row_id)) - { + let Some(document_key) = self.documents.document_key_for_doc_id(doc) else { continue; - } + }; if let Some(slop) = phrase_slop { // Park every clause's iterator on this doc so @@ -3320,7 +3481,7 @@ impl<'a, S: Scorer> Wand<'a, S> { } if candidates.insert( - ScoredDoc::new(row_id, score), + ScoredDoc::new(document_key, score), doc_length, u64::from(doc), wins.iter().zip(self.lead.iter()).zip(offs.iter()).map( @@ -3362,14 +3523,7 @@ impl<'a, S: Scorer> Wand<'a, S> { metrics.record_and_full_scores(stats.full_scores); metrics.record_freqs_collected(stats.freqs_collected); - let to_addr = |row_id_slot: u64| { - if docs_has_row_ids { - CandidateAddr::RowId(row_id_slot) - } else { - CandidateAddr::Pending(row_id_slot as u32) - } - }; - candidates.into_candidates(to_addr) + candidates.into_candidates(|key| self.documents.candidate_from_key(key)) } fn and_move_to_next_block(&mut self, target: u64) { @@ -4263,15 +4417,12 @@ mod tests { } assert_eq!(collector.num_frequency_slots(), LIMIT); - let mut candidates = collector.into_candidates(CandidateAddr::RowId)?; + let mut candidates = collector.into_candidates(|key| key)?; candidates.sort_unstable_by_key(|candidate| candidate.posting_doc_id); assert_eq!(candidates.len(), LIMIT); for (candidate, expected_doc) in candidates.iter().zip(NUM_DOCS - LIMIT..NUM_DOCS) { assert_eq!(candidate.posting_doc_id, expected_doc as u64); - assert!(matches!( - candidate.addr, - CandidateAddr::RowId(row_id) if row_id == expected_doc as u64 - )); + assert_eq!(candidate.document, expected_doc as u64); let expected_freqs = (0..expected_doc % 4 + 1) .map(|term| (term as u32, expected_doc as u32)) .collect::>(); @@ -4652,11 +4803,7 @@ mod tests { params.phrase_slop = Some(0); let error = wand - .search( - ¶ms, - Arc::new(RowAddrMask::default()), - &NoOpMetricsCollector, - ) + .search(¶ms, &NoOpMetricsCollector) .expect_err("corrupt packed positions should fail the phrase search"); let message = error.to_string(); assert!( @@ -4666,13 +4813,10 @@ mod tests { assert!(message.contains("corrupt"), "{message}"); } - fn sorted_candidate_row_ids(candidates: Vec) -> Vec { + fn sorted_candidate_row_ids(candidates: Vec>) -> Vec { let mut row_ids = candidates .into_iter() - .map(|candidate| match candidate.addr { - CandidateAddr::RowId(row_id) => row_id, - CandidateAddr::Pending(doc_id) => doc_id as u64, - }) + .map(|candidate| candidate.document) .collect::>(); row_ids.sort_unstable(); row_ids @@ -4713,11 +4857,7 @@ mod tests { let mut wand = Wand::new(Operator::And, postings.into_iter(), &docs, bm25); // This should trigger the bug when the second posting list becomes empty let result = wand - .search( - &FtsSearchParams::default(), - Arc::new(RowAddrMask::default()), - &NoOpMetricsCollector, - ) + .search(&FtsSearchParams::default(), &NoOpMetricsCollector) .unwrap(); assert_eq!(result.len(), 0); // Should not panic } @@ -4770,7 +4910,7 @@ mod tests { let floor = shared_floor.cloned().unwrap_or_else(new_floor); Wand::new(Operator::Or, postings.into_iter(), &docs, UnitScorer) .with_shared_threshold(floor) - .search(¶ms, Arc::new(RowAddrMask::default()), &metrics) + .search(¶ms, &metrics) .unwrap(); } metrics.0.load(Ordering::Relaxed) @@ -4842,11 +4982,7 @@ mod tests { // but the sum of block max scores is less than the threshold, wand.threshold = 1.5; - let result = wand.search( - &FtsSearchParams::default(), - Arc::new(RowAddrMask::default()), - &NoOpMetricsCollector, - ); + let result = wand.search(&FtsSearchParams::default(), &NoOpMetricsCollector); assert!(result.is_ok()); } @@ -4894,20 +5030,8 @@ mod tests { scored: scored.clone(), }, ); - let hits = wand - .search( - ¶ms, - Arc::new(RowAddrMask::default()), - &NoOpMetricsCollector, - ) - .unwrap(); - let mut row_ids = hits - .iter() - .map(|hit| match hit.addr { - CandidateAddr::RowId(r) => r, - CandidateAddr::Pending(_) => panic!("row_id should be set in this path"), - }) - .collect::>(); + let hits = wand.search(¶ms, &NoOpMetricsCollector).unwrap(); + let mut row_ids = hits.iter().map(|hit| hit.document).collect::>(); row_ids.sort_unstable(); (row_ids, scored.load(Ordering::Relaxed)) }; @@ -4955,13 +5079,7 @@ mod tests { let mut wand = Wand::new(Operator::Or, postings.into_iter(), &docs, UnitScorer); let metrics = PanicOnAndMetrics::new(); - let candidates = wand - .search( - &FtsSearchParams::default(), - Arc::new(RowAddrMask::default()), - &metrics, - ) - .unwrap(); + let candidates = wand.search(&FtsSearchParams::default(), &metrics).unwrap(); assert_eq!(sorted_candidate_row_ids(candidates), vec![0, 1, 2, 4, 5]); assert!(metrics.comparisons.load(Ordering::Relaxed) > 0); @@ -5251,7 +5369,6 @@ mod tests { let result = wand .search( &FtsSearchParams::new().with_limit(Some(1)), - Arc::new(RowAddrMask::default()), &NoOpMetricsCollector, ) .unwrap(); @@ -5472,14 +5589,7 @@ mod tests { &docs, InverseDocLengthScorer, ); - sorted_candidate_row_ids( - wand.search( - ¶ms, - Arc::new(RowAddrMask::default()), - &NoOpMetricsCollector, - ) - .unwrap(), - ) + sorted_candidate_row_ids(wand.search(¶ms, &NoOpMetricsCollector).unwrap()) }; let compressed = run(true); @@ -5775,13 +5885,15 @@ mod tests { let result = wand .search( &FtsSearchParams::new().with_limit(Some(1)), - Arc::new(RowAddrMask::default()), &NoOpMetricsCollector, ) .unwrap(); - let addrs = result.into_iter().map(|doc| doc.addr).collect::>(); - assert!(matches!(addrs.as_slice(), [CandidateAddr::RowId(0)])); + let addrs = result + .into_iter() + .map(|doc| doc.document) + .collect::>(); + assert_eq!(addrs, vec![0]); let scored = scored.load(Ordering::Relaxed); // The bulk path evaluates 63 doc weights up front to fill its // frequency-bound prune LUT; those are bound computations, not @@ -5830,15 +5942,14 @@ mod tests { ); let metrics = CountAndSearchStats::default(); let result = wand - .search( - &FtsSearchParams::new().with_limit(Some(1)), - Arc::new(RowAddrMask::default()), - &metrics, - ) + .search(&FtsSearchParams::new().with_limit(Some(1)), &metrics) .unwrap(); - let addrs = result.into_iter().map(|doc| doc.addr).collect::>(); - assert!(matches!(addrs.as_slice(), [CandidateAddr::RowId(0)])); + let addrs = result + .into_iter() + .map(|doc| doc.document) + .collect::>(); + assert_eq!(addrs, vec![0]); let candidates_seen = metrics.candidates_seen.load(Ordering::Relaxed); let candidates_pruned_before_return = metrics @@ -5894,13 +6005,15 @@ mod tests { let result = wand .search( &FtsSearchParams::new().with_limit(Some(1)), - Arc::new(RowAddrMask::default()), &NoOpMetricsCollector, ) .unwrap(); - let addrs = result.into_iter().map(|doc| doc.addr).collect::>(); - assert!(matches!(addrs.as_slice(), [CandidateAddr::RowId(1)])); + let addrs = result + .into_iter() + .map(|doc| doc.document) + .collect::>(); + assert_eq!(addrs, vec![1]); } #[rstest] @@ -5972,7 +6085,7 @@ mod tests { ); wand.threshold = 0.5; - let selected = vec![RowAddress::from(1_u64), RowAddress::from(2_u64)]; + let selected = vec![(1_u64, 1_u64), (2_u64, 2_u64)]; let result = wand .flat_search( &FtsSearchParams::default(), @@ -5983,10 +6096,7 @@ mod tests { let matched = result .into_iter() - .map(|doc| match doc.addr { - CandidateAddr::RowId(r) => r, - CandidateAddr::Pending(_) => panic!("row_id should be set in this path"), - }) + .map(|doc| doc.document) .collect::>(); assert_eq!(matched, vec![2]); } @@ -6039,7 +6149,10 @@ mod tests { ); wand.threshold = 0.5; - let selected = vec![RowAddress::from(100_u64)]; + let selected = docs + .doc_ids(100) + .map(|doc_id| (doc_id, 100_u64)) + .collect::>(); let result = wand .flat_search( &FtsSearchParams::default(), @@ -6048,13 +6161,14 @@ mod tests { ) .unwrap(); - // flat_search resolves the prefilter against the DocSet, so the single - // match comes back as a concrete RowId(100) rather than a deferred - // Pending addr. Asserting on the whole result avoids a never-taken - // match arm that would otherwise read as uncovered. - let addrs = result.into_iter().map(|doc| doc.addr).collect::>(); + // The legacy adapter resolves the prefilter to every owned document, + // while the candidate identity remains row 100. + let addrs = result + .into_iter() + .map(|doc| doc.document) + .collect::>(); assert!( - matches!(addrs.as_slice(), [CandidateAddr::RowId(100)]), + addrs.as_slice() == [100], "expected exactly row 100, got {addrs:?}" ); } @@ -6078,6 +6192,23 @@ mod tests { ); } + #[test] + fn test_modern_doc_id_validation_checks_layout_and_upper_bound() { + let compressed = generate_posting_list(vec![0, 4], 1.0, None, true); + let posting = PostingIterator::new(String::from("term"), 0, 0, compressed, 5); + posting + .validate_modern_doc_ids(5) + .expect("largest DocId is inside the document table"); + let error = posting.validate_modern_doc_ids(4).unwrap_err(); + assert!(error.to_string().contains("DocId 4")); + assert!(error.to_string().contains("[0, 4)")); + + let plain = generate_posting_list(vec![0], 1.0, None, false); + let posting = PostingIterator::new(String::from("legacy"), 0, 0, plain, 1); + let error = posting.validate_modern_doc_ids(1).unwrap_err(); + assert!(error.to_string().contains("legacy row-address layout")); + } + #[test] fn test_v3_without_impacts_uses_conservative_quantized_score_bound() { let exact_doc_length = 300; @@ -6380,7 +6511,7 @@ mod tests { params.phrase_slop = Some(slop); } - let normalize = |result: Vec| { + let normalize = |result: Vec>| { let mut rows = result .into_iter() .map(|candidate| { @@ -6388,10 +6519,7 @@ mod tests { candidate.posting_doc_id, candidate.doc_length, candidate.freqs, - match candidate.addr { - CandidateAddr::RowId(row_id) => row_id, - CandidateAddr::Pending(doc_id) => u64::from(doc_id), - }, + candidate.document, ) }) .collect::>(); @@ -6407,14 +6535,7 @@ mod tests { UnitScorer, ) .with_bulk_and_mode(mode); - let rows = normalize( - wand.search( - ¶ms, - Arc::new(RowAddrMask::default()), - &NoOpMetricsCollector, - ) - .unwrap(), - ); + let rows = normalize(wand.search(¶ms, &NoOpMetricsCollector).unwrap()); let used_bulk = wand.bulk_and_searches > 0; (rows, used_bulk) }; From 2e359e8d374df91c082930221fc3b10186849772 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 20 Jul 2026 22:38:18 +0800 Subject: [PATCH 02/10] perf(fts): cache modern posting validation --- rust/lance-index/src/scalar/inverted/index.rs | 111 ++++++++++++++++-- 1 file 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 422c0128194..34f62ec4c8a 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -1862,7 +1862,7 @@ impl InvertedPartition { let token_file = store.open_index_file(&token_file_path(id)).await?; let tokens = TokenSet::load(token_file, token_set_format).await?; let invert_list_file = store.open_index_file(&posting_file_path(id)).await?; - let inverted_list = PostingListReader::try_new(invert_list_file, index_cache).await?; + let mut inverted_list = PostingListReader::try_new(invert_list_file, index_cache).await?; let docs_path = doc_file_path(id); let docs_reader = store.open_index_file(&docs_path).await?; let docs = PartitionDocuments::try_new( @@ -1872,6 +1872,7 @@ impl InvertedPartition { frag_reuse_index, inverted_list.block_size() == MAX_POSTING_BLOCK_SIZE, )?; + inverted_list.modern_num_docs = Some(docs.len()); Ok(Self { id, @@ -2330,12 +2331,6 @@ impl InvertedPartition { .try_collect::>() .await?; - if self.docs.modern().is_some() { - for (_, token, _, posting) in &loaded_postings { - validate_modern_posting_doc_ids(posting, token, num_docs)?; - } - } - let needs_union = loaded_postings .windows(2) .any(|window| window[0].2 == window[1].2); @@ -2986,6 +2981,12 @@ pub struct PostingListReader { /// index or relying on persisted grouping metadata. grouping: PostingGrouping, + /// Modern postings contain dense DocIds into the partition document table. + /// Cache successful boundary validation per immutable token so repeated + /// queries do not decode the final posting block again. + modern_doc_id_validations: Option>, + modern_num_docs: Option, + index_cache: WeakLanceCache, } @@ -3059,7 +3060,16 @@ impl DeepSizeOf for PostingListReader { }) .unwrap_or(0), }; - metadata_size + self.grouping.deep_size_of_children(context) + let validation_size = self + .modern_doc_id_validations + .as_ref() + .map(|validations| { + validations + .len() + .saturating_mul(std::mem::size_of::()) + }) + .unwrap_or(0); + metadata_size + self.grouping.deep_size_of_children(context) + validation_size } } @@ -3093,6 +3103,12 @@ impl PostingListReader { let is_legacy_layout = matches!(&metadata, PostingMetadata::LegacyV1 { .. }); let grouping = PostingGrouping::for_reader(is_legacy_layout, reader.num_rows()); + let modern_doc_id_validations = (!is_legacy_layout).then(|| { + (0..reader.num_rows()) + .map(|_| AtomicBool::new(false)) + .collect::>() + .into() + }); Ok(Self { reader, @@ -3103,6 +3119,8 @@ impl PostingListReader { block_size, positions_layout, grouping, + modern_doc_id_validations, + modern_num_docs: None, index_cache: WeakLanceCache::from(index_cache), }) } @@ -3388,6 +3406,8 @@ impl PostingListReader { .clone(), }; + self.ensure_modern_posting_validated(token_id, &posting)?; + if is_phrase_query && !posting.has_position() { // hit the cache and when the cache was populated, the positions column was not loaded let positions = self.read_positions(token_id).await?; @@ -3397,6 +3417,26 @@ impl PostingListReader { Ok(posting) } + fn ensure_modern_posting_validated(&self, token_id: u32, posting: &PostingList) -> Result<()> { + let (Some(validations), Some(num_docs)) = + (&self.modern_doc_id_validations, self.modern_num_docs) + else { + return Ok(()); + }; + let validation = validations.get(token_id as usize).ok_or_else(|| { + Error::index(format!( + "modern FTS token id {token_id} is outside validation state [0, {})", + validations.len() + )) + })?; + if validation.load(Ordering::Acquire) { + return Ok(()); + } + validate_modern_posting_doc_ids(posting, &format!("token id {token_id}"), num_docs)?; + validation.store(true, Ordering::Release); + Ok(()) + } + /// Map a token id to its cache group's row range `[start, end)`, or `None` /// when grouping is not available so the caller falls back to the per-token /// path. In v2 the token id is the row offset, so the group range is also @@ -12234,6 +12274,61 @@ mod tests { posting.iter().map(|(doc, freq, _)| (doc, freq)).collect() } + #[tokio::test] + async fn test_modern_posting_validation_is_cached_per_token() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + let mut builder = InnerBuilder::new(0, false, TokenSetFormat::default()); + builder.tokens.add("term".to_owned()); + let mut valid_builder = PostingListBuilder::new(false); + valid_builder.add(0, PositionRecorder::Count(1)); + builder.posting_lists.push(valid_builder); + builder.docs.append(1000, 1); + builder.write(store.as_ref()).await.unwrap(); + + let reader = store.open_index_file(&posting_file_path(0)).await.unwrap(); + let mut posting_reader = PostingListReader::try_new(reader, &LanceCache::no_cache()) + .await + .unwrap(); + posting_reader.modern_num_docs = Some(1); + let validation = &posting_reader + .modern_doc_id_validations + .as_ref() + .expect("modern readers have per-token validation state")[0]; + assert!(!validation.load(Ordering::Acquire)); + + let mut corrupt_builder = PostingListBuilder::new(false); + corrupt_builder.add(1, PositionRecorder::Count(1)); + let corrupt_batch = corrupt_builder.to_batch(vec![1.0]).unwrap(); + let corrupt_posting = PostingList::from_batch(&corrupt_batch, Some(1.0), Some(1)).unwrap(); + let error = posting_reader + .ensure_modern_posting_validated(0, &corrupt_posting) + .unwrap_err(); + assert!(matches!(error, Error::Index { .. })); + assert!(error.to_string().contains("DocId 1")); + assert!(error.to_string().contains("[0, 1)")); + assert!(!validation.load(Ordering::Acquire)); + + let first = posting_reader + .posting_list(0, false, &NoOpMetricsCollector) + .await + .unwrap(); + assert_eq!(posting_entries(&first), vec![(0, 1)]); + assert!(validation.load(Ordering::Acquire)); + + let second = posting_reader + .posting_list(0, false, &NoOpMetricsCollector) + .await + .unwrap(); + assert_eq!(posting_entries(&second), vec![(0, 1)]); + assert!(validation.load(Ordering::Acquire)); + } + /// Runtime synthetic grouping must return correct posting lists for every /// token, including across synthetic group boundaries. #[tokio::test] From 0ccddb193835c3ff5148a3181d34294b27029424 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 21 Jul 2026 01:48:58 +0800 Subject: [PATCH 03/10] perf(index): optimize filtered FTS document access --- .../src/scalar/inverted/documents.rs | 364 ++++++++++++++++-- rust/lance-index/src/scalar/inverted/index.rs | 59 ++- 2 files changed, 368 insertions(+), 55 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/documents.rs b/rust/lance-index/src/scalar/inverted/documents.rs index 3d8b66ad6ae..8736f560c2a 100644 --- a/rust/lance-index/src/scalar/inverted/documents.rs +++ b/rust/lance-index/src/scalar/inverted/documents.rs @@ -13,8 +13,10 @@ use std::sync::{Arc, OnceLock}; use arrow::buffer::ScalarBuffer; use arrow_array::{Array, RecordBatch, UInt32Array, UInt64Array}; use lance_core::deepsize::DeepSizeOf; +use lance_core::utils::address::RowAddress; +use lance_core::utils::tokio::spawn_cpu; use lance_core::{Error, ROW_ID, Result}; -use lance_select::RowAddrMask; +use lance_select::{RowAddrMask, RowAddrSelection, RowAddrTreeMap}; use object_store::path::Path; use roaring::RoaringBitmap; use tokio::sync::OnceCell; @@ -185,6 +187,140 @@ impl DeepSizeOf for AddressValues { } } +/// A compact address-sorted view of the live DocIds in a projection. +/// +/// Most newly built partitions already store addresses in DocId order, so the +/// identity variant adds no per-document memory. Remapped or otherwise +/// unsorted projections keep only a u32 permutation instead of duplicating +/// the u64 address column. +#[derive(Debug)] +enum AddressDocIdLookup { + Identity, + Sorted(Box<[u32]>), +} + +impl DeepSizeOf for AddressDocIdLookup { + fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize { + match self { + Self::Identity => 0, + Self::Sorted(doc_ids) => std::mem::size_of_val(doc_ids.as_ref()), + } + } +} + +impl AddressDocIdLookup { + fn build(projection: &VersionAddressProjection) -> Self { + if projection.live_docs.is_none() + && (1..projection.addresses.len()) + .all(|index| projection.addresses.get(index - 1) <= projection.addresses.get(index)) + { + return Self::Identity; + } + + let mut doc_ids = match projection.live_docs.as_ref() { + Some(live_docs) => live_docs.iter().collect::>(), + None => (0..projection.addresses.len() as u32).collect::>(), + }; + doc_ids.sort_unstable_by_key(|&doc_id| projection.addresses.get(doc_id as usize)); + Self::Sorted(doc_ids.into_boxed_slice()) + } + + fn len(&self, projection: &VersionAddressProjection) -> usize { + match self { + Self::Identity => projection.addresses.len(), + Self::Sorted(doc_ids) => doc_ids.len(), + } + } + + fn doc_id_at(&self, position: usize) -> u32 { + match self { + Self::Identity => position as u32, + Self::Sorted(doc_ids) => doc_ids[position], + } + } + + fn address_at(&self, projection: &VersionAddressProjection, position: usize) -> u64 { + projection.addresses.get(self.doc_id_at(position) as usize) + } + + fn partition_point( + &self, + projection: &VersionAddressProjection, + mut predicate: impl FnMut(u64) -> bool, + ) -> usize { + let mut left = 0; + let mut right = self.len(projection); + while left < right { + let middle = left + (right - left) / 2; + if predicate(self.address_at(projection, middle)) { + left = middle + 1; + } else { + right = middle; + } + } + left + } + + fn insert_address_range( + &self, + projection: &VersionAddressProjection, + start: u64, + end: u64, + selected: &mut RoaringBitmap, + ) { + let first = self.partition_point(projection, |address| address < start); + let after_last = self.partition_point(projection, |address| address <= end); + for position in first..after_last { + selected.insert(self.doc_id_at(position)); + } + } + + fn matching_doc_ids( + &self, + projection: &VersionAddressProjection, + addresses: &RowAddrTreeMap, + ) -> RoaringBitmap { + let mut selected = RoaringBitmap::new(); + for (&fragment_id, selection) in addresses.iter() { + match selection { + RowAddrSelection::Full => { + let start = u64::from(RowAddress::new_from_parts(fragment_id, 0)); + let end = u64::from(RowAddress::new_from_parts(fragment_id, u32::MAX)); + self.insert_address_range(projection, start, end, &mut selected); + } + RowAddrSelection::Partial(offsets) => { + let mut offsets = offsets.iter(); + while let Some(range) = offsets.next_range() { + let start = + u64::from(RowAddress::new_from_parts(fragment_id, *range.start())); + let end = u64::from(RowAddress::new_from_parts(fragment_id, *range.end())); + self.insert_address_range(projection, start, end, &mut selected); + } + } + } + } + selected + } + + fn visibility( + &self, + projection: &VersionAddressProjection, + mask: &RowAddrMask, + ) -> DocVisibility { + match mask { + RowAddrMask::AllowList(allowed) => { + DocVisibility::Selected(self.matching_doc_ids(projection, allowed)) + } + RowAddrMask::BlockList(blocked) => { + let blocked = self.matching_doc_ids(projection, blocked); + let mut selected = projection.live_doc_ids(); + selected -= &blocked; + DocVisibility::Selected(selected) + } + } + } +} + /// Addresses projected into the dataset version that opened the index. #[derive(Debug)] pub(super) struct VersionAddressProjection { @@ -192,6 +328,7 @@ pub(super) struct VersionAddressProjection { /// `None` means every slot is live. Deleted documents retain their DocId /// slot but are absent from this bitmap. live_docs: Option, + doc_ids_by_address: OnceCell>, } impl DeepSizeOf for VersionAddressProjection { @@ -202,6 +339,11 @@ impl DeepSizeOf for VersionAddressProjection { .as_ref() .map(|live| live.serialized_size()) .unwrap_or(0) + + self + .doc_ids_by_address + .get() + .map(|lookup| lookup.deep_size_of_children(context)) + .unwrap_or(0) } } @@ -229,6 +371,7 @@ impl VersionAddressProjection { return Ok(Self { addresses: AddressValues::Shared(raw.values().clone()), live_docs: None, + doc_ids_by_address: OnceCell::new(), }); }; @@ -250,6 +393,7 @@ impl VersionAddressProjection { Ok(Self { addresses: AddressValues::Owned(addresses), live_docs: Some(live_docs), + doc_ids_by_address: OnceCell::new(), }) } @@ -265,19 +409,32 @@ impl VersionAddressProjection { } } - fn visibility(&self, mask: &RowAddrMask) -> DocVisibility { - if mask.is_select_all() && self.live_docs.is_none() { - return DocVisibility::All; - } - let selected = (0..self.addresses.len()) - .filter_map(|doc_id| { - let doc_id = DocId::new(doc_id as u32); - self.address(doc_id) - .filter(|address| mask.selected(*address)) - .map(|_| doc_id.get()) + fn live_doc_ids(&self) -> RoaringBitmap { + self.live_docs + .clone() + .unwrap_or_else(|| (0..self.addresses.len() as u32).collect()) + } + + async fn doc_ids_by_address(self: &Arc) -> Result> { + self.doc_ids_by_address + .get_or_try_init(|| { + let projection = self.clone(); + async move { + spawn_cpu(move || Result::Ok(Arc::new(AddressDocIdLookup::build(&projection)))) + .await + } }) - .collect(); - DocVisibility::Selected(selected) + .await + .cloned() + } + + async fn materialize_visibility( + self: &Arc, + mask: Arc, + ) -> Result { + let lookup = self.doc_ids_by_address().await?; + let projection = self.clone(); + spawn_cpu(move || Result::Ok(lookup.visibility(&projection, &mask))).await } } @@ -286,6 +443,10 @@ impl VersionAddressProjection { pub(super) enum DocVisibility { All, Selected(RoaringBitmap), + Filtered { + projection: Arc, + mask: Arc, + }, } impl DocVisibility { @@ -294,6 +455,9 @@ impl DocVisibility { match self { Self::All => true, Self::Selected(selected) => selected.contains(doc_id.get()), + Self::Filtered { projection, mask } => projection + .address(doc_id) + .is_some_and(|address| mask.selected(address)), } } @@ -301,6 +465,7 @@ impl DocVisibility { match self { Self::All => total_docs, Self::Selected(selected) => selected.len() as usize, + Self::Filtered { .. } => total_docs, } } @@ -310,8 +475,8 @@ impl DocVisibility { pub(crate) fn iter(&self) -> Option + '_> { match self { - Self::All => None, Self::Selected(selected) => Some(selected.iter().map(DocId::new)), + Self::All | Self::Filtered { .. } => None, } } } @@ -550,11 +715,26 @@ impl PartitionDocuments { .cloned() } - pub(crate) async fn visibility(&self, mask: &RowAddrMask) -> Result { + pub(crate) async fn visibility( + &self, + mask: Arc, + materialize_selected: bool, + ) -> Result { + if mask.max_len() == Some(0) { + return Ok(DocVisibility::Selected(RoaringBitmap::new())); + } if mask.is_select_all() && self.remapper.is_none() { - Ok(DocVisibility::All) + return Ok(DocVisibility::All); + } + + let projection = self.projection().await?; + if mask.is_select_all() { + return Ok(DocVisibility::Selected(projection.live_doc_ids())); + } + if materialize_selected { + projection.materialize_visibility(mask).await } else { - Ok(self.projection().await?.visibility(mask)) + Ok(DocVisibility::Filtered { projection, mask }) } } @@ -628,7 +808,7 @@ impl PartitionDocuments { format!("{ROW_ID} contains null values"), )); } - let addresses = address_column.values().to_vec(); + let addresses = address_column.values(); if use_bulk && addresses.len() != self.num_docs { return Err(corrupt_docs( &self.path, @@ -658,7 +838,7 @@ impl PartitionDocuments { } else { let by_doc_id = unique .into_iter() - .zip(addresses) + .zip(addresses.iter().copied()) .collect::>(); doc_ids .iter() @@ -681,7 +861,7 @@ impl PartitionDocuments { pub(crate) async fn prewarm(&self) -> Result<()> { self.lengths().await?; - self.projection().await?; + self.projection().await?.doc_ids_by_address().await?; Ok(()) } } @@ -1061,49 +1241,127 @@ mod tests { } #[test] - fn visibility_uses_doc_ids() { + fn live_documents_use_doc_ids() { let projection = VersionAddressProjection { addresses: AddressValues::Owned(vec![10, 20, 30]), live_docs: Some(RoaringBitmap::from_iter([0, 2])), + doc_ids_by_address: OnceCell::new(), }; - let mask = RowAddrMask::all_rows(); - let DocVisibility::Selected(selected) = projection.visibility(&mask) else { - panic!("remapped projection must preserve its live-doc bitmap") - }; + let selected = projection.live_doc_ids(); assert_eq!(selected.iter().collect::>(), vec![0, 2]); } - #[test] - fn remap_preserves_doc_id_slots_and_filters_in_current_address_domain() { + #[tokio::test] + async fn remap_preserves_doc_id_slots_and_filters_in_current_address_domain() { let raw = UInt64Array::from(vec![10, 20, 30, 40]); let remapper = TestRemapper { mapping: HashMap::from([(10, Some(100)), (20, None), (30, Some(300))]), }; - let projection = VersionAddressProjection::try_new(&raw, 4, Some(&remapper), "docs") - .expect("valid projection"); + let projection = Arc::new( + VersionAddressProjection::try_new(&raw, 4, Some(&remapper), "docs") + .expect("valid projection"), + ); assert_eq!(projection.address(DocId::new(0)), Some(100)); assert_eq!(projection.address(DocId::new(1)), None); assert_eq!(projection.address(DocId::new(2)), Some(300)); assert_eq!(projection.address(DocId::new(3)), Some(40)); - let DocVisibility::Selected(all_live) = projection.visibility(&RowAddrMask::all_rows()) - else { - panic!("a remapped projection must retain a live-doc bitmap") - }; + let all_live = projection.live_doc_ids(); assert_eq!(all_live.iter().collect::>(), vec![0, 2, 3]); - let allowed = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([100, 40])); - let DocVisibility::Selected(selected) = projection.visibility(&allowed) else { + let allowed = Arc::new(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([ + 100, 40, + ]))); + let DocVisibility::Selected(selected) = projection + .materialize_visibility(allowed) + .await + .expect("valid allow-list") + else { panic!("allow-list must compile to DocIds") }; assert_eq!(selected.iter().collect::>(), vec![0, 3]); - let blocked = RowAddrMask::from_block(RowAddrTreeMap::from_iter([300])); - let DocVisibility::Selected(selected) = projection.visibility(&blocked) else { + let first_lookup = projection + .doc_ids_by_address() + .await + .expect("cached address lookup"); + let blocked = Arc::new(RowAddrMask::from_block(RowAddrTreeMap::from_iter([300]))); + let DocVisibility::Selected(selected) = projection + .materialize_visibility(blocked) + .await + .expect("valid block-list") + else { panic!("block-list must compile to DocIds") }; assert_eq!(selected.iter().collect::>(), vec![0, 3]); + let second_lookup = projection + .doc_ids_by_address() + .await + .expect("cached address lookup"); + assert!(Arc::ptr_eq(&first_lookup, &second_lookup)); + } + + #[tokio::test] + async fn materialized_visibility_handles_unsorted_duplicate_addresses_and_full_fragments() { + let row_address = |fragment_id, row_offset| { + u64::from(RowAddress::new_from_parts(fragment_id, row_offset)) + }; + let projection = Arc::new(VersionAddressProjection { + addresses: AddressValues::Owned(vec![ + row_address(2, 5), + row_address(1, 2), + row_address(1, 2), + row_address(1, 4), + ]), + live_docs: None, + doc_ids_by_address: OnceCell::new(), + }); + + let allowed = Arc::new(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([ + row_address(1, 2), + row_address(1, 3), + row_address(1, 4), + ]))); + let DocVisibility::Selected(selected) = projection + .materialize_visibility(allowed) + .await + .expect("valid allow-list") + else { + panic!("allow-list must compile to DocIds") + }; + assert_eq!(selected.iter().collect::>(), vec![1, 2, 3]); + + let mut full_fragment = RowAddrTreeMap::new(); + full_fragment.insert_fragment(2); + let DocVisibility::Selected(selected) = projection + .materialize_visibility(Arc::new(RowAddrMask::from_allowed(full_fragment))) + .await + .expect("valid full-fragment allow-list") + else { + panic!("allow-list must compile to DocIds") + }; + assert_eq!(selected.iter().collect::>(), vec![0]); + } + + #[test] + fn lazy_visibility_projects_only_candidate_doc_ids() { + let projection = Arc::new(VersionAddressProjection { + addresses: AddressValues::Owned(vec![10, 20, 30]), + live_docs: Some(RoaringBitmap::from_iter([0, 2])), + doc_ids_by_address: OnceCell::new(), + }); + let visibility = DocVisibility::Filtered { + projection: projection.clone(), + mask: Arc::new(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([ + 20, 30, + ]))), + }; + + assert!(!visibility.selected(DocId::new(0))); + assert!(!visibility.selected(DocId::new(1))); + assert!(visibility.selected(DocId::new(2))); + assert!(!projection.doc_ids_by_address.initialized()); } #[test] @@ -1192,6 +1450,38 @@ mod tests { assert!(!documents.projection_loaded()); } + #[tokio::test] + async fn prewarm_loads_document_columns_and_address_lookup_once() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + Some("10"), + ) + .await; + let (counting, counts) = counted_store(store, path); + let documents = open_documents(counting, path, None).await.unwrap(); + + documents.prewarm().await.unwrap(); + let projection = documents.projection().await.unwrap(); + let first_lookup = projection.doc_ids_by_address().await.unwrap(); + documents.prewarm().await.unwrap(); + let second_lookup = projection.doc_ids_by_address().await.unwrap(); + + assert!(documents.lengths_loaded()); + assert!(documents.projection_loaded()); + assert!(matches!( + first_lookup.as_ref(), + AddressDocIdLookup::Identity + )); + assert!(Arc::ptr_eq(&first_lookup, &second_lookup)); + assert_eq!(counts.length_rows.load(Ordering::Relaxed), 3); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 3); + } + #[tokio::test] async fn invalid_or_mismatched_stats_are_corruption_not_fallback() { let (_directory, store) = test_store(); diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 34f62ec4c8a..74e32ce6f25 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -1144,7 +1144,17 @@ impl InvertedIndex { let documents = part.docs.modern().cloned().ok_or_else(|| { Error::internal("modern index contains legacy partition documents") })?; - let visibility = Arc::new(documents.visibility(mask.as_ref()).await?); + let materialize_selected = operator == Operator::Or + && mask.max_len().is_some_and(|selected| { + u128::from(selected).saturating_mul(100) + <= u128::from(*FLAT_SEARCH_PERCENT_THRESHOLD) + .saturating_mul(documents.len() as u128) + }); + let visibility = Arc::new( + documents + .visibility(mask.clone(), materialize_selected) + .await?, + ); if visibility.is_empty() { return Result::Ok((partition_ordinal, PartitionCandidates::empty())); } @@ -1213,16 +1223,26 @@ impl InvertedIndex { .or_default() .push((rank, doc_id)); } - for (partition_ordinal, entries) in by_partition { - let documents = self.partitions[partition_ordinal] - .docs - .modern() - .expect("modern search only groups modern partitions"); - let doc_ids = entries - .iter() - .map(|(_, doc_id)| *doc_id) - .collect::>(); - let resolved = documents.resolve_addresses(&doc_ids).await?; + let address_reads = by_partition + .into_iter() + .map(|(partition_ordinal, entries)| { + let documents = self.partitions[partition_ordinal] + .docs + .modern() + .cloned() + .expect("modern search only groups modern partitions"); + async move { + let doc_ids = entries + .iter() + .map(|(_, doc_id)| *doc_id) + .collect::>(); + let resolved = documents.resolve_addresses(&doc_ids).await?; + Result::Ok((entries, resolved)) + } + }); + let mut address_reads = + stream::iter(address_reads).buffer_unordered(self.store.io_parallelism().max(1)); + while let Some((entries, resolved)) = address_reads.try_next().await? { for ((rank, _), address) in entries.into_iter().zip(resolved) { addresses[rank] = address; } @@ -9591,15 +9611,21 @@ mod tests { let all_rows = RowAddrMask::all_rows(); assert!(matches!( - documents.visibility(&all_rows).await.unwrap(), + documents + .visibility(Arc::new(all_rows), false) + .await + .unwrap(), DocVisibility::All )); assert!(!documents.projection_loaded()); let filtered = RowAddrMask::allow_nothing(); - let visibility = documents.visibility(&filtered).await.unwrap(); + let visibility = documents + .visibility(Arc::new(filtered), true) + .await + .unwrap(); assert!(visibility.is_empty()); - assert!(documents.projection_loaded()); + assert!(!documents.projection_loaded()); assert_eq!( documents .resolve_addresses(&[DocId::new(0), DocId::new(99)]) @@ -10427,10 +10453,7 @@ mod tests { let documents = partition.docs.modern().unwrap(); let lengths = documents.lengths().await.unwrap(); - let visibility = documents - .visibility(NoFilter.mask().as_ref()) - .await - .unwrap(); + let visibility = documents.visibility(NoFilter.mask(), false).await.unwrap(); partition .bm25_search_modern( lengths.as_ref(), From 82f3e795d9d1156cdffe526d792f835b6dbd5f00 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 21 Jul 2026 11:40:48 +0800 Subject: [PATCH 04/10] perf(fts): specialize fully prewarmed document access --- .../src/scalar/inverted/documents.rs | 153 +++++++++++--- rust/lance-index/src/scalar/inverted/index.rs | 191 +++++++++++++++--- 2 files changed, 291 insertions(+), 53 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/documents.rs b/rust/lance-index/src/scalar/inverted/documents.rs index 8736f560c2a..95da7dc947c 100644 --- a/rust/lance-index/src/scalar/inverted/documents.rs +++ b/rust/lance-index/src/scalar/inverted/documents.rs @@ -491,6 +491,7 @@ pub(super) struct PartitionDocuments { remapper: Option>, lengths: OnceCell>, projection: OnceCell>, + prewarm_complete: OnceCell<()>, } /// Load-boundary discriminator between the read-only legacy representation and @@ -631,6 +632,7 @@ impl PartitionDocuments { remapper, lengths: OnceCell::new(), projection: OnceCell::new(), + prewarm_complete: OnceCell::new(), }) } @@ -643,7 +645,6 @@ impl PartitionDocuments { self.lengths.initialized() } - #[cfg(test)] pub(crate) fn projection_loaded(&self) -> bool { self.projection.initialized() } @@ -652,6 +653,33 @@ impl PartitionDocuments { self.store.open_index_file(&self.path).await } + fn lengths_from_batch(&self, batch: &RecordBatch) -> Result> { + let column = required_u32_column(batch, NUM_TOKEN_COL, &self.path)?; + if column.null_count() != 0 { + return Err(corrupt_docs( + &self.path, + format!("{NUM_TOKEN_COL} contains null values"), + )); + } + Ok(Arc::new(DocLengths::try_new( + column.values().clone(), + self.num_docs, + self.persisted_total_tokens, + self.quantized_scoring, + &self.path, + )?)) + } + + fn projection_from_batch(&self, batch: &RecordBatch) -> Result> { + let column = required_u64_column(batch, ROW_ID, &self.path)?; + Ok(Arc::new(VersionAddressProjection::try_new( + column, + self.num_docs, + self.remapper.as_deref(), + &self.path, + )?)) + } + pub(crate) async fn stats(&self) -> Result { let total_tokens = match self.persisted_total_tokens { Some(total_tokens) => total_tokens, @@ -679,20 +707,7 @@ impl PartitionDocuments { let batch = reader .read_range(0..self.num_docs, Some(&[NUM_TOKEN_COL])) .await?; - let column = required_u32_column(&batch, NUM_TOKEN_COL, &self.path)?; - if column.null_count() != 0 { - return Err(corrupt_docs( - &self.path, - format!("{NUM_TOKEN_COL} contains null values"), - )); - } - Ok(Arc::new(DocLengths::try_new( - column.values().clone(), - self.num_docs, - self.persisted_total_tokens, - self.quantized_scoring, - &self.path, - )?)) + self.lengths_from_batch(&batch) }) .await .cloned() @@ -703,13 +718,7 @@ impl PartitionDocuments { .get_or_try_init(|| async { let reader = self.reader().await?; let batch = reader.read_range(0..self.num_docs, Some(&[ROW_ID])).await?; - let column = required_u64_column(&batch, ROW_ID, &self.path)?; - Ok(Arc::new(VersionAddressProjection::try_new( - column, - self.num_docs, - self.remapper.as_deref(), - &self.path, - )?)) + self.projection_from_batch(&batch) }) .await .cloned() @@ -859,9 +868,56 @@ impl PartitionDocuments { DocSet::load(self.reader().await?, false, self.remapper.clone()).await } + /// Resolve one DocId synchronously when its current-version projection is resident. + pub(crate) fn cached_row_address(&self, doc_id: DocId) -> Result> { + let Some(projection) = self.projection.get() else { + return Ok(None); + }; + if doc_id.as_usize() >= self.num_docs { + return Err(corrupt_docs( + &self.path, + format!( + "candidate DocId {} is outside [0, {})", + doc_id.get(), + self.num_docs + ), + )); + } + projection + .address(doc_id) + .map(RowAddress::new_from_u64) + .map(Some) + .ok_or_else(|| { + corrupt_docs( + &self.path, + format!("candidate DocId {} is not live", doc_id.get()), + ) + }) + } + pub(crate) async fn prewarm(&self) -> Result<()> { - self.lengths().await?; - self.projection().await?.doc_ids_by_address().await?; + self.prewarm_complete + .get_or_try_init(|| async { + if self.lengths.get().is_none() && self.projection.get().is_none() { + let reader = self.reader().await?; + let batch = reader + .read_range(0..self.num_docs, Some(&[ROW_ID, NUM_TOKEN_COL])) + .await?; + let lengths = self.lengths_from_batch(&batch)?; + let projection = self.projection_from_batch(&batch)?; + + // A concurrent single-column request may win either OnceCell. + // Awaiting the accessors below joins that initialization without + // replacing the already-published value. + let _ = self.lengths.set(lengths); + let _ = self.projection.set(projection); + } + + self.lengths().await?; + self.projection().await?.doc_ids_by_address().await?; + Result::Ok(()) + }) + .await?; Ok(()) } } @@ -955,6 +1011,7 @@ mod tests { #[derive(Debug, Default)] struct DocumentReadCounts { + open_calls: AtomicUsize, range_calls: AtomicUsize, ranges_calls: AtomicUsize, rows: AtomicUsize, @@ -1069,6 +1126,7 @@ mod tests { async fn open_index_file(&self, name: &str) -> Result> { let reader = self.inner.open_index_file(name).await?; if name == self.target { + self.counts.open_calls.fetch_add(1, Ordering::Relaxed); Ok(Arc::new(CountingReader { inner: reader, counts: self.counts.clone(), @@ -1465,7 +1523,11 @@ mod tests { let (counting, counts) = counted_store(store, path); let documents = open_documents(counting, path, None).await.unwrap(); - documents.prewarm().await.unwrap(); + futures::future::join_all((0..8).map(|_| documents.prewarm())) + .await + .into_iter() + .collect::>>() + .unwrap(); let projection = documents.projection().await.unwrap(); let first_lookup = projection.doc_ids_by_address().await.unwrap(); documents.prewarm().await.unwrap(); @@ -1478,6 +1540,37 @@ mod tests { AddressDocIdLookup::Identity )); assert!(Arc::ptr_eq(&first_lookup, &second_lookup)); + assert_eq!( + documents.cached_row_address(DocId::new(1)).unwrap(), + Some(RowAddress::new_from_u64(20)) + ); + assert_eq!(counts.open_calls.load(Ordering::Relaxed), 2); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 1); + assert_eq!(counts.length_rows.load(Ordering::Relaxed), 3); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 3); + } + + #[tokio::test] + async fn prewarm_reuses_an_already_loaded_document_column() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + Some("10"), + ) + .await; + let (counting, counts) = counted_store(store, path); + let documents = open_documents(counting, path, None).await.unwrap(); + + documents.lengths().await.unwrap(); + documents.prewarm().await.unwrap(); + documents.prewarm().await.unwrap(); + + assert_eq!(counts.open_calls.load(Ordering::Relaxed), 3); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 2); assert_eq!(counts.length_rows.load(Ordering::Relaxed), 3); assert_eq!(counts.address_rows.load(Ordering::Relaxed), 3); } @@ -1598,5 +1691,15 @@ mod tests { .to_string() .contains("_rowid contains null") ); + assert!( + documents + .prewarm() + .await + .unwrap_err() + .to_string() + .contains("_rowid contains null") + ); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); } } diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 74e32ce6f25..c33c06ed132 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -290,6 +290,17 @@ impl PartitionCandidates { } } +struct ModernSearchRequest<'a> { + tokens: Arc, + params: Arc, + operator: Operator, + mask: Arc, + metrics: Arc, + scorer: &'a MemBM25Scorer, + impact_scorer: Arc, + limit: usize, +} + fn push_scored_key( candidates: &mut BinaryHeap>, limit: usize, @@ -970,7 +981,7 @@ impl InvertedIndex { ) .await } else { - self.bm25_search_modern( + self.bm25_search_modern(ModernSearchRequest { tokens, params, operator, @@ -979,7 +990,7 @@ impl InvertedIndex { scorer, impact_scorer, limit, - ) + }) .await } } @@ -1086,18 +1097,59 @@ impl InvertedIndex { .unzip()) } - #[allow(clippy::too_many_arguments)] async fn bm25_search_modern( &self, - tokens: Arc, - params: Arc, - operator: Operator, - mask: Arc, - metrics: Arc, - scorer: &MemBM25Scorer, - impact_scorer: Arc, - limit: usize, + request: ModernSearchRequest<'_>, + ) -> Result<(Vec, Vec)> { + // Select a concrete completion path before candidate search. The + // fully resident future never builds deferred address-read state, while + // a cold query keeps DocIds until its final bounded I/O phase. + if self.has_resident_document_projections() { + self.bm25_search_modern_resident(request).await + } else { + self.bm25_search_modern_deferred(request).await + } + } + + fn has_resident_document_projections(&self) -> bool { + self.partitions.iter().all(|partition| { + partition + .docs + .modern() + .is_some_and(|documents| documents.projection_loaded()) + }) + } + + async fn bm25_search_modern_resident( + &self, + request: ModernSearchRequest<'_>, + ) -> Result<(Vec, Vec)> { + let ranked = self.bm25_search_modern_candidates(request).await?; + self.resolve_resident_modern_candidates(ranked) + } + + async fn bm25_search_modern_deferred( + &self, + request: ModernSearchRequest<'_>, ) -> Result<(Vec, Vec)> { + let ranked = self.bm25_search_modern_candidates(request).await?; + self.resolve_deferred_modern_candidates(ranked).await + } + + async fn bm25_search_modern_candidates( + &self, + request: ModernSearchRequest<'_>, + ) -> Result>> { + let ModernSearchRequest { + tokens, + params, + operator, + mask, + metrics, + scorer, + impact_scorer, + limit, + } = request; if self.partitions.len() > u32::MAX as usize { return Err(Error::index(format!( "FTS partition count {} exceeds candidate identity capacity", @@ -1212,7 +1264,44 @@ impl InvertedIndex { } } - let ranked = ranked.into_sorted_vec(); + Ok(ranked.into_sorted_vec()) + } + + fn resolve_resident_modern_candidates( + &self, + ranked: Vec>, + ) -> Result<(Vec, Vec)> { + let mut addresses = Vec::with_capacity(ranked.len()); + let mut scores = Vec::with_capacity(ranked.len()); + for Reverse(candidate) in ranked { + let partition_ordinal = (candidate.row_id >> 32) as usize; + let doc_id = DocId::new(candidate.row_id as u32); + let documents = self + .partitions + .get(partition_ordinal) + .and_then(|partition| partition.docs.modern()) + .ok_or_else(|| { + Error::internal(format!( + "resident FTS candidate DocId {} references missing modern partition ordinal {partition_ordinal}", + doc_id.get() + )) + })?; + let address = documents.cached_row_address(doc_id)?.ok_or_else(|| { + Error::internal(format!( + "resident FTS candidate DocId {} in partition ordinal {partition_ordinal} lost its cached address projection", + doc_id.get() + )) + })?; + addresses.push(address.into()); + scores.push(candidate.score.0); + } + Ok((addresses, scores)) + } + + async fn resolve_deferred_modern_candidates( + &self, + ranked: Vec>, + ) -> Result<(Vec, Vec)> { let mut addresses = vec![0_u64; ranked.len()]; let mut by_partition = BTreeMap::>::new(); for (rank, Reverse(candidate)) in ranked.iter().enumerate() { @@ -1223,23 +1312,27 @@ impl InvertedIndex { .or_default() .push((rank, doc_id)); } - let address_reads = by_partition - .into_iter() - .map(|(partition_ordinal, entries)| { - let documents = self.partitions[partition_ordinal] - .docs - .modern() - .cloned() - .expect("modern search only groups modern partitions"); - async move { - let doc_ids = entries - .iter() - .map(|(_, doc_id)| *doc_id) - .collect::>(); - let resolved = documents.resolve_addresses(&doc_ids).await?; - Result::Ok((entries, resolved)) - } + let mut address_reads = Vec::with_capacity(by_partition.len()); + for (partition_ordinal, entries) in by_partition { + let documents = self + .partitions + .get(partition_ordinal) + .and_then(|partition| partition.docs.modern()) + .cloned() + .ok_or_else(|| { + Error::internal(format!( + "deferred FTS candidates reference missing modern partition ordinal {partition_ordinal}" + )) + })?; + address_reads.push(async move { + let doc_ids = entries + .iter() + .map(|(_, doc_id)| *doc_id) + .collect::>(); + let resolved = documents.resolve_addresses(&doc_ids).await?; + Result::Ok((entries, resolved)) }); + } let mut address_reads = stream::iter(address_reads).buffer_unordered(self.store.io_parallelism().max(1)); while let Some((entries, resolved)) = address_reads.try_next().await? { @@ -10425,6 +10518,48 @@ mod tests { (tmpdir, index) } + #[tokio::test] + async fn test_prewarmed_modern_search_uses_resident_address_projection() { + let (_tmpdir, index) = load_global_scoring_test_index(true).await; + let tokens = Arc::new(Tokens::new(vec!["alpha".to_owned()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(2))); + + assert!(!index.has_resident_document_projections()); + let deferred = index + .bm25_search( + tokens.clone(), + params.clone(), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + assert!(!index.has_resident_document_projections()); + + for partition in &index.partitions { + partition.docs.modern().unwrap().prewarm().await.unwrap(); + } + assert!(index.has_resident_document_projections()); + + let resident = index + .bm25_search( + tokens, + params, + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + assert_eq!(resident, deferred); + assert_eq!(resident.0.len(), 2); + assert!(resident.0.contains(&100)); + assert!(resident.0.contains(&200)); + } + async fn search_test_impact_partition( partition: &InvertedPartition, tokens: &Tokens, From 1477b3c30f56324eb44689cd57475ed5d8aecbeb Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 21 Jul 2026 16:23:20 +0800 Subject: [PATCH 05/10] perf(fts): make prewarm query-ready --- docs/src/format/index/scalar/fts.md | 9 + .../src/scalar/inverted/documents.rs | 183 ++++++++- rust/lance-index/src/scalar/inverted/index.rs | 363 +++++++++++++++--- 3 files changed, 506 insertions(+), 49 deletions(-) diff --git a/docs/src/format/index/scalar/fts.md b/docs/src/format/index/scalar/fts.md index d5c75158011..563d7ebf994 100644 --- a/docs/src/format/index/scalar/fts.md +++ b/docs/src/format/index/scalar/fts.md @@ -34,6 +34,15 @@ An FTS index may contain multiple partitions. Each partition has its own set of | `_rowid` | UInt64 | false | Document row ID | | `_num_tokens` | UInt32 | false | Number of tokens in the document | +Partitioned `docs.lance` files may include the optional schema metadata key +`total_tokens`. Its decimal `UInt64` value is the sum of `_num_tokens` in that +file. `_num_tokens` remains the canonical per-document data. Readers use the +metadata value to construct exact corpus statistics without scanning the column; +when the key is absent, they compute the sum from `_num_tokens`. Writers produce +the key from the same document table in the same file commit. A present value +that cannot be parsed, or that differs when `_num_tokens` is subsequently +loaded, is file corruption. + ### FTS List File Schema | Column | Type | Nullable | Description | diff --git a/rust/lance-index/src/scalar/inverted/documents.rs b/rust/lance-index/src/scalar/inverted/documents.rs index 95da7dc947c..d39f6a1c688 100644 --- a/rust/lance-index/src/scalar/inverted/documents.rs +++ b/rust/lance-index/src/scalar/inverted/documents.rs @@ -147,6 +147,10 @@ impl DocLengths { ) } + fn scoring_ready(&self) -> bool { + !self.quantized_scoring || self.norms.get().is_some() + } + #[inline] pub(crate) fn scoring(&self, doc_id: DocId) -> u32 { match self.scoring_norms() { @@ -561,6 +565,13 @@ impl PartitionDocumentStore { } } + pub(crate) fn query_ready(&self) -> bool { + match self { + Self::Legacy(_) => true, + Self::Modern(docs) => docs.query_ready(), + } + } + pub(crate) async fn load_build_docset(&self) -> Result { match self { Self::Legacy(docs) => Ok((**docs).clone()), @@ -649,6 +660,18 @@ impl PartitionDocuments { self.projection.initialized() } + pub(crate) fn query_ready(&self) -> bool { + self.prewarm_complete.initialized() + && self + .lengths + .get() + .is_some_and(|lengths| lengths.scoring_ready()) + && self + .projection + .get() + .is_some_and(|projection| projection.doc_ids_by_address.initialized()) + } + async fn reader(&self) -> Result> { self.store.open_index_file(&self.path).await } @@ -863,6 +886,31 @@ impl PartitionDocuments { } } + /// Estimated Arrow payload retained while resolving these final candidates. + /// The estimate is used to cap cross-partition read concurrency; a single + /// oversized partition is still allowed to make progress. + pub(crate) fn estimated_address_read_bytes(&self, doc_ids: &[DocId]) -> usize { + if doc_ids.is_empty() || self.projection.initialized() { + return 0; + } + if self.remapper.is_some() { + return self.num_docs.saturating_mul(std::mem::size_of::()); + } + + // Avoid repeating the sort/dedup performed by `resolve_addresses`. + // Treating every requested DocId as distinct and disjoint can only + // overestimate the payload, which makes the concurrency bound safer. + let point_bytes = doc_ids.len().saturating_mul(std::mem::size_of::()); + let use_bulk = doc_ids.len() > MAX_POINT_ADDRESS_RANGES + || point_bytes > MAX_POINT_ADDRESS_BYTES + || doc_ids.len() >= self.num_docs; + if use_bulk { + self.num_docs.saturating_mul(std::mem::size_of::()) + } else { + point_bytes + } + } + /// Materialize the build-side table for rewrite/update operations. pub(crate) async fn load_build_docset(&self) -> Result { DocSet::load(self.reader().await?, false, self.remapper.clone()).await @@ -913,7 +961,12 @@ impl PartitionDocuments { let _ = self.projection.set(projection); } - self.lengths().await?; + let lengths = self.lengths().await?; + spawn_cpu(move || { + let _ = lengths.scoring_norms(); + Result::Ok(()) + }) + .await?; self.projection().await?.doc_ids_by_address().await?; Result::Ok(()) }) @@ -994,6 +1047,7 @@ fn corrupt_docs(path: &str, message: impl Into) -> Error { mod tests { use std::collections::HashMap; use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; use arrow_array::{ArrayRef, RecordBatch}; use arrow_schema::{DataType, Field, Schema}; @@ -1003,6 +1057,7 @@ mod tests { use lance_io::object_store::ObjectStore; use lance_select::RowAddrTreeMap; use roaring::RoaringTreemap; + use tokio::sync::Notify; use crate::scalar::lance_format::LanceIndexStore; use crate::scalar::{IndexFile, IndexWriter}; @@ -1031,9 +1086,32 @@ mod tests { } } + const PAUSE_ONCE: usize = 1; + const FAIL_ONCE: usize = 2; + + #[derive(Debug, Default)] + struct ReadFault { + action: AtomicUsize, + started: Notify, + } + + impl ReadFault { + async fn apply(&self) -> Result<()> { + match self.action.swap(0, Ordering::AcqRel) { + PAUSE_ONCE => { + self.started.notify_one(); + std::future::pending::>().await + } + FAIL_ONCE => Err(Error::io("injected document read failure")), + _ => Ok(()), + } + } + } + struct CountingReader { inner: Arc, counts: Arc, + fault: Option>, } #[async_trait] @@ -1053,6 +1131,9 @@ mod tests { ) -> Result { self.counts.range_calls.fetch_add(1, Ordering::Relaxed); self.counts.record(range.len(), projection); + if let Some(fault) = &self.fault { + fault.apply().await?; + } self.inner.read_range(range, projection).await } @@ -1064,6 +1145,9 @@ mod tests { self.counts.ranges_calls.fetch_add(1, Ordering::Relaxed); self.counts .record(ranges.iter().map(Range::len).sum(), projection); + if let Some(fault) = &self.fault { + fault.apply().await?; + } self.inner.read_ranges(ranges, projection).await } @@ -1089,6 +1173,7 @@ mod tests { inner: Arc, target: String, counts: Arc, + fault: Option>, } impl DeepSizeOf for CountingStore { @@ -1108,6 +1193,7 @@ mod tests { inner: self.inner.clone(), target: self.target.clone(), counts: self.counts.clone(), + fault: self.fault.clone(), }) } @@ -1130,6 +1216,7 @@ mod tests { Ok(Arc::new(CountingReader { inner: reader, counts: self.counts.clone(), + fault: self.fault.clone(), })) } else { Ok(reader) @@ -1141,6 +1228,7 @@ mod tests { inner: self.inner.with_io_priority(io_priority), target: self.target.clone(), counts: self.counts.clone(), + fault: self.fault.clone(), }) } @@ -1239,11 +1327,32 @@ mod tests { inner, target: target.to_owned(), counts: counts.clone(), + fault: None, }), counts, ) } + fn faulting_store( + inner: Arc, + target: &str, + action: usize, + ) -> (Arc, Arc) { + let fault = Arc::new(ReadFault { + action: AtomicUsize::new(action), + started: Notify::new(), + }); + ( + Arc::new(CountingStore { + inner, + target: target.to_owned(), + counts: Arc::new(DocumentReadCounts::default()), + fault: Some(fault.clone()), + }), + fault, + ) + } + #[derive(Debug)] struct TestRemapper { mapping: HashMap>, @@ -1535,6 +1644,7 @@ mod tests { assert!(documents.lengths_loaded()); assert!(documents.projection_loaded()); + assert!(documents.query_ready()); assert!(matches!( first_lookup.as_ref(), AddressDocIdLookup::Identity @@ -1550,6 +1660,72 @@ mod tests { assert_eq!(counts.address_rows.load(Ordering::Relaxed), 3); } + #[tokio::test] + async fn prewarm_materializes_quantized_norms_before_becoming_query_ready() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 300, 5]), + Some("307"), + ) + .await; + let reader = store.open_index_file(path).await.unwrap(); + let documents = + PartitionDocuments::try_new(store, path.to_owned(), reader.as_ref(), None, true) + .unwrap(); + + assert!(!documents.query_ready()); + documents.prewarm().await.unwrap(); + let lengths = documents.lengths().await.unwrap(); + assert!(lengths.scoring_ready()); + assert_eq!(lengths.scoring_norms().unwrap().len(), 3); + assert!(documents.query_ready()); + } + + #[tokio::test] + async fn cancelled_or_failed_prewarm_can_retry_without_partial_publication() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + Some("10"), + ) + .await; + + let (pausing, pause) = faulting_store(store.clone(), path, PAUSE_ONCE); + let documents = Arc::new(open_documents(pausing, path, None).await.unwrap()); + let task = tokio::spawn({ + let documents = documents.clone(); + async move { documents.prewarm().await } + }); + tokio::time::timeout(Duration::from_secs(5), pause.started.notified()) + .await + .expect("prewarm should reach the injected pending read"); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); + assert!(!documents.query_ready()); + documents.prewarm().await.unwrap(); + assert!(documents.query_ready()); + + let (failing, _fault) = faulting_store(store, path, FAIL_ONCE); + let documents = open_documents(failing, path, None).await.unwrap(); + let error = documents.prewarm().await.unwrap_err(); + assert!(error.to_string().contains("injected document read failure")); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); + assert!(!documents.query_ready()); + documents.prewarm().await.unwrap(); + assert!(documents.query_ready()); + } + #[tokio::test] async fn prewarm_reuses_an_already_loaded_document_column() { let (_directory, store) = test_store(); @@ -1630,6 +1806,7 @@ mod tests { assert_eq!(counts.rows.load(Ordering::Relaxed), 0); let point_ids = [DocId::new(5), DocId::new(6), DocId::new(10), DocId::new(5)]; + assert_eq!(documents.estimated_address_read_bytes(&point_ids), 32); assert_eq!( documents.resolve_addresses(&point_ids).await.unwrap(), vec![1005, 1006, 1010, 1005] @@ -1640,6 +1817,10 @@ mod tests { let bulk_ids = (0..=512).step_by(2).map(DocId::new).collect::>(); assert_eq!(bulk_ids.len(), MAX_POINT_ADDRESS_RANGES + 1); + assert_eq!( + documents.estimated_address_read_bytes(&bulk_ids), + num_docs as usize * std::mem::size_of::() + ); let resolved = documents.resolve_addresses(&bulk_ids).await.unwrap(); assert_eq!(resolved.first(), Some(&1000)); assert_eq!(resolved.last(), Some(&1512)); diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index d583568226a..2e3b96b4a83 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -18,6 +18,7 @@ use std::{ use crate::metrics::NoOpMetricsCollector; use crate::prefilter::NoFilter; use crate::scalar::registry::{TrainingCriteria, TrainingOrdering}; +use crate::vector::graph::OrderedFloat; use arrow::array::{FixedSizeListBuilder, Float32Builder, Int32Builder}; use arrow::datatypes::{self, Float32Type, Int32Type, UInt64Type}; use arrow::{ @@ -50,7 +51,10 @@ use lance_core::{Error, ROW_ID, ROW_ID_FIELD, Result}; use lance_select::{RowAddrMask, RowAddrTreeMap}; use roaring::RoaringBitmap; use std::sync::LazyLock; -use tokio::{sync::OnceCell, task::spawn_blocking}; +use tokio::{ + sync::{Mutex, OnceCell}, + task::spawn_blocking, +}; use tracing::{info, instrument, warn}; use super::documents::{ @@ -296,6 +300,80 @@ struct ModernSearchRequest<'a> { limit: usize, } +/// Typed identity for one modern candidate after partition-local scoring. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct PartitionDocId { + partition_ordinal: u32, + doc_id: DocId, +} + +impl PartitionDocId { + fn try_new(partition_ordinal: usize, doc_id: DocId) -> Result { + Ok(Self { + partition_ordinal: u32::try_from(partition_ordinal).map_err(|_| { + Error::index(format!( + "FTS partition ordinal {partition_ordinal} exceeds candidate identity capacity" + )) + })?, + doc_id, + }) + } + + fn partition_ordinal(self) -> usize { + self.partition_ordinal as usize + } +} + +#[derive(Debug, Clone)] +struct ScoredPartitionDoc { + document: PartitionDocId, + score: OrderedFloat, +} + +impl ScoredPartitionDoc { + fn new(document: PartitionDocId, score: f32) -> Self { + Self { + document, + score: OrderedFloat(score), + } + } +} + +impl PartialEq for ScoredPartitionDoc { + fn eq(&self, other: &Self) -> bool { + self.score == other.score + } +} + +impl Eq for ScoredPartitionDoc {} + +impl PartialOrd for ScoredPartitionDoc { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for ScoredPartitionDoc { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.score.cmp(&other.score) + } +} + +const MAX_CONCURRENT_ADDRESS_READ_BYTES: usize = 64 * 1024 * 1024; + +fn address_read_concurrency(io_parallelism: usize, largest_read_bytes: usize) -> usize { + let io_parallelism = io_parallelism.max(1); + if largest_read_bytes == 0 { + return io_parallelism; + } + io_parallelism.min( + MAX_CONCURRENT_ADDRESS_READ_BYTES + .checked_div(largest_read_bytes) + .unwrap_or(0) + .max(1), + ) +} + fn push_scored_key( candidates: &mut BinaryHeap>, limit: usize, @@ -313,6 +391,23 @@ fn push_scored_key( } } +fn push_scored_partition_doc( + candidates: &mut BinaryHeap>, + limit: usize, + document: PartitionDocId, + score: f32, +) { + if candidates.len() < limit { + candidates.push(Reverse(ScoredPartitionDoc::new(document, score))); + } else if candidates + .peek() + .is_some_and(|candidate| candidate.0.score.0 < score) + { + candidates.pop(); + candidates.push(Reverse(ScoredPartitionDoc::new(document, score))); + } +} + fn rescore_partition_candidates( partition: PartitionCandidates, scorer: &MemBM25Scorer, @@ -578,6 +673,18 @@ pub(super) fn parse_format_version_from_metadata( } } +#[derive(Debug, Default)] +struct InvertedPrewarmState { + query_ready: bool, + positions_ready: bool, +} + +impl InvertedPrewarmState { + fn satisfies(&self, with_position: bool) -> bool { + self.query_ready && (!with_position || self.positions_ready) + } +} + #[derive(Clone)] pub struct InvertedIndex { params: InvertedIndexParams, @@ -587,6 +694,7 @@ pub struct InvertedIndex { format_version: InvertedListFormatVersion, pub(crate) partitions: Vec>, corpus_stats: Arc>, + prewarm_state: Arc>, // Fragments which are contained in the index, but no longer in the dataset. // These should be pruned at search time since we don't prune them at update time. deleted_fragments: RoaringBitmap, @@ -1135,7 +1243,7 @@ impl InvertedIndex { async fn bm25_search_modern_candidates( &self, request: ModernSearchRequest<'_>, - ) -> Result>> { + ) -> Result>> { let ModernSearchRequest { tokens, params, @@ -1255,8 +1363,12 @@ impl InvertedIndex { let mut parts = stream::iter(parts).buffer_unordered(get_num_compute_intensive_cpus()); while let Some((partition_ordinal, partition)) = parts.try_next().await? { for (doc_id, score) in rescore_partition_candidates(partition, scorer, &mut idf_cache) { - let key = ((partition_ordinal as u64) << 32) | u64::from(doc_id.get()); - push_scored_key(&mut ranked, limit, key, score); + push_scored_partition_doc( + &mut ranked, + limit, + PartitionDocId::try_new(partition_ordinal, doc_id)?, + score, + ); } } @@ -1265,13 +1377,13 @@ impl InvertedIndex { fn resolve_resident_modern_candidates( &self, - ranked: Vec>, + ranked: Vec>, ) -> Result<(Vec, Vec)> { let mut addresses = Vec::with_capacity(ranked.len()); let mut scores = Vec::with_capacity(ranked.len()); for Reverse(candidate) in ranked { - let partition_ordinal = (candidate.row_id >> 32) as usize; - let doc_id = DocId::new(candidate.row_id as u32); + let partition_ordinal = candidate.document.partition_ordinal(); + let doc_id = candidate.document.doc_id; let documents = self .partitions .get(partition_ordinal) @@ -1296,19 +1408,20 @@ impl InvertedIndex { async fn resolve_deferred_modern_candidates( &self, - ranked: Vec>, + ranked: Vec>, ) -> Result<(Vec, Vec)> { let mut addresses = vec![0_u64; ranked.len()]; let mut by_partition = BTreeMap::>::new(); for (rank, Reverse(candidate)) in ranked.iter().enumerate() { - let partition_ordinal = (candidate.row_id >> 32) as usize; - let doc_id = DocId::new(candidate.row_id as u32); + let partition_ordinal = candidate.document.partition_ordinal(); + let doc_id = candidate.document.doc_id; by_partition .entry(partition_ordinal) .or_default() .push((rank, doc_id)); } let mut address_reads = Vec::with_capacity(by_partition.len()); + let mut largest_read_bytes = 0; for (partition_ordinal, entries) in by_partition { let documents = self .partitions @@ -1320,17 +1433,19 @@ impl InvertedIndex { "deferred FTS candidates reference missing modern partition ordinal {partition_ordinal}" )) })?; + let doc_ids = entries + .iter() + .map(|(_, doc_id)| *doc_id) + .collect::>(); + largest_read_bytes = + largest_read_bytes.max(documents.estimated_address_read_bytes(&doc_ids)); address_reads.push(async move { - let doc_ids = entries - .iter() - .map(|(_, doc_id)| *doc_id) - .collect::>(); let resolved = documents.resolve_addresses(&doc_ids).await?; Result::Ok((entries, resolved)) }); } - let mut address_reads = - stream::iter(address_reads).buffer_unordered(self.store.io_parallelism().max(1)); + let concurrency = address_read_concurrency(self.store.io_parallelism(), largest_read_bytes); + let mut address_reads = stream::iter(address_reads).buffer_unordered(concurrency); while let Some((entries, resolved)) = address_reads.try_next().await? { for ((rank, _), address) in entries.into_iter().zip(resolved) { addresses[rank] = address; @@ -1404,6 +1519,7 @@ impl InvertedIndex { token_set_format: TokenSetFormat::Arrow, })], corpus_stats: Arc::new(OnceCell::new()), + prewarm_state: Arc::new(Mutex::new(InvertedPrewarmState::default())), deleted_fragments: RoaringBitmap::new(), })) } @@ -1525,6 +1641,7 @@ impl InvertedIndex { format_version, partitions, corpus_stats: Arc::new(OnceCell::new()), + prewarm_state: Arc::new(Mutex::new(InvertedPrewarmState::default())), deleted_fragments, })) } @@ -1743,7 +1860,17 @@ fn prewarm_chunk_ranges( impl InvertedIndex { pub async fn prewarm_with_options(&self, options: &FtsPrewarmOptions) -> Result<()> { - let with_position = options.with_position; + let mut state = self.prewarm_state.lock().await; + if state.satisfies(options.with_position) { + return Ok(()); + } + self.prewarm_query_state(options.with_position).await?; + state.query_ready = true; + state.positions_ready |= options.with_position; + Ok(()) + } + + async fn prewarm_query_state(&self, with_position: bool) -> Result<()> { let chunk_concurrency = self.store.io_parallelism().max(1); let prewarm_started = Instant::now(); info!( @@ -1795,8 +1922,20 @@ impl InvertedIndex { "fts partition prewarm finished" ); } + self.aggregate_corpus_stats().await?; + let query_ready = self.partitions.iter().all(|partition| { + partition.docs.query_ready() + && partition.inverted_list.modern_posting_validation_ready() + }); + if !query_ready { + return Err(Error::internal( + "FTS prewarm completed without publishing a query-ready document and posting state" + .to_owned(), + )); + } info!( partition_count = self.partitions.len(), + query_ready, elapsed_ms = prewarm_started.elapsed().as_millis() as u64, "fts index prewarm finished" ); @@ -3094,7 +3233,7 @@ pub struct PostingListReader { /// Modern postings contain dense DocIds into the partition document table. /// Cache successful boundary validation per immutable token so repeated /// queries do not decode the final posting block again. - modern_doc_id_validations: Option>, + modern_doc_id_validations: Option]>>, modern_num_docs: Option, index_cache: WeakLanceCache, @@ -3176,7 +3315,7 @@ impl DeepSizeOf for PostingListReader { .map(|validations| { validations .len() - .saturating_mul(std::mem::size_of::()) + .saturating_mul(std::mem::size_of::>()) }) .unwrap_or(0); metadata_size + self.grouping.deep_size_of_children(context) + validation_size @@ -3215,7 +3354,7 @@ impl PostingListReader { let grouping = PostingGrouping::for_reader(is_legacy_layout, reader.num_rows()); let modern_doc_id_validations = (!is_legacy_layout).then(|| { (0..reader.num_rows()) - .map(|_| AtomicBool::new(false)) + .map(|_| OnceCell::new()) .collect::>() .into() }); @@ -3516,7 +3655,8 @@ impl PostingListReader { .clone(), }; - self.ensure_modern_posting_validated(token_id, &posting)?; + self.ensure_modern_posting_validated(token_id, &posting) + .await?; if is_phrase_query && !posting.has_position() { // hit the cache and when the cache was populated, the positions column was not loaded @@ -3527,7 +3667,11 @@ impl PostingListReader { Ok(posting) } - fn ensure_modern_posting_validated(&self, token_id: u32, posting: &PostingList) -> Result<()> { + async fn ensure_modern_posting_validated( + &self, + token_id: u32, + posting: &PostingList, + ) -> Result<()> { let (Some(validations), Some(num_docs)) = (&self.modern_doc_id_validations, self.modern_num_docs) else { @@ -3539,12 +3683,42 @@ impl PostingListReader { validations.len() )) })?; - if validation.load(Ordering::Acquire) { + validation + .get_or_try_init(|| async { + Self::validate_modern_posting(token_id, posting, num_docs) + }) + .await + .map(|_| ()) + } + + fn validate_modern_posting( + token_id: u32, + posting: &PostingList, + num_docs: usize, + ) -> Result<()> { + validate_modern_posting_doc_ids(posting, &format!("token id {token_id}"), num_docs) + } + + async fn publish_modern_posting_validated(&self, token_id: u32) -> Result<()> { + let Some(validations) = &self.modern_doc_id_validations else { return Ok(()); - } - validate_modern_posting_doc_ids(posting, &format!("token id {token_id}"), num_docs)?; - validation.store(true, Ordering::Release); - Ok(()) + }; + let validation = validations.get(token_id as usize).ok_or_else(|| { + Error::index(format!( + "modern FTS token id {token_id} is outside validation state [0, {})", + validations.len() + )) + })?; + validation + .get_or_try_init(|| async { Result::Ok(()) }) + .await + .map(|_| ()) + } + + fn modern_posting_validation_ready(&self) -> bool { + self.modern_doc_id_validations + .as_ref() + .is_none_or(|validations| validations.iter().all(|state| state.get().is_some())) } /// Map a token id to its cache group's row range `[start, end)`, or `None` @@ -3834,6 +4008,7 @@ impl PostingListReader { let posting_tail_codec = state.posting_tail_codec; let block_size = state.block_size; let positions_layout = state.positions_layout; + let num_docs = self.modern_num_docs; let posting_lists = spawn_blocking(move || { let ctx = PrewarmBuildCtx { max_scores: max_scores.as_deref().map(|v| v.as_slice()), @@ -3848,7 +4023,13 @@ impl PostingListReader { offsets: chunk_offsets.as_deref(), end_row: chunk_end_row, }; - Self::build_prewarm_posting_lists_chunk(chunk_batch, chunk, &ctx) + let posting_lists = Self::build_prewarm_posting_lists_chunk(chunk_batch, chunk, &ctx)?; + if let Some(num_docs) = num_docs { + for (token_id, posting) in &posting_lists { + Self::validate_modern_posting(*token_id, posting, num_docs)?; + } + } + Result::Ok(posting_lists) }) .await .map_err(|err| { @@ -3856,6 +4037,9 @@ impl PostingListReader { "Failed to build prewarm posting lists in blocking task: {err}" )) })??; + for (token_id, _) in &posting_lists { + self.publish_modern_posting_validated(*token_id).await?; + } // The chunk yields its token range as contiguous ascending ids from // `tok_start`; the group publish path relies on this to index the lists. debug_assert_eq!(posting_lists.len(), chunk_token_count); @@ -3885,8 +4069,25 @@ impl PostingListReader { let ranges = grouping.ranges_for_chunk(tok_start, tok_end, token_count); let posting_tail_codec = self.posting_tail_codec; let block_size = self.block_size; + let num_docs = self.modern_num_docs; + let (chunk_max_scores, chunk_lengths) = match &self.metadata { + PostingMetadata::V2 { metadata } => { + let loaded = metadata.get().ok_or_else(|| { + Error::internal("packed prewarm requires loaded posting metadata".to_owned()) + })?; + ( + loaded.max_scores[tok_start..tok_end].to_vec(), + loaded.lengths[tok_start..tok_end].to_vec(), + ) + } + PostingMetadata::LegacyV1 { .. } => { + return Err(Error::internal( + "packed prewarm is not supported for legacy posting metadata".to_owned(), + )); + } + }; - spawn_blocking(move || { + let groups = spawn_blocking(move || { let mut groups = Vec::with_capacity(ranges.len()); for (start, end) in ranges { let start_usize = start as usize; @@ -3894,15 +4095,29 @@ impl PostingListReader { let local_start = start_usize - tok_start; let group_len = end_usize - start_usize; let group_batch = chunk_batch.slice(local_start, group_len).shrink_to_fit()?; - groups.push(( - start, - end, - PostingListGroup::new_packed_with_block_size( - group_batch, - posting_tail_codec, - block_size, - )?, - )); + let group = PostingListGroup::new_packed_with_block_size( + group_batch, + posting_tail_codec, + block_size, + )?; + if let Some(num_docs) = num_docs { + for token_id in start..end { + let chunk_slot = token_id as usize - tok_start; + let posting = group + .posting_list( + (token_id - start) as usize, + Some(chunk_max_scores[chunk_slot]), + Some(chunk_lengths[chunk_slot]), + )? + .ok_or_else(|| { + Error::index(format!( + "token {token_id} is missing from prewarm posting group [{start}, {end})" + )) + })?; + Self::validate_modern_posting(token_id, &posting, num_docs)?; + } + } + groups.push((start, end, group)); } Result::Ok(groups) }) @@ -3911,7 +4126,13 @@ impl PostingListReader { Error::internal(format!( "Failed to build packed prewarm posting groups in blocking task: {err}" )) - })? + })??; + for (start, end, _) in &groups { + for token_id in *start..*end { + self.publish_modern_posting_validated(token_id).await?; + } + } + Ok(groups) } /// Strip positions into their own per-token cache entries (the posting cache @@ -7772,6 +7993,17 @@ mod tests { use super::*; + #[test] + fn address_read_concurrency_respects_payload_budget() { + assert_eq!(address_read_concurrency(64, 0), 64); + assert_eq!(address_read_concurrency(64, 8 * 1024 * 1024), 8); + assert_eq!(address_read_concurrency(64, 16 * 1024 * 1024), 4); + assert_eq!( + address_read_concurrency(64, 2 * MAX_CONCURRENT_ADDRESS_READ_BYTES), + 1 + ); + } + #[derive(Debug)] struct MetadataAccessDeniedStore { inner: Arc, @@ -10561,10 +10793,41 @@ mod tests { .unwrap(); assert!(!index.has_resident_document_projections()); - for partition in &index.partitions { - partition.docs.modern().unwrap().prewarm().await.unwrap(); - } + index.partitions[0] + .docs + .modern() + .unwrap() + .prewarm() + .await + .unwrap(); + assert!(index.partitions[0].docs.query_ready()); + assert!(!index.has_resident_document_projections()); + let partially_resident = index + .bm25_search( + tokens.clone(), + params.clone(), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + assert_eq!(partially_resident, deferred); + + let prewarm_options = FtsPrewarmOptions::default(); + futures::future::join_all((0..8).map(|_| index.prewarm_with_options(&prewarm_options))) + .await + .into_iter() + .collect::>>() + .unwrap(); assert!(index.has_resident_document_projections()); + assert!(index.corpus_stats.initialized()); + assert!(index.partitions.iter().all(|partition| { + partition.docs.query_ready() + && partition.inverted_list.modern_posting_validation_ready() + })); + assert!(index.prewarm_state.lock().await.satisfies(false)); let resident = index .bm25_search( @@ -12480,7 +12743,7 @@ mod tests { .modern_doc_id_validations .as_ref() .expect("modern readers have per-token validation state")[0]; - assert!(!validation.load(Ordering::Acquire)); + assert!(validation.get().is_none()); let mut corrupt_builder = PostingListBuilder::new(false); corrupt_builder.add(1, PositionRecorder::Count(1)); @@ -12488,25 +12751,26 @@ mod tests { let corrupt_posting = PostingList::from_batch(&corrupt_batch, Some(1.0), Some(1)).unwrap(); let error = posting_reader .ensure_modern_posting_validated(0, &corrupt_posting) + .await .unwrap_err(); assert!(matches!(error, Error::Index { .. })); assert!(error.to_string().contains("DocId 1")); assert!(error.to_string().contains("[0, 1)")); - assert!(!validation.load(Ordering::Acquire)); + assert!(validation.get().is_none()); let first = posting_reader .posting_list(0, false, &NoOpMetricsCollector) .await .unwrap(); assert_eq!(posting_entries(&first), vec![(0, 1)]); - assert!(validation.load(Ordering::Acquire)); + assert!(validation.get().is_some()); let second = posting_reader .posting_list(0, false, &NoOpMetricsCollector) .await .unwrap(); assert_eq!(posting_entries(&second), vec![(0, 1)]); - assert!(validation.load(Ordering::Acquire)); + assert!(validation.get().is_some()); } /// Runtime synthetic grouping must return correct posting lists for every @@ -12533,7 +12797,8 @@ mod tests { let reader = store.open_index_file(&posting_file_path(0)).await.unwrap(); let cache = LanceCache::no_cache(); - let posting_reader = PostingListReader::try_new(reader, &cache).await.unwrap(); + let mut posting_reader = PostingListReader::try_new(reader, &cache).await.unwrap(); + posting_reader.modern_num_docs = Some(num_tokens as usize); assert!( matches!( &posting_reader.grouping, @@ -12586,7 +12851,8 @@ mod tests { // A real (strong) cache must outlive the reader's weak handle so the // prewarmed entries are still resolvable below. let cache = LanceCache::with_capacity(1 << 20); - let posting_reader = PostingListReader::try_new(reader, &cache).await.unwrap(); + let mut posting_reader = PostingListReader::try_new(reader, &cache).await.unwrap(); + posting_reader.modern_num_docs = Some(num_tokens as usize); assert!( matches!( &posting_reader.grouping, @@ -12599,6 +12865,7 @@ mod tests { .prewarm_posting_lists(false, 2) .await .unwrap(); + assert!(posting_reader.modern_posting_validation_ready()); for token in 0..num_tokens { let (start, end) = posting_reader.group_range_for_token(token).unwrap(); From c69853ff76a37caa6ed1d1f06cd8d1dc1e93e704 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 21 Jul 2026 17:10:12 +0800 Subject: [PATCH 06/10] perf(fts): keep validated posting reads synchronous --- rust/lance-index/src/scalar/inverted/index.rs | 24 +++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 2e3b96b4a83..672dab9805d 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -3655,8 +3655,10 @@ impl PostingListReader { .clone(), }; - self.ensure_modern_posting_validated(token_id, &posting) - .await?; + if !self.modern_posting_is_validated(token_id)? { + self.ensure_modern_posting_validated(token_id, &posting) + .await?; + } if is_phrase_query && !posting.has_position() { // hit the cache and when the cache was populated, the positions column was not loaded @@ -3691,6 +3693,21 @@ impl PostingListReader { .map(|_| ()) } + #[inline] + fn modern_posting_is_validated(&self, token_id: u32) -> Result { + let (Some(validations), Some(_)) = (&self.modern_doc_id_validations, self.modern_num_docs) + else { + return Ok(true); + }; + let validation = validations.get(token_id as usize).ok_or_else(|| { + Error::index(format!( + "modern FTS token id {token_id} is outside validation state [0, {})", + validations.len() + )) + })?; + Ok(validation.get().is_some()) + } + fn validate_modern_posting( token_id: u32, posting: &PostingList, @@ -12744,6 +12761,7 @@ mod tests { .as_ref() .expect("modern readers have per-token validation state")[0]; assert!(validation.get().is_none()); + assert!(!posting_reader.modern_posting_is_validated(0).unwrap()); let mut corrupt_builder = PostingListBuilder::new(false); corrupt_builder.add(1, PositionRecorder::Count(1)); @@ -12757,6 +12775,7 @@ mod tests { assert!(error.to_string().contains("DocId 1")); assert!(error.to_string().contains("[0, 1)")); assert!(validation.get().is_none()); + assert!(!posting_reader.modern_posting_is_validated(0).unwrap()); let first = posting_reader .posting_list(0, false, &NoOpMetricsCollector) @@ -12764,6 +12783,7 @@ mod tests { .unwrap(); assert_eq!(posting_entries(&first), vec![(0, 1)]); assert!(validation.get().is_some()); + assert!(posting_reader.modern_posting_is_validated(0).unwrap()); let second = posting_reader .posting_list(0, false, &NoOpMetricsCollector) From 711028537fa470a8c055ba352a52445a93ee0635 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 21 Jul 2026 17:38:13 +0800 Subject: [PATCH 07/10] perf(fts): publish prewarm fast paths --- rust/lance-index/src/scalar/inverted/index.rs | 51 +++++++++++++++---- 1 file changed, 41 insertions(+), 10 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 672dab9805d..a487ffcc047 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -695,6 +695,8 @@ pub struct InvertedIndex { pub(crate) partitions: Vec>, corpus_stats: Arc>, prewarm_state: Arc>, + /// One-way publication: document projections never leave memory after loading. + document_projections_resident: Arc, // Fragments which are contained in the index, but no longer in the dataset. // These should be pruned at search time since we don't prune them at update time. deleted_fragments: RoaringBitmap, @@ -1216,12 +1218,20 @@ impl InvertedIndex { } fn has_resident_document_projections(&self) -> bool { - self.partitions.iter().all(|partition| { + if self.document_projections_resident.load(Ordering::Acquire) { + return true; + } + let resident = self.partitions.iter().all(|partition| { partition .docs .modern() .is_some_and(|documents| documents.projection_loaded()) - }) + }); + if resident { + self.document_projections_resident + .store(true, Ordering::Release); + } + resident } async fn bm25_search_modern_resident( @@ -1306,11 +1316,9 @@ impl InvertedIndex { <= u128::from(*FLAT_SEARCH_PERCENT_THRESHOLD) .saturating_mul(documents.len() as u128) }); - let visibility = Arc::new( - documents - .visibility(mask.clone(), materialize_selected) - .await?, - ); + let visibility = documents + .visibility(mask.clone(), materialize_selected) + .await?; if visibility.is_empty() { return Result::Ok((partition_ordinal, PartitionCandidates::empty())); } @@ -1336,7 +1344,7 @@ impl InvertedIndex { let candidates = spawn_cpu(move || { part_for_wand.bm25_search_modern( lengths.as_ref(), - visibility.as_ref(), + &visibility, params.as_ref(), operator, postings, @@ -1520,6 +1528,7 @@ impl InvertedIndex { })], corpus_stats: Arc::new(OnceCell::new()), prewarm_state: Arc::new(Mutex::new(InvertedPrewarmState::default())), + document_projections_resident: Arc::new(AtomicBool::new(false)), deleted_fragments: RoaringBitmap::new(), })) } @@ -1642,6 +1651,7 @@ impl InvertedIndex { partitions, corpus_stats: Arc::new(OnceCell::new()), prewarm_state: Arc::new(Mutex::new(InvertedPrewarmState::default())), + document_projections_resident: Arc::new(AtomicBool::new(false)), deleted_fragments, })) } @@ -3234,6 +3244,8 @@ pub struct PostingListReader { /// Cache successful boundary validation per immutable token so repeated /// queries do not decode the final posting block again. modern_doc_id_validations: Option]>>, + /// Skips per-token readiness checks once the whole immutable table is validated. + modern_postings_validated: AtomicBool, modern_num_docs: Option, index_cache: WeakLanceCache, @@ -3369,6 +3381,7 @@ impl PostingListReader { positions_layout, grouping, modern_doc_id_validations, + modern_postings_validated: AtomicBool::new(false), modern_num_docs: None, index_cache: WeakLanceCache::from(index_cache), }) @@ -3695,6 +3708,9 @@ impl PostingListReader { #[inline] fn modern_posting_is_validated(&self, token_id: u32) -> Result { + if self.modern_postings_validated.load(Ordering::Acquire) { + return Ok(true); + } let (Some(validations), Some(_)) = (&self.modern_doc_id_validations, self.modern_num_docs) else { return Ok(true); @@ -3733,9 +3749,18 @@ impl PostingListReader { } fn modern_posting_validation_ready(&self) -> bool { - self.modern_doc_id_validations + if self.modern_postings_validated.load(Ordering::Acquire) { + return true; + } + let ready = self + .modern_doc_id_validations .as_ref() - .is_none_or(|validations| validations.iter().all(|state| state.get().is_some())) + .is_none_or(|validations| validations.iter().all(|state| state.get().is_some())); + if ready { + self.modern_postings_validated + .store(true, Ordering::Release); + } + ready } /// Map a token id to its cache group's row range `[start, end)`, or `None` @@ -10839,6 +10864,7 @@ mod tests { .collect::>>() .unwrap(); assert!(index.has_resident_document_projections()); + assert!(index.document_projections_resident.load(Ordering::Acquire)); assert!(index.corpus_stats.initialized()); assert!(index.partitions.iter().all(|partition| { partition.docs.query_ready() @@ -12886,6 +12912,11 @@ mod tests { .await .unwrap(); assert!(posting_reader.modern_posting_validation_ready()); + assert!( + posting_reader + .modern_postings_validated + .load(Ordering::Acquire) + ); for token in 0..num_tokens { let (start, end) = posting_reader.group_range_for_token(token).unwrap(); From 23ff9980932b67285d9af4d363b47be89a71b0a8 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 21 Jul 2026 18:35:53 +0800 Subject: [PATCH 08/10] perf(fts): bypass resident document futures --- .../src/scalar/inverted/documents.rs | 41 ++++++++++++++++--- rust/lance-index/src/scalar/inverted/index.rs | 25 ++++++++--- 2 files changed, 55 insertions(+), 11 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/documents.rs b/rust/lance-index/src/scalar/inverted/documents.rs index d39f6a1c688..02e8617cdd2 100644 --- a/rust/lance-index/src/scalar/inverted/documents.rs +++ b/rust/lance-index/src/scalar/inverted/documents.rs @@ -736,6 +736,11 @@ impl PartitionDocuments { .cloned() } + /// Return resident lengths without entering the asynchronous singleflight path. + pub(crate) fn cached_lengths(&self) -> Option> { + self.lengths.get().cloned() + } + pub(crate) async fn projection(&self) -> Result> { self.projection .get_or_try_init(|| async { @@ -752,11 +757,8 @@ impl PartitionDocuments { mask: Arc, materialize_selected: bool, ) -> Result { - if mask.max_len() == Some(0) { - return Ok(DocVisibility::Selected(RoaringBitmap::new())); - } - if mask.is_select_all() && self.remapper.is_none() { - return Ok(DocVisibility::All); + if let Some(visibility) = self.immediate_visibility(mask.clone(), materialize_selected) { + return Ok(visibility); } let projection = self.projection().await?; @@ -770,6 +772,30 @@ impl PartitionDocuments { } } + /// Resolve visibility without I/O or CPU-pool work when all required state is resident. + pub(crate) fn immediate_visibility( + &self, + mask: Arc, + materialize_selected: bool, + ) -> Option { + if mask.max_len() == Some(0) { + return Some(DocVisibility::Selected(RoaringBitmap::new())); + } + if mask.is_select_all() && self.remapper.is_none() { + return Some(DocVisibility::All); + } + + let projection = self.projection.get()?.clone(); + if mask.is_select_all() { + return Some(DocVisibility::Selected(projection.live_doc_ids())); + } + if materialize_selected { + None + } else { + Some(DocVisibility::Filtered { projection, mask }) + } + } + /// Resolve final global top-k DocIds to current row addresses. pub(crate) async fn resolve_addresses(&self, doc_ids: &[DocId]) -> Result> { if doc_ids.is_empty() { @@ -1645,6 +1671,11 @@ mod tests { assert!(documents.lengths_loaded()); assert!(documents.projection_loaded()); assert!(documents.query_ready()); + assert_eq!(documents.cached_lengths().unwrap().total_tokens(), 10); + assert!(matches!( + documents.immediate_visibility(Arc::new(RowAddrMask::all_rows()), false), + Some(DocVisibility::All) + )); assert!(matches!( first_lookup.as_ref(), AddressDocIdLookup::Identity diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index a487ffcc047..ed51282e19c 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -1272,7 +1272,9 @@ impl InvertedIndex { } // Ensures old partitioned files without persisted stats populate their // fallback stats before any synchronous local scorer is constructed. - self.aggregate_corpus_stats().await?; + if self.corpus_stats.get().is_none() { + self.aggregate_corpus_stats().await?; + } let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); let local_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); @@ -1316,13 +1318,22 @@ impl InvertedIndex { <= u128::from(*FLAT_SEARCH_PERCENT_THRESHOLD) .saturating_mul(documents.len() as u128) }); - let visibility = documents - .visibility(mask.clone(), materialize_selected) - .await?; + let visibility = + match documents.immediate_visibility(mask.clone(), materialize_selected) { + Some(visibility) => visibility, + None => { + documents + .visibility(mask.clone(), materialize_selected) + .await? + } + }; if visibility.is_empty() { return Result::Ok((partition_ordinal, PartitionCandidates::empty())); } - let lengths = documents.lengths().await?; + let lengths = match documents.cached_lengths() { + Some(lengths) => lengths, + None => documents.lengths().await?, + }; let max_position = postings .iter() .map(|posting| posting.term_index() as usize) @@ -1943,6 +1954,8 @@ impl InvertedIndex { .to_owned(), )); } + self.document_projections_resident + .store(true, Ordering::Release); info!( partition_count = self.partitions.len(), query_ready, @@ -10863,8 +10876,8 @@ mod tests { .into_iter() .collect::>>() .unwrap(); - assert!(index.has_resident_document_projections()); assert!(index.document_projections_resident.load(Ordering::Acquire)); + assert!(index.has_resident_document_projections()); assert!(index.corpus_stats.initialized()); assert!(index.partitions.iter().all(|partition| { partition.docs.query_ready() From 03e17747a094552838f5251a0439d337af8d8eb1 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 21 Jul 2026 18:51:49 +0800 Subject: [PATCH 09/10] perf(fts): specialize unfiltered document scoring --- .../src/scalar/inverted/documents.rs | 4 ++ rust/lance-index/src/scalar/inverted/index.rs | 33 +++++++---- rust/lance-index/src/scalar/inverted/wand.rs | 58 +++++++++++++++++-- 3 files changed, 80 insertions(+), 15 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/documents.rs b/rust/lance-index/src/scalar/inverted/documents.rs index 02e8617cdd2..c415111e131 100644 --- a/rust/lance-index/src/scalar/inverted/documents.rs +++ b/rust/lance-index/src/scalar/inverted/documents.rs @@ -454,6 +454,10 @@ pub(super) enum DocVisibility { } impl DocVisibility { + pub(crate) fn is_all(&self) -> bool { + matches!(self, Self::All) + } + #[inline] pub(crate) fn selected(&self, doc_id: DocId) -> bool { match self { diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index ed51282e19c..9b7bfe0c430 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -2790,16 +2790,29 @@ impl InvertedPartition { metrics: &dyn MetricsCollector, shared_threshold: Arc, ) -> Result>> { - let documents = ModernWandDocuments::new(lengths, visibility); - self.bm25_search_with_documents( - &documents, - params, - operator, - postings, - impact_scorer, - metrics, - shared_threshold, - ) + if visibility.is_all() { + let documents = ModernWandDocuments::all(lengths); + self.bm25_search_with_documents( + &documents, + params, + operator, + postings, + impact_scorer, + metrics, + shared_threshold, + ) + } else { + let documents = ModernWandDocuments::filtered(lengths, visibility); + self.bm25_search_with_documents( + &documents, + params, + operator, + postings, + impact_scorer, + metrics, + shared_threshold, + ) + } } #[instrument(level = "debug", skip_all)] diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index d9f23b37316..dcd8f5ac0d0 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -1561,13 +1561,61 @@ pub(super) trait WandDocuments { fn flat_doc_length(&self, doc_id: u64, document_key: u64, compressed: bool) -> u32; } -pub(super) struct ModernWandDocuments<'a> { +pub(super) trait ModernVisibility { + fn selected(&self, doc_id: DocId) -> bool; + fn len(&self, total_docs: usize) -> usize; + fn iter(&self) -> Option + '_>>; +} + +pub(super) struct AllModernDocuments; + +impl ModernVisibility for AllModernDocuments { + #[inline] + fn selected(&self, _doc_id: DocId) -> bool { + true + } + + fn len(&self, total_docs: usize) -> usize { + total_docs + } + + fn iter(&self) -> Option + '_>> { + None + } +} + +impl ModernVisibility for &DocVisibility { + #[inline] + fn selected(&self, doc_id: DocId) -> bool { + DocVisibility::selected(self, doc_id) + } + + fn len(&self, total_docs: usize) -> usize { + DocVisibility::len(self, total_docs) + } + + fn iter(&self) -> Option + '_>> { + DocVisibility::iter(self) + .map(|doc_ids| Box::new(doc_ids) as Box>) + } +} + +pub(super) struct ModernWandDocuments<'a, V> { lengths: &'a DocLengths, - visibility: &'a DocVisibility, + visibility: V, +} + +impl<'a> ModernWandDocuments<'a, AllModernDocuments> { + pub(crate) fn all(lengths: &'a DocLengths) -> Self { + Self { + lengths, + visibility: AllModernDocuments, + } + } } -impl<'a> ModernWandDocuments<'a> { - pub(crate) fn new(lengths: &'a DocLengths, visibility: &'a DocVisibility) -> Self { +impl<'a> ModernWandDocuments<'a, &'a DocVisibility> { + pub(crate) fn filtered(lengths: &'a DocLengths, visibility: &'a DocVisibility) -> Self { Self { lengths, visibility, @@ -1575,7 +1623,7 @@ impl<'a> ModernWandDocuments<'a> { } } -impl WandDocuments for ModernWandDocuments<'_> { +impl WandDocuments for ModernWandDocuments<'_, V> { type Candidate = DocId; fn len(&self) -> usize { From 964bfe243e64de2382a5801389f7d5529a527806 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 21 Jul 2026 19:11:19 +0800 Subject: [PATCH 10/10] perf(fts): skip resident corpus synchronization --- rust/lance-index/src/scalar/inverted/index.rs | 55 +++++++++++++++++-- 1 file changed, 49 insertions(+), 6 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 9b7bfe0c430..bde7c7aa8fe 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -1246,6 +1246,14 @@ impl InvertedIndex { &self, request: ModernSearchRequest<'_>, ) -> Result<(Vec, Vec)> { + // Old partitioned files without persisted stats populate their + // fallback stats before deferred candidate orchestration. A resident + // search can skip this full-index synchronization: standard prewarm + // has already initialized it, while any independently resident + // partition loads its lengths before constructing a local scorer. + if self.corpus_stats.get().is_none() { + self.aggregate_corpus_stats().await?; + } let ranked = self.bm25_search_modern_candidates(request).await?; self.resolve_deferred_modern_candidates(ranked).await } @@ -1270,12 +1278,6 @@ impl InvertedIndex { self.partitions.len() ))); } - // Ensures old partitioned files without persisted stats populate their - // fallback stats before any synchronous local scorer is constructed. - if self.corpus_stats.get().is_none() { - self.aggregate_corpus_stats().await?; - } - let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); let local_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); let parts = self @@ -10915,6 +10917,47 @@ mod tests { assert!(resident.0.contains(&200)); } + #[tokio::test] + async fn test_resident_modern_search_loads_partition_stats_without_global_stats() { + let (_tmpdir, index) = load_global_scoring_test_index(false).await; + assert!(index.corpus_stats.get().is_none()); + assert!( + index + .partitions + .iter() + .all(|partition| partition.docs.cached_stats().is_none()) + ); + + for partition in &index.partitions { + partition.docs.modern().unwrap().projection().await.unwrap(); + } + assert!(index.has_resident_document_projections()); + + let scorer = MemBM25Scorer::new(56_100, 112, HashMap::from([("alpha".to_owned(), 2)])); + let result = index + .bm25_search( + Arc::new(Tokens::new(vec!["alpha".to_owned()], DocType::Text)), + Arc::new(FtsSearchParams::new().with_limit(Some(2))), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + Some(&scorer), + ) + .await + .unwrap(); + + assert_eq!(result.0.len(), 2); + assert!(result.0.contains(&100)); + assert!(result.0.contains(&200)); + assert!(index.corpus_stats.get().is_none()); + assert!( + index + .partitions + .iter() + .all(|partition| partition.docs.cached_stats().is_some()) + ); + } + async fn search_test_impact_partition( partition: &InvertedPartition, tokens: &Tokens,