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/Cargo.toml b/rust/lance-index/Cargo.toml index 4387dc299f7..ec128ee1b85 100644 --- a/rust/lance-index/Cargo.toml +++ b/rust/lance-index/Cargo.toml @@ -12,7 +12,7 @@ categories.workspace = true rust-version.workspace = true [dependencies] -arc-swap.workspace = true +arc-swap = { workspace = true, features = ["weak"] } arrow.workspace = true arrow-array.workspace = true arrow-ipc.workspace = true diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index f51ff18a103..d67c164f6cb 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 588c8822942..83c2bf7b97d 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..37bb79a56cf --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/documents.rs @@ -0,0 +1,2164 @@ +// 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::borrow::Cow; +use std::sync::{Arc, OnceLock, Weak}; + +use arc_swap::ArcSwapWeak; +use arrow::buffer::ScalarBuffer; +use arrow_array::{Array, RecordBatch, UInt32Array, UInt64Array}; +use lance_core::cache::{CacheKey, WeakLanceCache}; +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, RowAddrSelection, RowAddrTreeMap}; +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"; + +/// 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, +} + +/// One partition's immutable `DocId -> row address` column. +/// +/// The independently weighed cache entry is the only long-lived owner. The +/// address projection keeps a weak handle and query-local guards upgrade it, +/// so cache eviction releases the column once in-flight queries finish. +#[derive(Debug)] +pub(super) struct CachedDocRowIds { + pub(crate) row_ids: Arc, +} + +impl DeepSizeOf for CachedDocRowIds { + fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize { + self.row_ids.len() * std::mem::size_of::() + } +} + +/// Cache key for one partition's [`CachedDocRowIds`]. +#[derive(Debug, Clone)] +pub(super) struct DocRowIdsKey { + pub(crate) partition_id: u64, +} + +impl CacheKey for DocRowIdsKey { + type ValueType = CachedDocRowIds; + + fn key(&self) -> Cow<'_, str> { + format!("doc-row-ids-{}", self.partition_id).into() + } + + fn type_name() -> &'static str { + "DocRowIds" + } +} + +/// 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(), + ) + } + + 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() { + Some(norms) => dequantize_doc_length(norms[doc_id.as_usize()]), + None => self.exact(doc_id), + } + } +} + +#[derive(Debug)] +enum AddressValues { + Shared { len: usize }, + Owned(Arc>), +} + +impl AddressValues { + fn len(&self) -> usize { + match self { + Self::Shared { len } => *len, + Self::Owned(values) => values.len(), + } + } +} + +impl DeepSizeOf for AddressValues { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + match self { + Self::Shared { .. } => 0, + Self::Owned(values) => values.deep_size_of_children(context), + } + } +} + +/// 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: &ResidentAddressProjection) -> Self { + if projection.projection.live_docs.is_none() + && (1..projection.len()).all(|index| { + projection.stored_address(index - 1) <= projection.stored_address(index) + }) + { + return Self::Identity; + } + + let mut doc_ids = match projection.projection.live_docs.as_ref() { + Some(live_docs) => live_docs.iter().collect::>(), + None => (0..projection.len() as u32).collect::>(), + }; + doc_ids.sort_unstable_by_key(|&doc_id| projection.stored_address(doc_id as usize)); + Self::Sorted(doc_ids.into_boxed_slice()) + } + + fn len(&self, projection: &ResidentAddressProjection) -> usize { + match self { + Self::Identity => projection.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: &ResidentAddressProjection, position: usize) -> u64 { + projection.stored_address(self.doc_id_at(position) as usize) + } + + fn partition_point( + &self, + projection: &ResidentAddressProjection, + 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: &ResidentAddressProjection, + 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: &ResidentAddressProjection, + 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: &ResidentAddressProjection, + 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 { + addresses: AddressValues, + /// `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>, +} + +/// A query-scoped projection guard. Shared addresses remain alive only while +/// this guard or their independently weighed cache entry owns the Arrow column. +#[derive(Debug, Clone)] +pub(super) struct ResidentAddressProjection { + projection: Arc, + addresses: ResidentAddressValues, +} + +#[derive(Debug, Clone)] +enum ResidentAddressValues { + Shared(Arc), + Owned(Arc>), +} + +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) + + self + .doc_ids_by_address + .get() + .map(|lookup| lookup.deep_size_of_children(context)) + .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 { len: raw.len() }, + live_docs: None, + doc_ids_by_address: OnceCell::new(), + }); + }; + + 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(Arc::new(addresses)), + live_docs: Some(live_docs), + doc_ids_by_address: OnceCell::new(), + }) + } + + fn resident( + self: &Arc, + shared_addresses: Option>, + ) -> Option { + match &self.addresses { + AddressValues::Shared { len } => { + let shared_addresses = shared_addresses?; + debug_assert_eq!(shared_addresses.len(), *len); + Some(ResidentAddressProjection { + projection: self.clone(), + addresses: ResidentAddressValues::Shared(shared_addresses), + }) + } + AddressValues::Owned(values) => Some(ResidentAddressProjection { + projection: self.clone(), + addresses: ResidentAddressValues::Owned(values.clone()), + }), + } + } +} + +impl ResidentAddressProjection { + fn len(&self) -> usize { + self.projection.addresses.len() + } + + fn stored_address(&self, index: usize) -> u64 { + match &self.addresses { + ResidentAddressValues::Shared(values) => values.value(index), + ResidentAddressValues::Owned(values) => values[index], + } + } + + fn address(&self, doc_id: DocId) -> Option { + if self + .projection + .live_docs + .as_ref() + .is_some_and(|live| !live.contains(doc_id.get())) + { + None + } else { + Some(self.stored_address(doc_id.as_usize())) + } + } + + fn live_doc_ids(&self) -> RoaringBitmap { + self.projection + .live_docs + .clone() + .unwrap_or_else(|| (0..self.len() as u32).collect()) + } + + async fn doc_ids_by_address(&self) -> Result> { + self.projection + .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 + } + }) + .await + .cloned() + } + + async fn materialize_visibility(self, mask: Arc) -> Result { + let lookup = self.doc_ids_by_address().await?; + let projection = self; + spawn_cpu(move || Result::Ok(lookup.visibility(&projection, &mask))).await + } +} + +/// Query-local selection in the partition-local DocId domain. +#[derive(Debug, Clone)] +pub(super) enum DocVisibility { + All, + Selected(RoaringBitmap), + Filtered { + projection: ResidentAddressProjection, + mask: Arc, + }, +} + +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 { + 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)), + } + } + + pub(crate) fn len(&self, total_docs: usize) -> usize { + match self { + Self::All => total_docs, + Self::Selected(selected) => selected.len() as usize, + Self::Filtered { .. } => total_docs, + } + } + + 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::Selected(selected) => Some(selected.iter().map(DocId::new)), + Self::All | Self::Filtered { .. } => None, + } + } +} + +/// Modern query-side document state for one partition. +pub(super) struct PartitionDocuments { + store: Arc, + path: String, + partition_id: u64, + index_cache: WeakLanceCache, + num_docs: usize, + persisted_total_tokens: Option, + quantized_scoring: bool, + remapper: Option>, + lengths: OnceCell>, + projection: OnceCell>, + shared_addresses: ArcSwapWeak, + prewarm_complete: 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) 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()), + 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, + partition_id: u64, + index_cache: WeakLanceCache, + 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, + partition_id, + index_cache, + num_docs, + persisted_total_tokens, + quantized_scoring, + remapper, + lengths: OnceCell::new(), + projection: OnceCell::new(), + shared_addresses: ArcSwapWeak::from(Weak::new()), + prewarm_complete: 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() + } + + #[cfg(test)] + pub(crate) fn address_buffer_handle(&self) -> Weak { + self.shared_addresses.load_full() + } + + pub(crate) fn projection_resident(&self) -> bool { + self.resident_address_projection().is_some() + } + + 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()) + && self.projection_resident() + } + + async fn reader(&self) -> Result> { + self.store.open_index_file(&self.path).await + } + + async fn row_ids_column(&self) -> Result> { + let store = self.store.clone(); + let path = self.path.clone(); + let num_docs = self.num_docs; + let cached = self + .index_cache + .get_or_insert_with_key( + DocRowIdsKey { + partition_id: self.partition_id, + }, + || async move { + let reader = store.open_index_file(&path).await?; + let batch = reader.read_range(0..num_docs, Some(&[ROW_ID])).await?; + let row_ids = required_u64_column(&batch, ROW_ID, &path)?; + if row_ids.null_count() != 0 { + return Err(corrupt_docs( + &path, + format!("{ROW_ID} contains null values"), + )); + } + if row_ids.len() != num_docs { + return Err(corrupt_docs( + &path, + format!( + "{ROW_ID} has {} rows but the file footer reports {num_docs}", + row_ids.len() + ), + )); + } + Ok(CachedDocRowIds { + row_ids: Arc::new(row_ids.clone()), + }) + }, + ) + .await?; + let row_ids = cached.row_ids.clone(); + self.shared_addresses.store(Arc::downgrade(&row_ids)); + Ok(row_ids) + } + + 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, + )?)) + } + + 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?; + self.lengths_from_batch(&batch) + }) + .await + .cloned() + } + + /// Return resident lengths without entering the asynchronous singleflight path. + pub(crate) fn cached_lengths(&self) -> Option> { + self.lengths.get().cloned() + } + + fn resident_address_projection(&self) -> Option { + let projection = self.projection.get()?.clone(); + let shared_addresses = match &projection.addresses { + AddressValues::Shared { .. } => self.shared_addresses.load().upgrade(), + AddressValues::Owned(_) => None, + }; + projection.resident(shared_addresses) + } + + pub(crate) async fn address_projection(&self) -> Result { + if let Some(projection) = self.resident_address_projection() { + return Ok(projection); + } + + let row_ids = self.row_ids_column().await?; + let projection = self + .projection + .get_or_try_init(|| async { + Result::Ok(Arc::new(VersionAddressProjection::try_new( + row_ids.as_ref(), + self.num_docs, + self.remapper.as_deref(), + &self.path, + )?)) + }) + .await + .cloned()?; + projection.resident(Some(row_ids)).ok_or_else(|| { + Error::internal(format!( + "address projection for {} could not bind its cache-managed ROW_ID column", + self.path + )) + }) + } + + pub(crate) async fn visibility( + &self, + mask: Arc, + materialize_selected: bool, + ) -> Result { + if let Some(visibility) = self.immediate_visibility(mask.clone(), materialize_selected) { + return Ok(visibility); + } + + let projection = self.address_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(DocVisibility::Filtered { projection, mask }) + } + } + + /// 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.resident_address_projection()?; + 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() { + return Ok(Vec::new()); + } + self.validate_doc_ids(doc_ids)?; + if let Some(projection) = self.resident_address_projection() { + return self.resolve_projected_addresses(&projection, doc_ids); + } + if self.remapper.is_some() { + let projection = self.address_projection().await?; + return self.resolve_projected_addresses(&projection, doc_ids); + } + + let row_ids = self.row_ids_column().await?; + Ok(doc_ids + .iter() + .map(|doc_id| row_ids.value(doc_id.as_usize())) + .collect()) + } + + fn validate_doc_ids(&self, doc_ids: &[DocId]) -> Result<()> { + 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 + ), + )); + } + } + Ok(()) + } + + fn resolve_projected_addresses( + &self, + projection: &ResidentAddressProjection, + doc_ids: &[DocId], + ) -> Result> { + 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() + } + + /// Resolve addresses synchronously when the cache-managed projection buffer + /// is resident. The returned `None` asks the caller to use the async reload path. + pub(crate) fn cached_row_addresses(&self, doc_ids: &[DocId]) -> Result>> { + self.validate_doc_ids(doc_ids)?; + let Some(projection) = self.resident_address_projection() else { + return Ok(None); + }; + self.resolve_projected_addresses(&projection, doc_ids) + .map(Some) + } + + /// Estimated Arrow payload retained while loading this partition's cached + /// row-address column. + /// 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_resident() { + return 0; + } + self.num_docs.saturating_mul(std::mem::size_of::()) + } + + /// 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.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 row_ids = + Arc::new(required_u64_column(&batch, ROW_ID, &self.path)?.clone()); + if row_ids.null_count() != 0 { + return Err(corrupt_docs( + &self.path, + format!("{ROW_ID} contains null values"), + )); + } + let projection = Arc::new(VersionAddressProjection::try_new( + row_ids.as_ref(), + self.num_docs, + self.remapper.as_deref(), + &self.path, + )?); + let cached_row_ids = Arc::new(CachedDocRowIds { + row_ids: row_ids.clone(), + }); + self.index_cache + .insert_with_key( + &DocRowIdsKey { + partition_id: self.partition_id, + }, + cached_row_ids, + ) + .await; + self.shared_addresses.store(Arc::downgrade(&row_ids)); + + // 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); + } + + let lengths = self.lengths().await?; + spawn_cpu(move || { + let _ = lengths.scoring_norms(); + Result::Ok(()) + }) + .await?; + self.address_projection() + .await? + .doc_ids_by_address() + .await?; + Result::Ok(()) + }) + .await?; + if !self.projection_resident() { + self.address_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 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::ops::Range; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + use arrow_array::{ArrayRef, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use async_trait::async_trait; + use lance_core::cache::{LanceCache, QuickCacheBackend}; + use lance_core::utils::tempfile::TempObjDir; + 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}; + + use super::*; + + #[derive(Debug, Default)] + struct DocumentReadCounts { + open_calls: AtomicUsize, + 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); + } + } + } + + 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] + 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); + if let Some(fault) = &self.fault { + fault.apply().await?; + } + 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); + if let Some(fault) = &self.fault { + fault.apply().await?; + } + 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, + fault: Option>, + } + + 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(), + fault: self.fault.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 { + self.counts.open_calls.fetch_add(1, Ordering::Relaxed); + Ok(Arc::new(CountingReader { + inner: reader, + counts: self.counts.clone(), + fault: self.fault.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(), + fault: self.fault.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, Arc) { + let directory = TempObjDir::default(); + let cache = Arc::new(LanceCache::with_capacity(1024 * 1024)); + test_store_with_cache(directory, cache) + } + + fn eviction_test_store() -> (TempObjDir, Arc, Arc) { + let directory = TempObjDir::default(); + let cache = Arc::new(LanceCache::with_backend(Arc::new( + QuickCacheBackend::with_capacity(1024 * 1024), + ))); + test_store_with_cache(directory, cache) + } + + fn test_store_with_cache( + directory: TempObjDir, + cache: Arc, + ) -> (TempObjDir, Arc, Arc) { + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + directory.clone(), + cache.clone(), + )); + (directory, store, cache) + } + + 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, + index_cache: &LanceCache, + remapper: Option>, + ) -> Result { + let reader = store.open_index_file(path).await?; + PartitionDocuments::try_new( + store, + path.to_owned(), + 0, + WeakLanceCache::from(index_cache), + 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(), + 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>, + } + + 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 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 live_documents_use_doc_ids() { + let projection = Arc::new(VersionAddressProjection { + addresses: AddressValues::Owned(Arc::new(vec![10, 20, 30])), + live_docs: Some(RoaringBitmap::from_iter([0, 2])), + doc_ids_by_address: OnceCell::new(), + }); + let projection = projection.resident(None).unwrap(); + let selected = projection.live_doc_ids(); + assert_eq!(selected.iter().collect::>(), vec![0, 2]); + } + + #[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 = Arc::new( + VersionAddressProjection::try_new(&raw, 4, Some(&remapper), "docs") + .expect("valid projection"), + ); + let projection = projection.resident(None).unwrap(); + + 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 all_live = projection.live_doc_ids(); + assert_eq!(all_live.iter().collect::>(), vec![0, 2, 3]); + + let allowed = Arc::new(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([ + 100, 40, + ]))); + let DocVisibility::Selected(selected) = projection + .clone() + .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 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 + .clone() + .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(Arc::new(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 projection = projection.resident(None).unwrap(); + + 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 + .clone() + .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(Arc::new(vec![10, 20, 30])), + live_docs: Some(RoaringBitmap::from_iter([0, 2])), + doc_ids_by_address: OnceCell::new(), + }); + let resident = projection.resident(None).unwrap(); + let visibility = DocVisibility::Filtered { + projection: resident, + 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] + 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, cache) = 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, cache.as_ref(), 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, cache) = 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, cache.as_ref(), 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 prewarm_loads_document_columns_and_address_lookup_once() { + let (_directory, store, cache) = 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, cache.as_ref(), None) + .await + .unwrap(); + + futures::future::join_all((0..8).map(|_| documents.prewarm())) + .await + .into_iter() + .collect::>>() + .unwrap(); + let projection = documents.address_projection().await.unwrap(); + let first_lookup = projection.doc_ids_by_address().await.unwrap(); + documents.prewarm().await.unwrap(); + let second_lookup = documents + .address_projection() + .await + .unwrap() + .doc_ids_by_address() + .await + .unwrap(); + + 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 + )); + assert!(Arc::ptr_eq(&first_lookup, &second_lookup)); + assert_eq!( + documents.cached_row_addresses(&[DocId::new(1)]).unwrap(), + Some(vec![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); + assert!( + cache + .get_with_key(&DocRowIdsKey { partition_id: 0 }) + .await + .is_some() + ); + } + + #[tokio::test] + async fn filtered_visibility_releases_addresses_after_cache_eviction() { + let (_directory, store, cache) = eviction_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, cache.as_ref(), None) + .await + .unwrap(); + let mask = Arc::new(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([20]))); + + let visibility = documents.visibility(mask.clone(), false).await.unwrap(); + assert!(!visibility.selected(DocId::new(0))); + assert!(visibility.selected(DocId::new(1))); + assert!(!visibility.selected(DocId::new(2))); + + let weak_addresses = documents.address_buffer_handle(); + + cache.clear().await; + assert!(weak_addresses.upgrade().is_some()); + assert!(documents.projection_resident()); + + drop(visibility); + assert!(weak_addresses.upgrade().is_none()); + assert!(!documents.projection_resident()); + + let reloaded = documents.visibility(mask, false).await.unwrap(); + assert!(reloaded.selected(DocId::new(1))); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 6); + } + + #[tokio::test] + async fn prewarm_reloads_addresses_after_cache_eviction() { + let (_directory, store, cache) = eviction_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, cache.as_ref(), None) + .await + .unwrap(); + + documents.prewarm().await.unwrap(); + assert!(documents.query_ready()); + let weak_addresses = documents.address_buffer_handle(); + + cache.clear().await; + assert!(weak_addresses.upgrade().is_none()); + assert!(documents.projection_loaded()); + assert!(!documents.projection_resident()); + assert!(!documents.query_ready()); + + documents.prewarm().await.unwrap(); + assert!(documents.query_ready()); + assert_eq!( + documents + .cached_row_addresses(&[DocId::new(2), DocId::new(0)]) + .unwrap(), + Some(vec![30, 10]) + ); + assert_eq!(counts.length_rows.load(Ordering::Relaxed), 3); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 6); + } + + #[tokio::test] + async fn prewarm_materializes_quantized_norms_before_becoming_query_ready() { + let (_directory, store, cache) = 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(), + 0, + WeakLanceCache::from(cache.as_ref()), + 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, cache) = 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, cache.as_ref(), 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, cache.as_ref(), 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, cache) = 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, cache.as_ref(), 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); + } + + #[tokio::test] + async fn invalid_or_mismatched_stats_are_corruption_not_fallback() { + let (_directory, store, cache) = 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", cache.as_ref(), 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", cache.as_ref(), 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_reuses_the_cached_row_id_column() { + let (_directory, store, cache) = 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, cache.as_ref(), 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.estimated_address_read_bytes(&point_ids), + num_docs as usize * std::mem::size_of::() + ); + assert_eq!( + documents.resolve_addresses(&point_ids).await.unwrap(), + vec![1005, 1006, 1010, 1005] + ); + assert_eq!(counts.ranges_calls.load(Ordering::Relaxed), 0); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 1); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 600); + assert!( + cache + .get_with_key(&DocRowIdsKey { partition_id: 0 }) + .await + .is_some() + ); + + let bulk_ids = (0..=512).step_by(2).map(DocId::new).collect::>(); + 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)); + assert_eq!(counts.ranges_calls.load(Ordering::Relaxed), 0); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 1); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 600); + assert!(!documents.projection_loaded()); + } + + #[tokio::test] + async fn final_address_resolution_reloads_after_cache_eviction() { + let (_directory, store, _cache) = 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 no_retention_cache = LanceCache::no_cache(); + let documents = open_documents(counting, path, &no_retention_cache, None) + .await + .unwrap(); + + for _ in 0..2 { + assert_eq!( + documents + .resolve_addresses(&[DocId::new(2), DocId::new(0)]) + .await + .unwrap(), + vec![30, 10] + ); + } + assert_eq!(counts.ranges_calls.load(Ordering::Relaxed), 0); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 2); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 6); + assert!(!documents.projection_loaded()); + } + + #[tokio::test] + async fn document_column_nulls_are_reported_as_corruption() { + let (_directory, store, cache) = 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", cache.as_ref(), 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", cache.as_ref(), None) + .await + .unwrap(); + assert!( + documents + .resolve_addresses(&[DocId::new(1)]) + .await + .unwrap_err() + .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 a6db0e355c9..1029f5c7093 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,19 +51,25 @@ 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::{ + 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::{ builder::{ - BLOCK_SIZE, doc_file_path, inverted_list_schema_for_version_with_block_size_and_impacts, - posting_file_path, token_file_path, + BLOCK_SIZE, ScoredDoc, doc_file_path, + inverted_list_schema_for_version_with_block_size_and_impacts, posting_file_path, + token_file_path, }, iter::PlainPostingListIterator, query::*, @@ -81,7 +88,6 @@ use crate::scalar::{ OldIndexDataFilter, RowIdRemapper, ScalarIndex, ScalarIndexParams, SearchResult, TokenQuery, UpdateCriteria, }; -use crate::vector::graph::OrderedFloat; use crate::{FtsPrewarmOptions, Index}; use crate::{prefilter::PreFilter, scalar::inverted::iter::take_fst_keys}; use std::str::FromStr; @@ -98,16 +104,16 @@ pub const INVERT_LIST_FILE: &str = "invert.lance"; pub const DOCS_FILE: &str = "docs.lance"; pub const METADATA_FILE: &str = "metadata.lance"; -/// Partitions searched per cpu-pool task in `bm25_search`. Chunking bounds -/// the task rate and lets the shared top-k floor propagate between the -/// partitions a thread scores back-to-back. `LANCE_FTS_SEARCH_CHUNK=1` -/// restores one-task-per-partition. +/// Partitions searched per CPU-pool task. Each chunk loads concurrently and +/// then scores sequentially so query concurrency does not flood the pool with +/// one small task per partition. `LANCE_FTS_SEARCH_CHUNK=1` restores the +/// per-partition task shape. fn fts_search_chunk() -> usize { static CHUNK: LazyLock = LazyLock::new(|| { std::env::var("LANCE_FTS_SEARCH_CHUNK") .ok() - .and_then(|v| v.parse().ok()) - .filter(|&n| n >= 1) + .and_then(|value| value.parse().ok()) + .filter(|&value| value >= 1) .unwrap_or(16) }); *CHUNK @@ -282,10 +288,184 @@ pub fn validate_format_version_block_size( } #[derive(Debug)] -struct PartitionCandidates { +struct PartitionCandidates { tokens_by_position: Vec, grouped_expansions: Vec, - candidates: Vec, + candidates: Vec>, +} + +struct ModernSearchRequest<'a> { + tokens: Arc, + params: Arc, + operator: Operator, + mask: Arc, + metrics: Arc, + scorer: &'a MemBM25Scorer, + impact_scorer: Arc, + 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, + 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 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, + 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)] @@ -296,6 +476,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 { @@ -477,6 +678,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, @@ -486,6 +699,10 @@ pub struct InvertedIndex { format_version: InvertedListFormatVersion, pub(crate) partitions: Vec>, corpus_stats: Arc>, + prewarm_state: Arc>, + /// Optimistic fast-path hint. Cache eviction can make it stale; the + /// resident resolver clears it when a weak projection upgrade misses. + 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, @@ -656,9 +873,6 @@ impl InvertedIndex { /// expansions, not just the raw query tokens — otherwise /// `query_weight(expanded_token)` returns 0 and the BM25 contribution /// of every expanded match is discarded. - /// - /// `metrics` is forwarded to the per-token metadata cache boundary so the - /// caller's per-query cache counters see the reads triggered here. pub async fn bm25_base_scorer( &self, query_tokens: &Tokens, @@ -699,10 +913,6 @@ impl InvertedIndex { Ok(MemBM25Scorer::new(total_tokens, num_docs, token_docs)) } - /// Collect the `(total_tokens, num_docs, per_term_df)` triple used to - /// combine per-segment BM25 statistics into a global scorer. `metrics` - /// is threaded to record per-token metadata cache activity in the - /// caller's per-query counters. pub async fn bm25_stats_for_terms( &self, terms: &[String], @@ -715,29 +925,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,75 +1100,48 @@ impl InvertedIndex { if limit == 0 { return Ok((Vec::new(), Vec::new())); } - - /// Global-heap entry carrying the score plus where its row_id comes - /// from. `Pending` candidates keep (slot, doc_id) through the top-k - /// so only the final survivors pay row_id resolution. - struct MergedCandidate { - score: OrderedFloat, - /// Index into the per-query partition list for `Pending` addresses. - slot: u32, - addr: CandidateAddr, - } - - impl PartialEq for MergedCandidate { - fn eq(&self, other: &Self) -> bool { - self.score == other.score - } - } - - impl Eq for MergedCandidate {} - - impl PartialOrd for MergedCandidate { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } - } - - impl Ord for MergedCandidate { - fn cmp(&self, other: &Self) -> std::cmp::Ordering { - self.score.cmp(&other.score) - } - } - - fn push_scored_candidate( - candidates: &mut BinaryHeap>, - limit: usize, - slot: u32, - addr: CandidateAddr, - score: f32, - ) { - let candidate = MergedCandidate { - score: OrderedFloat(score), - slot, - addr, - }; - if candidates.len() < limit { - candidates.push(Reverse(candidate)); - } else if candidates.peek().unwrap().0.score.0 < score { - candidates.pop(); - candidates.push(Reverse(candidate)); - } - } - 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(ModernSearchRequest { + 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`). Only sound for the - // impact (global) scorer, whose scores are comparable across - // partitions; legacy BM25 scores are partition-local scale, so those - // partitions each get a private floor below. + #[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())); - // Partitions are processed in chunks: each chunk loads its postings - // and scoring DocSets asynchronously, then ONE cpu-pool task searches - // the whole chunk sequentially. Chunks pipeline against each other - // through buffer_unordered (chunk i searches while chunk i+1 loads). - // Chunking bounds the cpu-pool task rate (vs one tiny task per - // partition) and lets the shared top-k floor propagate immediately - // between partitions searched on the same thread. + let io_parallelism = self.store.io_parallelism(); let parts = self .partitions .chunks(fts_search_chunk()) @@ -964,17 +1154,19 @@ impl InvertedIndex { let impact_scorer = impact_scorer.clone(); let impact_shared_threshold = impact_shared_threshold.clone(); async move { - // Load the chunk's partitions concurrently; only the - // search below is batched onto one cpu task. let loads = chunk.into_iter().map(|part| { 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(); 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(), @@ -983,20 +1175,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(None); } - let docs_for_wand = - part.docs.docs_for_wand(operator, mask.as_ref()).await?; let max_position = postings .iter() .map(|posting| posting.term_index() as usize) @@ -1004,11 +1185,14 @@ 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 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 { Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())) @@ -1016,263 +1200,501 @@ impl InvertedIndex { let wand_scorer = use_global_scorer.then(|| impact_scorer.clone()); Result::Ok(Some(( part, - docs_for_wand, + docs, postings, wand_scorer, - partition_threshold, + threshold, tokens_by_position, grouped_expansions, ))) } }); - let loaded: Vec<_> = stream::iter(loads) - .buffer_unordered(self.store.io_parallelism()) + let loaded = stream::iter(loads) + .buffer_unordered(io_parallelism) .try_collect::>() .await? .into_iter() .flatten() - .collect(); + .collect::>(); if loaded.is_empty() { return Result::Ok(Vec::new()); } - let params = params.clone(); - let mask = mask.clone(); - let metrics = metrics.clone(); + let results = spawn_cpu(move || { let mut results = Vec::with_capacity(loaded.len()); for ( part, - docs_for_wand, + docs, postings, wand_scorer, - partition_threshold, + threshold, tokens_by_position, grouped_expansions, ) in loaded { - let candidates = part.bm25_search( - docs_for_wand.as_ref(), + let candidates = part.bm25_search_legacy( + docs.as_ref(), params.as_ref(), operator, - mask.clone(), + mask.as_ref(), postings, wand_scorer, metrics.as_ref(), - partition_threshold, + threshold, )?; - results.push(( - part, - PartitionCandidates { - tokens_by_position, - grouped_expansions, - candidates, - }, - )); + results.push(PartitionCandidates { + tokens_by_position, + grouped_expansions, + candidates, + }); } - std::result::Result::<_, Error>::Ok(results) + Result::Ok(results) }) .await?; Result::Ok(results) } }) .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().min(32)) .map_ok(|results| stream::iter(results.into_iter().map(Result::Ok))) .try_flatten(); - let mut idf_cache: HashMap = HashMap::new(); - // Partitions that produced candidates, indexed by the `slot` carried - // in deferred heap entries. - let mut resolving_parts: Vec> = Vec::new(); - while let Some((part, res)) = parts.try_next().await? { - if res.candidates.is_empty() { - continue; - } - let slot = resolving_parts.len() as u32; - resolving_parts.push(part); - let PartitionCandidates { - tokens_by_position, - grouped_expansions, - 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); - } - push_scored_candidate(&mut candidates, limit, slot, 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); - } - 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); - } - } - push_scored_candidate(&mut candidates, limit, slot, addr, score); - } - } + async fn bm25_search_modern( + &self, + 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 } + } - // Resolve row_ids only for the candidates that survived the global - // top-k: group deferred survivors per partition and batch-resolve — - // at most `limit` lookups regardless of how many partitions - // contributed candidates. - /// One partition's surviving deferred candidates: positions in the - /// merged result list paired with the doc_ids to resolve. - type DeferredGroup = Vec<(usize, u32)>; - let sorted = candidates.into_sorted_vec(); - let mut row_ids = Vec::with_capacity(sorted.len()); - let mut scores = Vec::with_capacity(sorted.len()); - let mut deferred: HashMap = HashMap::new(); - for (pos, Reverse(candidate)) in sorted.into_iter().enumerate() { - scores.push(candidate.score.0); - match candidate.addr { - CandidateAddr::RowId(row_id) => row_ids.push(row_id), - CandidateAddr::Pending(doc_id) => { - deferred - .entry(candidate.slot) - .or_default() - .push((pos, doc_id)); - // Placeholder, overwritten by the batch resolution below. - row_ids.push(0); - } - } + fn has_resident_document_projections(&self) -> bool { + if self.document_projections_resident.load(Ordering::Acquire) { + return true; } - if !deferred.is_empty() { - let groups = deferred - .into_iter() - .map(|(slot, entries)| (resolving_parts[slot as usize].clone(), entries)) - .collect::>(); - let resolved: Vec<(DeferredGroup, Vec)> = - stream::iter(groups.into_iter().map(|(part, entries)| async move { - let doc_ids: Vec = entries.iter().map(|&(_, doc_id)| doc_id).collect(); - let resolved = part.docs.resolve_row_ids(&doc_ids).await?; - Result::Ok((entries, resolved)) - })) - .buffer_unordered(get_num_compute_intensive_cpus()) - .try_collect() - .await?; - for (entries, resolved) in resolved { - for ((pos, _), row_id) in entries.into_iter().zip(resolved) { - row_ids[pos] = row_id; - } - } + let resident = self.document_projections_resident_now(); + if resident { + self.document_projections_resident + .store(true, Ordering::Release); } - Ok((row_ids, scores)) + resident } - async fn load_legacy_index( - store: Arc, - frag_reuse_index: Option>, - index_cache: &LanceCache, - ) -> Result> { - log::warn!("loading legacy FTS index"); - let tokens_fut = tokio::spawn({ - let store = store.clone(); - async move { - let token_reader = store.open_index_file(TOKENS_FILE).await?; - let tokenizer = token_reader - .schema() - .metadata - .get("tokenizer") - .map(|s| serde_json::from_str::(s)) - .transpose()? - .unwrap_or_default(); - let tokens = TokenSet::load(token_reader, TokenSetFormat::Arrow).await?; - Result::Ok((tokenizer, tokens)) - } - }); - let invert_list_fut = tokio::spawn({ - let store = store.clone(); - let index_cache_clone = index_cache.clone(); - async move { - let invert_list_reader = store.open_index_file(INVERT_LIST_FILE).await?; - let invert_list = - PostingListReader::try_new(invert_list_reader, &index_cache_clone).await?; - Result::Ok(Arc::new(invert_list)) - } - }); - let docs_fut = tokio::spawn({ - let store = store.clone(); - async move { - let docs_reader = store.open_index_file(DOCS_FILE).await?; - let docs = DocSet::load(docs_reader, true, frag_reuse_index).await?; - Result::Ok(docs) - } - }); + fn document_projections_resident_now(&self) -> bool { + self.partitions.iter().all(|partition| { + partition + .docs + .modern() + .is_some_and(|documents| documents.projection_resident()) + }) + } - let (tokenizer_config, tokens) = tokens_fut.await??; - let inverted_list = invert_list_fut.await??; - let docs = docs_fut.await??; + async fn bm25_search_modern_resident( + &self, + request: ModernSearchRequest<'_>, + ) -> Result<(Vec, Vec)> { + let ranked = self.bm25_search_modern_candidates(request).await?; + if let Some(result) = self.resolve_resident_modern_candidates(&ranked)? { + return Ok(result); + } + self.document_projections_resident + .store(false, Ordering::Release); + self.resolve_deferred_modern_candidates(ranked).await + } - let tokenizer = tokenizer_config.build()?; + async fn bm25_search_modern_deferred( + &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 + } - Ok(Arc::new(Self { - params: tokenizer_config, - store: store.clone(), - tokenizer, - token_set_format: TokenSetFormat::Arrow, - format_version: InvertedListFormatVersion::V1, - partitions: vec![Arc::new(InvertedPartition { - id: 0, + 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", + self.partitions.len() + ))); + } + let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + let io_parallelism = self.store.io_parallelism(); + let parts = self + .partitions + .chunks(fts_search_chunk()) + .enumerate() + .map(|(chunk_ordinal, chunk)| { + let first_partition_ordinal = chunk_ordinal * fts_search_chunk(); + let chunk = chunk + .iter() + .cloned() + .enumerate() + .map(|(offset, part)| (first_partition_ordinal + offset, part)) + .collect::>(); + 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(); + async move { + let loads = chunk.into_iter().map(|(partition_ordinal, part)| { + 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(); + 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(None); + } + let documents = part.docs.modern().cloned().ok_or_else(|| { + Error::internal("modern index contains legacy partition documents") + })?; + 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 = 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(None); + } + 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) + .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(); + } + let use_global_scorer = impact_safe || exact_scoring_required; + let threshold = if use_global_scorer { + impact_shared_threshold + } else { + Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())) + }; + let wand_scorer = use_global_scorer.then(|| impact_scorer.clone()); + Result::Ok(Some(( + partition_ordinal, + part, + lengths, + visibility, + postings, + wand_scorer, + threshold, + tokens_by_position, + grouped_expansions, + ))) + } + }); + let loaded = stream::iter(loads) + .buffer_unordered(io_parallelism) + .try_collect::>() + .await? + .into_iter() + .flatten() + .collect::>(); + if loaded.is_empty() { + return Result::Ok(Vec::new()); + } + + let results = spawn_cpu(move || { + let mut results = Vec::with_capacity(loaded.len()); + for ( + partition_ordinal, + part, + lengths, + visibility, + postings, + wand_scorer, + threshold, + tokens_by_position, + grouped_expansions, + ) in loaded + { + let candidates = part.bm25_search_modern( + lengths.as_ref(), + &visibility, + params.as_ref(), + operator, + postings, + wand_scorer, + metrics.as_ref(), + threshold, + )?; + results.push(( + partition_ordinal, + PartitionCandidates { + tokens_by_position, + grouped_expansions, + candidates, + }, + )); + } + Result::Ok(results) + }) + .await?; + Result::Ok(results) + } + }) + .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().min(32)) + .map_ok(|results| stream::iter(results.into_iter().map(Result::Ok))) + .try_flatten(); + while let Some((partition_ordinal, partition)) = parts.try_next().await? { + for (doc_id, score) in rescore_partition_candidates(partition, scorer, &mut idf_cache) { + push_scored_partition_doc( + &mut ranked, + limit, + PartitionDocId::try_new(partition_ordinal, doc_id)?, + score, + ); + } + } + + Ok(ranked.into_sorted_vec()) + } + + fn resolve_resident_modern_candidates( + &self, + ranked: &[Reverse], + ) -> Result, Vec)>> { + let mut addresses = vec![0; ranked.len()]; + let mut by_partition = BTreeMap::>::new(); + for (rank, Reverse(candidate)) in ranked.iter().enumerate() { + 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)); + } + for (partition_ordinal, entries) in by_partition { + let documents = self + .partitions + .get(partition_ordinal) + .and_then(|partition| partition.docs.modern()) + .ok_or_else(|| { + Error::internal(format!( + "resident FTS candidates reference missing modern partition ordinal {partition_ordinal}" + )) + })?; + let doc_ids = entries + .iter() + .map(|(_, doc_id)| *doc_id) + .collect::>(); + let Some(resolved) = documents.cached_row_addresses(&doc_ids)? else { + return Ok(None); + }; + for ((rank, _), address) in entries.into_iter().zip(resolved) { + addresses[rank] = address; + } + } + let scores = ranked + .iter() + .map(|Reverse(candidate)| candidate.score.0) + .collect(); + Ok(Some((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() { + 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 + .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}" + )) + })?; + 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 resolved = documents.resolve_addresses(&doc_ids).await?; + Result::Ok((entries, resolved)) + }); + } + 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; + } + } + let scores = ranked + .into_iter() + .map(|Reverse(candidate)| candidate.score.0) + .collect(); + Ok((addresses, scores)) + } + + async fn load_legacy_index( + store: Arc, + frag_reuse_index: Option>, + index_cache: &LanceCache, + ) -> Result> { + log::warn!("loading legacy FTS index"); + let tokens_fut = tokio::spawn({ + let store = store.clone(); + async move { + let token_reader = store.open_index_file(TOKENS_FILE).await?; + let tokenizer = token_reader + .schema() + .metadata + .get("tokenizer") + .map(|s| serde_json::from_str::(s)) + .transpose()? + .unwrap_or_default(); + let tokens = TokenSet::load(token_reader, TokenSetFormat::Arrow).await?; + Result::Ok((tokenizer, tokens)) + } + }); + let invert_list_fut = tokio::spawn({ + let store = store.clone(); + let index_cache_clone = index_cache.clone(); + async move { + let invert_list_reader = store.open_index_file(INVERT_LIST_FILE).await?; + let invert_list = + PostingListReader::try_new(invert_list_reader, &index_cache_clone).await?; + Result::Ok(Arc::new(invert_list)) + } + }); + let docs_fut = tokio::spawn({ + let store = store.clone(); + async move { + let docs_reader = store.open_index_file(DOCS_FILE).await?; + let docs = DocSet::load(docs_reader, true, frag_reuse_index).await?; + Result::Ok(docs) + } + }); + + let (tokenizer_config, tokens) = tokens_fut.await??; + let inverted_list = invert_list_fut.await??; + let docs = docs_fut.await??; + + let tokenizer = tokenizer_config.build()?; + + Ok(Arc::new(Self { + params: tokenizer_config, + store: store.clone(), + tokenizer, + token_set_format: TokenSetFormat::Arrow, + format_version: InvertedListFormatVersion::V1, + partitions: vec![Arc::new(InvertedPartition { + id: 0, 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()), + prewarm_state: Arc::new(Mutex::new(InvertedPrewarmState::default())), + document_projections_resident: Arc::new(AtomicBool::new(false)), deleted_fragments: RoaringBitmap::new(), })) } 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`], @@ -1388,6 +1810,8 @@ impl InvertedIndex { format_version, 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, })) } @@ -1606,7 +2030,24 @@ 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) + && (self.is_legacy() || self.document_projections_resident_now()) + { + return Ok(()); + } + let with_position = options.with_position || state.positions_ready; + state.query_ready = false; + state.positions_ready = false; + self.document_projections_resident + .store(false, Ordering::Release); + self.prewarm_query_state(with_position).await?; + state.query_ready = true; + state.positions_ready = 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!( @@ -1640,12 +2081,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, @@ -1662,8 +2099,22 @@ 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(), + )); + } + self.document_projections_resident + .store(true, Ordering::Release); info!( partition_count = self.partitions.len(), + query_ready, elapsed_ms = prewarm_started.elapsed().as_millis() as u64, "fts index prewarm finished" ); @@ -1795,10 +2246,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, } @@ -1839,33 +2289,27 @@ 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?; - // 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 mut inverted_list = PostingListReader::try_new(invert_list_file, index_cache).await?; 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, id, WeakLanceCache::from(index_cache), - num_docs, - false, + docs_reader.as_ref(), frag_reuse_index, // 256-document blocks score with quantized document lengths. inverted_list.block_size() == MAX_POSTING_BLOCK_SIZE, - )); + )?; + inverted_list.modern_num_docs = Some(docs.len()); Ok(Self { id, store, tokens, inverted_list: Arc::new(inverted_list), - docs, + docs: PartitionDocumentStore::Modern(Arc::new(docs)), token_set_format, }) } @@ -1969,7 +2413,7 @@ impl InvertedPartition { doc_ids: &[u32], frequencies: &[u32], block_size: usize, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Vec { @@ -1995,7 +2439,7 @@ impl InvertedPartition { fn union_plain_posting_lists( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Result { @@ -2031,7 +2475,7 @@ impl InvertedPartition { fn union_plain_posting_lists_with_positions( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Result { @@ -2087,7 +2531,7 @@ impl InvertedPartition { fn union_compressed_posting_lists( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Result { @@ -2146,7 +2590,7 @@ impl InvertedPartition { fn union_compressed_posting_lists_with_positions( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Result { @@ -2210,7 +2654,7 @@ impl InvertedPartition { fn union_posting_lists( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, with_positions: bool, query_weight: f32, scorer: &MemBM25Scorer, @@ -2366,7 +2810,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 }; @@ -2405,7 +2854,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( @@ -2463,40 +2912,94 @@ 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> { - if postings.is_empty() { - return Ok(Vec::new()); - } + ) -> Result>> { + let documents = LegacyWandDocuments::new(docs, mask); + self.bm25_search_with_documents( + &documents, + params, + operator, + postings, + impact_scorer, + metrics, + shared_threshold, + ) + } - // 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) - .with_shared_threshold(shared_threshold); - wand.search(params, mask, metrics)? + #[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>> { + 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 scorer = IndexBM25Scorer::new(std::iter::once(self)); - let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) - .with_shared_threshold(shared_threshold); - wand.search(params, mask, metrics)? - }; - Ok(hits) + 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)] + #[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()); + } + + let hits = if let Some(scorer) = impact_scorer { + let mut wand = Wand::new(operator, postings.into_iter(), documents, scorer) + .with_shared_threshold(shared_threshold); + wand.search(params, metrics)? + } else { + let scorer = IndexBM25Scorer::new(std::iter::once(self)); + let mut wand = Wand::new(operator, postings.into_iter(), documents, scorer) + .with_shared_threshold(shared_threshold); + wand.search(params, metrics)? + }; + Ok(hits) } pub async fn into_builder(self) -> Result { @@ -2508,7 +3011,7 @@ impl InvertedPartition { self.inverted_list.block_size(), ); builder.tokens = self.tokens.into_mutable(); - builder.docs = self.docs.owned_docset().await?; + builder.docs = self.docs.load_build_docset().await?; builder .posting_lists @@ -2921,6 +3424,14 @@ 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]>>, + /// Skips per-token readiness checks once the whole immutable table is validated. + modern_postings_validated: AtomicBool, + modern_num_docs: Option, + index_cache: WeakLanceCache, } @@ -2994,7 +3505,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 } } @@ -3028,6 +3548,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(|_| OnceCell::new()) + .collect::>() + .into() + }); Ok(Self { reader, @@ -3038,6 +3564,9 @@ impl PostingListReader { block_size, positions_layout, grouping, + modern_doc_id_validations, + modern_postings_validated: AtomicBool::new(false), + modern_num_docs: None, index_cache: WeakLanceCache::from(index_cache), }) } @@ -3117,10 +3646,6 @@ impl PostingListReader { /// not been loaded yet, and never triggers the bulk load itself. The stats /// path uses this so a single-term `df` lookup costs O(1) bytes rather /// than O(num_unique_tokens). - /// - /// `metrics` is threaded through so callers holding a real - /// `MetricsCollector` (e.g. `MatchQueryExec`) record the per-token - /// `PostingMetadataKey` cache boundary in their per-query counters. pub(crate) async fn posting_len_for_token( &self, token_id: u32, @@ -3181,7 +3706,7 @@ impl PostingListReader { _ => metrics.record_index_cache_miss(), } } - let metadata = result.map(|(v, _)| v)?; + let metadata = result.map(|(value, _)| value)?; Ok((Some(metadata.max_score), Some(metadata.length))) } } @@ -3350,6 +3875,11 @@ impl PostingListReader { } }; + 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 let positions = self.read_positions(token_id, metrics).await?; @@ -3359,6 +3889,87 @@ impl PostingListReader { Ok(posting) } + 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 { + 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() + )) + })?; + validation + .get_or_try_init(|| async { + Self::validate_modern_posting(token_id, posting, num_docs) + }) + .await + .map(|_| ()) + } + + #[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); + }; + 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, + 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(()); + }; + 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 { + 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())); + 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` /// 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 @@ -3646,6 +4257,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()), @@ -3660,7 +4272,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| { @@ -3668,6 +4286,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); @@ -3697,8 +4318,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; @@ -3706,15 +4344,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) }) @@ -3723,7 +4375,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 @@ -6530,78 +7188,11 @@ impl NumTokens { } } -/// `Shared` is a zero-copy view of the partition's `DocRowIdsKey` cache -/// entry; hold it only transiently, or evicting the entry can no longer -/// free the memory. -#[derive(Debug, Clone)] -enum RowIds { - Owned(Vec), - Shared(ScalarBuffer), -} - -impl Default for RowIds { - fn default() -> Self { - Self::Owned(Vec::new()) - } -} - -impl std::ops::Deref for RowIds { - type Target = [u64]; - - fn deref(&self) -> &Self::Target { - match self { - Self::Owned(values) => values, - Self::Shared(values) => values, - } - } -} - -impl DeepSizeOf for RowIds { - fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - match self { - Self::Owned(values) => values.deep_size_of_children(context), - // Weighed by the DocRowIdsKey cache entry. - Self::Shared(_) => 0, - } - } -} - -impl RowIds { - fn with_capacity(capacity: usize) -> Self { - Self::Owned(Vec::with_capacity(capacity)) - } - - fn into_owned(self) -> Vec { - match self { - Self::Owned(values) => values, - Self::Shared(values) => values.to_vec(), - } - } - - fn push(&mut self, value: u64) { - match self { - Self::Owned(values) => values.push(value), - Self::Shared(values) => { - let mut owned = values.to_vec(); - owned.push(value); - *self = Self::Owned(owned); - } - } - } - - fn memory_size(&self) -> usize { - match self { - Self::Owned(values) => values.capacity() * std::mem::size_of::(), - Self::Shared(_) => 0, - } - } -} - // DocSet is a mapping from row ids to the number of tokens in the document // It's used to sort the documents by the bm25 score #[derive(Debug, Clone, Default)] pub struct DocSet { - row_ids: RowIds, + row_ids: Vec, num_tokens: NumTokens, // (row_id, doc_id) pairs sorted by row_id inv: Vec<(u64, u32)>, @@ -6779,7 +7370,7 @@ impl DocSet { total_tokens: u64, ) -> Self { Self { - row_ids: RowIds::default(), + row_ids: Vec::new(), num_tokens: NumTokens::Shared(num_tokens_col.values().clone()), inv: Vec::new(), total_tokens, @@ -6788,45 +7379,8 @@ impl DocSet { } } - /// Per-query view for the masked wand path: this num-tokens-only set - /// plus a transient borrow of the row-ids cache entry. - pub(crate) fn with_shared_row_ids(&self, row_ids: ScalarBuffer) -> Self { - let mut docs = self.clone(); - docs.row_ids = RowIds::Shared(row_ids); - docs - } - - /// Resident scoring set for a modern partition: shared `num_tokens` plus - /// `inv`; no row_ids — the column lives only in the cache entry. - pub(crate) fn from_cached_num_tokens_with_inv( - row_ids_col: &UInt64Array, - num_tokens_col: &arrow_array::UInt32Array, - total_tokens: u64, - ) -> Self { - let row_ids = row_ids_col.values(); - let mut inv: Vec<(u64, u32)> = row_ids - .iter() - .enumerate() - .map(|(doc_id, row_id)| (*row_id, doc_id as u32)) - .collect(); - if !row_ids.is_sorted() { - inv.sort_unstable_by_key(|entry| entry.0); - } - let mut docs = Self::from_cached_num_tokens(num_tokens_col, total_tokens); - docs.inv = inv; - docs - } - - /// True iff `doc_ids()` can answer reverse lookups; flat_search requires - /// this. - pub(crate) fn supports_reverse_lookup(&self) -> bool { - !self.inv.is_empty() || self.has_row_ids() - } - /// Build a `DocSet` from already-loaded `row_id` and `num_tokens` - /// arrow columns. Lets callers that have one column already in hand - /// (e.g. `LazyDocSet` after `total_tokens_num` pre-fetched - /// `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, @@ -6851,7 +7405,7 @@ impl DocSet { let total_tokens = num_tokens.iter().map(|&x| x as u64).sum(); return Ok(Self { - row_ids: RowIds::Owned(row_ids), + row_ids, num_tokens: NumTokens::Owned(num_tokens), inv: Vec::new(), total_tokens, @@ -6894,7 +7448,7 @@ impl DocSet { let total_tokens = num_tokens.iter().map(|&x| x as u64).sum(); return Ok(Self { - row_ids: RowIds::Owned(row_ids), + row_ids, num_tokens: NumTokens::Owned(num_tokens), inv, total_tokens, @@ -6915,7 +7469,7 @@ impl DocSet { } let total_tokens = num_tokens.iter().map(|&x| x as u64).sum(); Ok(Self { - row_ids: RowIds::Owned(row_ids), + row_ids, num_tokens: NumTokens::Owned(num_tokens), inv, total_tokens, @@ -6929,7 +7483,7 @@ impl DocSet { pub fn remap(&mut self, mapping: &RowAddrRemap) -> Vec { let mut removed = Vec::new(); let len = self.len(); - let row_ids = std::mem::replace(&mut self.row_ids, RowIds::with_capacity(len)).into_owned(); + let row_ids = std::mem::replace(&mut self.row_ids, Vec::with_capacity(len)); let num_tokens = std::mem::replace(&mut self.num_tokens, NumTokens::with_capacity(len)).into_owned(); self.invalidate_norms(); @@ -7021,7 +7575,7 @@ impl DocSet { } pub(crate) fn memory_size(&self) -> usize { - self.row_ids.memory_size() + self.row_ids.capacity() * std::mem::size_of::() + self.num_tokens.memory_size() + self.inv.capacity() * std::mem::size_of::<(u64, u32)>() } @@ -7662,7 +8216,7 @@ mod tests { use crate::scalar::inverted::document_tokenizer::DocType; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use futures::stream; - use lance_core::cache::LanceCache; + use lance_core::cache::{LanceCache, QuickCacheBackend}; use lance_core::utils::tempfile::TempObjDir; use lance_io::object_store::ObjectStore; @@ -7692,40 +8246,20 @@ mod tests { use std::sync::Arc; use std::sync::atomic::{AtomicU32, Ordering}; - use crate::scalar::inverted::lazy_docset::DocRowIdsKey; use crate::scalar::inverted::tokenizer::document_tokenizer::TextTokenizer; use lance_tokenizer::{Language, SimpleTokenizer, StopWordFilter, TextAnalyzer}; use super::*; - /// Test helper: resolve any `Pending` candidates a partition-level wand - /// walk emitted, so per-partition results can be asserted on real - /// row_ids. Production resolves only global top-k survivors inside - /// `InvertedIndex::bm25_search`. - async fn resolve_deferred_candidates( - docs: &LazyDocSet, - candidates: &mut [DocCandidate], - ) -> Result<()> { - let pending: Vec = candidates - .iter() - .filter_map(|c| match c.addr { - CandidateAddr::Pending(d) => Some(d), - CandidateAddr::RowId(_) => None, - }) - .collect(); - if pending.is_empty() { - return Ok(()); - } - let mut iter = docs.resolve_row_ids(&pending).await?.into_iter(); - for c in candidates { - if matches!(c.addr, CandidateAddr::Pending(_)) { - let r = iter.next().ok_or_else(|| { - Error::internal("resolve_row_ids returned fewer items than requested") - })?; - c.addr = CandidateAddr::RowId(r); - } - } - Ok(()) + #[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)] @@ -8795,510 +9329,13 @@ mod tests { // Check that we got results from both partitions assert!( - row_ids.contains(&100), - "Should contain row_id from partition 0" - ); - assert!( - row_ids.iter().any(|&id| id >= 200), - "Should contain row_id from partition 1" - ); - } - - /// Counts docs-file ROW_ID reads so the test below can assert the - /// caching contract of deferred row_id resolution: one full-column read - /// per resolving partition on first use, zero scattered single-row reads, - /// and zero additional reads on subsequent queries. - #[derive(Debug, Default)] - struct DocsRowIdReadCounter { - full_column_reads: std::sync::atomic::AtomicUsize, - scattered_reads: std::sync::atomic::AtomicUsize, - } - - #[cfg_attr(coverage, coverage(off))] - impl DocsRowIdReadCounter { - fn full_column_reads(&self) -> usize { - self.full_column_reads - .load(std::sync::atomic::Ordering::Relaxed) - } - fn scattered_reads(&self) -> usize { - self.scattered_reads - .load(std::sync::atomic::Ordering::Relaxed) - } - } - - struct DocsRowIdCountingReader { - inner: Arc, - counter: Arc, - } - - #[cfg_attr(coverage, coverage(off))] - #[async_trait] - impl IndexReader for DocsRowIdCountingReader { - async fn read_record_batch(&self, n: u64, batch_size: u64) -> Result { - self.inner.read_record_batch(n, batch_size).await - } - async fn read_global_buffer(&self, index: u32) -> Result { - self.inner.read_global_buffer(index).await - } - async fn read_range( - &self, - range: std::ops::Range, - projection: Option<&[&str]>, - ) -> Result { - if projection - .map(|cols| cols.contains(&lance_core::ROW_ID)) - .unwrap_or(false) - { - self.counter - .full_column_reads - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - } - self.inner.read_range(range, projection).await - } - async fn read_ranges( - &self, - ranges: &[std::ops::Range], - projection: Option<&[&str]>, - ) -> Result { - if projection - .map(|cols| cols.contains(&lance_core::ROW_ID)) - .unwrap_or(false) - { - self.counter - .scattered_reads - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - } - self.inner.read_ranges(ranges, projection).await - } - async fn num_batches(&self, batch_size: u64) -> u32 { - self.inner.num_batches(batch_size).await - } - fn num_rows(&self) -> usize { - self.inner.num_rows() - } - fn schema(&self) -> &lance_core::datatypes::Schema { - self.inner.schema() - } - } - - #[derive(Debug)] - struct DocsRowIdCountingStore { - inner: Arc, - counter: Arc, - } - - #[cfg_attr(coverage, coverage(off))] - impl DeepSizeOf for DocsRowIdCountingStore { - fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - self.inner.deep_size_of_children(context) - } - } - - #[cfg_attr(coverage, coverage(off))] - #[async_trait] - impl IndexStore for DocsRowIdCountingStore { - fn as_any(&self) -> &dyn std::any::Any { - self - } - fn clone_arc(&self) -> Arc { - Arc::new(Self { - inner: self.inner.clone(), - counter: self.counter.clone(), - }) - } - fn io_parallelism(&self) -> usize { - self.inner.io_parallelism() - } - fn with_io_priority(&self, io_priority: u64) -> Arc { - Arc::new(Self { - inner: self.inner.with_io_priority(io_priority), - counter: self.counter.clone(), - }) - } - async fn new_index_file( - &self, - name: &str, - schema: Arc, - ) -> Result> { - self.inner.new_index_file(name, schema).await - } - async fn open_index_file(&self, name: &str) -> Result> { - let reader = self.inner.open_index_file(name).await?; - if name.ends_with(DOCS_FILE) { - Ok(Arc::new(DocsRowIdCountingReader { - inner: reader, - counter: self.counter.clone(), - })) - } else { - Ok(reader) - } - } - async fn copy_index_file( - &self, - name: &str, - dest_store: &dyn IndexStore, - ) -> Result { - self.inner.copy_index_file(name, dest_store).await - } - async fn copy_index_file_to( - &self, - name: &str, - new_name: &str, - dest_store: &dyn IndexStore, - ) -> Result { - self.inner - .copy_index_file_to(name, new_name, dest_store) - .await - } - async fn rename_index_file( - &self, - name: &str, - new_name: &str, - ) -> Result { - self.inner.rename_index_file(name, new_name).await - } - async fn delete_index_file(&self, name: &str) -> Result<()> { - self.inner.delete_index_file(name).await - } - async fn list_files_with_sizes(&self) -> Result> { - self.inner.list_files_with_sizes().await - } - } - - /// Matching-partition count used by `write_many_partition_index`. - const MANY_PARTITIONS: u64 = 40; - - /// Writes `MANY_PARTITIONS` partitions whose only token is "pipeline" - /// (1-3 docs each), one extra partition whose only token does not match, - /// and the metadata file. Returns every matching doc's row_id. - async fn write_many_partition_index(store: &LanceIndexStore) -> Vec { - let mut expected_row_ids: Vec = Vec::with_capacity((MANY_PARTITIONS * 3) as usize); - for pid in 0..MANY_PARTITIONS { - let mut builder = InnerBuilder::new(pid, false, TokenSetFormat::default()); - builder.tokens.add("pipeline".to_owned()); - builder.posting_lists.push(PostingListBuilder::new(false)); - let ndocs = (pid % 3) + 1; - for d in 0..ndocs { - builder.posting_lists[0].add(d as u32, PositionRecorder::Count(1)); - let row_id = pid * 1000 + d; - builder.docs.append(row_id, 1); - expected_row_ids.push(row_id); - } - builder.write(store).await.unwrap(); - } - // A partition whose only token does NOT match the query. - let mut empty_builder = - InnerBuilder::new(MANY_PARTITIONS, false, TokenSetFormat::default()); - empty_builder.tokens.add("unrelated".to_owned()); - empty_builder - .posting_lists - .push(PostingListBuilder::new(false)); - empty_builder.posting_lists[0].add(0, PositionRecorder::Count(1)); - empty_builder.docs.append(999_999, 1); - empty_builder.write(store).await.unwrap(); - - let all_partitions: Vec = (0..=MANY_PARTITIONS).collect(); - let metadata = std::collections::HashMap::from_iter(vec![ - ( - "partitions".to_owned(), - serde_json::to_string(&all_partitions).unwrap(), - ), - ( - "params".to_owned(), - serde_json::to_string(&InvertedIndexParams::default()).unwrap(), - ), - ( - TOKEN_SET_FORMAT_KEY.to_owned(), - TokenSetFormat::default().to_string(), - ), - ]); - let mut writer = store - .new_index_file(METADATA_FILE, Arc::new(arrow_schema::Schema::empty())) - .await - .unwrap(); - writer.finish_with_metadata(metadata).await.unwrap(); - expected_row_ids - } - - #[tokio::test] - async fn test_bm25_search_many_partitions_resolves_exact_row_ids() { - // Regression test for deferred row_id resolution: with many partitions - // (including one whose only token does not match the query), a search - // must resolve every candidate's row_id exactly once and correctly. - // Also asserts the caching contract: each resolving partition reads the - // ROW_ID column exactly once (one full-column read, no scattered - // single-row reads), and subsequent queries perform no further reads. - let tmpdir = TempObjDir::default(); - let store = Arc::new(LanceIndexStore::new( - ObjectStore::local().into(), - tmpdir.clone(), - Arc::new(LanceCache::no_cache()), - )); - let expected_row_ids = write_many_partition_index(store.as_ref()).await; - - // Load through a store that counts docs-file ROW_ID reads. - let counter = Arc::new(DocsRowIdReadCounter::default()); - let inner_store: Arc = store.clone(); - let counting_store: Arc = Arc::new(DocsRowIdCountingStore { - inner: inner_store, - counter: counter.clone(), - }); - let cache = Arc::new(LanceCache::with_capacity(64 * 1024 * 1024)); - let index = InvertedIndex::load(counting_store, None, cache.as_ref()) - .await - .unwrap(); - assert_eq!(index.partitions.len(), 41); - assert_eq!(counter.full_column_reads(), 0); - assert_eq!(counter.scattered_reads(), 0); - - // Every matching doc must come back exactly once with a correct row_id. - let tokens = Arc::new(Tokens::new(vec!["pipeline".to_owned()], DocType::Text)); - let params = Arc::new(FtsSearchParams::new().with_limit(Some(1000))); - let prefilter = Arc::new(NoFilter); - let metrics = Arc::new(NoOpMetricsCollector); - let (row_ids, scores) = index - .bm25_search( - tokens.clone(), - params, - Operator::Or, - prefilter.clone(), - metrics.clone(), - None, - ) - .await - .unwrap(); - let mut got = row_ids.clone(); - got.sort_unstable(); - let mut want = expected_row_ids.clone(); - want.sort_unstable(); - assert_eq!(got, want, "every matching doc resolved exactly once"); - assert_eq!(scores.len(), row_ids.len()); - assert!(scores.iter().all(|s| *s > 0.0)); - - // Caching contract: one full-column ROW_ID read per resolving - // partition, no scattered single-row reads. The non-matching - // partition must not read its ROW_ID column at all. - assert_eq!( - counter.full_column_reads(), - MANY_PARTITIONS as usize, - "each resolving partition reads the ROW_ID column exactly once" - ); - assert_eq!( - counter.scattered_reads(), - 0, - "resolution must not issue scattered single-row reads" - ); - - // Top-k smaller than the total: exactly k results, and the cached - // columns serve resolution with no further docs-file reads. - let params_k = Arc::new(FtsSearchParams::new().with_limit(Some(7))); - let (row_ids_k, scores_k) = index - .bm25_search(tokens, params_k, Operator::Or, prefilter, metrics, None) - .await - .unwrap(); - assert_eq!(row_ids_k.len(), 7); - assert_eq!(scores_k.len(), 7); - assert!(row_ids_k.iter().all(|id| expected_row_ids.contains(id))); - assert_eq!( - counter.full_column_reads(), - MANY_PARTITIONS as usize, - "subsequent queries must not re-read the ROW_ID column" - ); - assert_eq!(counter.scattered_reads(), 0); - } - - #[tokio::test] - async fn test_bm25_search_resolves_only_topk_survivors_and_accounts_cache() { - // Late resolution: row_ids are resolved only for candidates that - // survive the global top-k, so with k much smaller than the partition - // count, only the partitions holding survivors load their ROW_ID - // column. Each loaded column must be visible to the index cache as - // its own weighed entry. - let tmpdir = TempObjDir::default(); - let store = Arc::new(LanceIndexStore::new( - ObjectStore::local().into(), - tmpdir.clone(), - Arc::new(LanceCache::no_cache()), - )); - let expected_row_ids = write_many_partition_index(store.as_ref()).await; - - let counter = Arc::new(DocsRowIdReadCounter::default()); - let inner_store: Arc = store.clone(); - let counting_store: Arc = Arc::new(DocsRowIdCountingStore { - inner: inner_store, - counter: counter.clone(), - }); - let cache = Arc::new(LanceCache::with_capacity(64 * 1024 * 1024)); - let index = InvertedIndex::load(counting_store, None, cache.as_ref()) - .await - .unwrap(); - - // k=3 with 40 matching partitions: survivors span at most 3 - // partitions, so at most 3 ROW_ID columns are read even though every - // matching partition produced candidates. - let tokens = Arc::new(Tokens::new(vec!["pipeline".to_owned()], DocType::Text)); - let params = Arc::new(FtsSearchParams::new().with_limit(Some(3))); - let (row_ids, scores) = index - .bm25_search( - tokens, - params, - Operator::Or, - Arc::new(NoFilter), - Arc::new(NoOpMetricsCollector), - None, - ) - .await - .unwrap(); - assert_eq!(row_ids.len(), 3); - assert_eq!(scores.len(), 3); - assert!(row_ids.iter().all(|id| expected_row_ids.contains(id))); - let reads = counter.full_column_reads(); - assert!( - (1..=3).contains(&reads), - "only partitions with surviving candidates load their ROW_ID column, got {reads}" - ); - assert_eq!(counter.scattered_reads(), 0); - - // Accounting: every loaded column is present in the index cache as a - // DocRowIds entry (weighed at insert, evictable under pressure). - // Partition caches are scoped with a `part-{id}` prefix on load. - let mut cached_entries = 0; - for partition_id in 0..=MANY_PARTITIONS { - if cache - .with_key_prefix(format!("part-{partition_id}").as_str()) - .get_with_key(&DocRowIdsKey { partition_id }) - .await - .is_some() - { - cached_entries += 1; - } - } - assert_eq!( - cached_entries, reads, - "each loaded ROW_ID column is its own index-cache entry" - ); - } - - #[tokio::test] - async fn test_prewarm_fills_row_ids_cache_entry_without_a_second_copy() { - let tmpdir = TempObjDir::default(); - let store = Arc::new(LanceIndexStore::new( - ObjectStore::local().into(), - tmpdir.clone(), - Arc::new(LanceCache::no_cache()), - )); - let expected_row_ids = write_many_partition_index(store.as_ref()).await; - - let counter = Arc::new(DocsRowIdReadCounter::default()); - let inner_store: Arc = store.clone(); - let counting_store: Arc = Arc::new(DocsRowIdCountingStore { - inner: inner_store, - counter: counter.clone(), - }); - let cache = Arc::new(LanceCache::with_capacity(64 * 1024 * 1024)); - let index = InvertedIndex::load(counting_store, None, cache.as_ref()) - .await - .unwrap(); - - index - .prewarm_with_options(&FtsPrewarmOptions::default()) - .await - .unwrap(); - let total_partitions = (MANY_PARTITIONS + 1) as usize; - assert_eq!(counter.full_column_reads(), total_partitions); - assert_eq!(counter.scattered_reads(), 0); - - for partition_id in 0..=MANY_PARTITIONS { - assert!( - cache - .with_key_prefix(format!("part-{partition_id}").as_str()) - .get_with_key(&DocRowIdsKey { partition_id }) - .await - .is_some(), - "prewarm must fill the DocRowIds entry for partition {partition_id}" - ); - } - for part in &index.partitions { - let resident = part.docs.ensure_loaded().await.unwrap(); - assert!(!resident.has_row_ids()); - assert!(resident.supports_reverse_lookup()); - } - - let tokens = Arc::new(Tokens::new(vec!["pipeline".to_owned()], DocType::Text)); - let params = Arc::new(FtsSearchParams::new().with_limit(Some(1000))); - let (row_ids, _) = index - .bm25_search( - tokens, - params, - Operator::Or, - Arc::new(NoFilter), - Arc::new(NoOpMetricsCollector), - None, - ) - .await - .unwrap(); - let mut got = row_ids; - got.sort_unstable(); - let mut want = expected_row_ids; - want.sort_unstable(); - assert_eq!(got, want); - assert_eq!(counter.full_column_reads(), total_partitions); - assert_eq!(counter.scattered_reads(), 0); - } - - #[tokio::test] - async fn test_row_ids_resolution_reloads_after_eviction() { - // no_cache retains nothing — the always-evicted worst case: every - // query must reload the column and stay correct. - let tmpdir = TempObjDir::default(); - let store = Arc::new(LanceIndexStore::new( - ObjectStore::local().into(), - tmpdir.clone(), - Arc::new(LanceCache::no_cache()), - )); - let expected_row_ids = write_many_partition_index(store.as_ref()).await; - - let counter = Arc::new(DocsRowIdReadCounter::default()); - let inner_store: Arc = store.clone(); - let counting_store: Arc = Arc::new(DocsRowIdCountingStore { - inner: inner_store, - counter: counter.clone(), - }); - let cache = LanceCache::no_cache(); - let index = InvertedIndex::load(counting_store, None, &cache) - .await - .unwrap(); - - let tokens = Arc::new(Tokens::new(vec!["pipeline".to_owned()], DocType::Text)); - let params = Arc::new(FtsSearchParams::new().with_limit(Some(1000))); - let mut want = expected_row_ids; - want.sort_unstable(); - let mut reads_after_first = 0; - for round in 0..2 { - let (row_ids, _) = index - .bm25_search( - tokens.clone(), - params.clone(), - Operator::Or, - Arc::new(NoFilter), - Arc::new(NoOpMetricsCollector), - None, - ) - .await - .unwrap(); - let mut got = row_ids; - got.sort_unstable(); - assert_eq!(got, want, "round {round} must stay correct"); - if round == 0 { - reads_after_first = counter.full_column_reads(); - assert!(reads_after_first > 0); - } - } + row_ids.contains(&100), + "Should contain row_id from partition 0" + ); assert!( - counter.full_column_reads() > reads_after_first, - "with nothing retained, the next query reloads the column" + row_ids.iter().any(|&id| id >= 200), + "Should contain row_id from partition 1" ); - assert_eq!(counter.scattered_reads(), 0); } #[tokio::test] @@ -10144,21 +10181,8 @@ mod tests { ); } - /// Guards the review fix that threads `Option<&dyn MetricsCollector>` - /// through `bm25_stats_for_terms → df_for_term → posting_len_for_token → - /// posting_metadata_for_token`. Cold stats for `N` tokens on one - /// partition must record exactly `N` misses (one per `PostingMetadataKey`) - /// and zero hits; a second call with the cache warm must flip that to - /// zero misses and `N` hits. If any hop in the chain drops the collector - /// this test regresses to `0/0` on both calls. #[tokio::test] async fn test_bm25_stats_for_terms_records_metadata_cache_stats() { - // Keep a live `LanceCache` clone in test scope so the - // `WeakLanceCache` inside `PostingListReader` can still upgrade after - // `load_counted_v2_index` returns. Without this the weak reference - // would collapse and every `posting_metadata_for_token` call would - // silently fall through to the "cache no longer available" path, - // making both cold and warm calls report identical miss counts. let cache = LanceCache::with_capacity(1024 * 1024); let (index, _counter, _tmpdir) = load_counted_v2_index(100, cache.clone()).await; assert!( @@ -10167,18 +10191,13 @@ mod tests { ); let terms = ["t0".to_string(), "t1".to_string(), "t2".to_string()]; - let cold = LocalMetricsCollector::default(); let cold_stats = index .bm25_stats_for_terms(&terms, Some(&cold)) .await .unwrap(); assert_eq!(cold_stats.2, vec![1, 1, 1]); - assert_eq!( - cold.index_cache_misses(), - terms.len(), - "expected one miss per (term, partition) on cold", - ); + assert_eq!(cold.index_cache_misses(), terms.len()); assert_eq!(cold.index_cache_hits(), 0); let warm = LocalMetricsCollector::default(); @@ -10188,11 +10207,7 @@ mod tests { .unwrap(); assert_eq!(warm_stats, cold_stats); assert_eq!(warm.index_cache_misses(), 0); - assert_eq!( - warm.index_cache_hits(), - terms.len(), - "expected one hit per (term, partition) on warm", - ); + assert_eq!(warm.index_cache_hits(), terms.len()); } #[tokio::test] @@ -10209,77 +10224,94 @@ 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(Operator::Or, &all_rows) - .await - .unwrap(); - assert!(Arc::ptr_eq(first, &wand_view)); + assert!(matches!( + documents + .visibility(Arc::new(all_rows), false) + .await + .unwrap(), + DocVisibility::All + )); + assert!(!documents.projection_loaded()); - // Flat-shaped mask (OR + tiny allow-list) → resident set with `inv`. let filtered = RowAddrMask::allow_nothing(); - let resident = partition - .docs - .docs_for_wand(Operator::Or, &filtered) + let visibility = documents + .visibility(Arc::new(filtered), true) .await .unwrap(); - assert!(!resident.has_row_ids()); - assert!(resident.supports_reverse_lookup()); - assert!(matches!(&resident.num_tokens, NumTokens::Shared(_))); - assert_eq!(resident.total_tokens_num(), 100); + assert!(visibility.is_empty()); + assert!(!documents.projection_loaded()); + assert_eq!( + documents + .resolve_addresses(&[DocId::new(0), DocId::new(99)]) + .await + .unwrap(), + [0, 99] + ); + } - // Masked non-flat walk (AND is never flat) → per-query view. - let masked_view = partition - .docs - .docs_for_wand(Operator::And, &filtered) + #[tokio::test] + 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 documents = index.partitions[0].docs.modern().unwrap(); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); + + 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(); - assert!(masked_view.has_row_ids()); - assert!(matches!(&masked_view.row_ids, RowIds::Shared(_))); - assert_eq!(masked_view.total_tokens_num(), 100); - assert_eq!( - partition.docs.resolve_row_ids(&[0, 99]).await.unwrap(), - [0, 99] - ); + assert!(row_ids.is_empty()); + assert!(scores.is_empty()); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); } #[tokio::test] - async fn test_concurrent_total_and_num_tokens_view_initialization() { + 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.clone(); + let docs = index.partitions[0].docs.modern().unwrap().clone(); - 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 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 totals = totals.into_iter().collect::>>().unwrap(); - assert_eq!(totals, vec![100; 8]); + 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] @@ -10651,7 +10683,9 @@ mod tests { .unwrap(); writer.finish_with_metadata(metadata).await.unwrap(); - let cache = Arc::new(LanceCache::with_capacity(4096)); + let cache = Arc::new(LanceCache::with_backend(Arc::new( + QuickCacheBackend::with_capacity(4096), + ))); let index = InvertedIndex::load(store.clone(), None, cache.as_ref()) .await .unwrap(); @@ -10700,6 +10734,31 @@ mod tests { ), "positions should be stored in the dedicated position cache" ); + + drop(positions); + drop(group); + cache.clear().await; + assert!( + inverted_list + .index_cache + .get_with_key(&PositionKey { token_id: 0 }) + .await + .is_none() + ); + + index + .prewarm_with_options(&FtsPrewarmOptions::default()) + .await + .unwrap(); + assert!(index.prewarm_state.lock().await.satisfies(true)); + assert!( + inverted_list + .index_cache + .get_with_key(&PositionKey { token_id: 0 }) + .await + .is_some(), + "re-prewarm after eviction must preserve the strongest requested mode" + ); } #[tokio::test] @@ -10975,7 +11034,7 @@ mod tests { async fn load_global_scoring_test_index( first_partition_has_impacts: bool, second_partition_has_impacts: bool, - ) -> (TempObjDir, Arc) { + ) -> (TempObjDir, Arc, Arc) { let tmpdir = TempObjDir::default(); let store = Arc::new(LanceIndexStore::new( ObjectStore::local().into(), @@ -11024,9 +11083,256 @@ mod tests { } write_test_metadata(&store, vec![0, 1], InvertedIndexParams::default()).await; - let cache = LanceCache::with_capacity(4096); - let index = InvertedIndex::load(store, None, &cache).await.unwrap(); - (tmpdir, index) + let cache = Arc::new(LanceCache::with_backend(Arc::new( + QuickCacheBackend::with_capacity(4096), + ))); + let index = InvertedIndex::load(store, None, cache.as_ref()) + .await + .unwrap(); + (tmpdir, cache, index) + } + + #[tokio::test] + async fn test_chunked_modern_search_preserves_cold_and_prewarmed_results() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + let matching_partitions = 17_u64; + for partition_id in 0..matching_partitions { + let mut builder = InnerBuilder::new(partition_id, false, TokenSetFormat::default()); + builder.tokens.add("pipeline".to_owned()); + builder.posting_lists.push(PostingListBuilder::new(false)); + builder.posting_lists[0].add(0, PositionRecorder::Count(1)); + builder.docs.append(partition_id * 1_000 + 7, 1); + builder.write(store.as_ref()).await.unwrap(); + } + let unmatched_partition = matching_partitions; + let mut builder = InnerBuilder::new(unmatched_partition, false, TokenSetFormat::default()); + builder.tokens.add("unrelated".to_owned()); + builder.posting_lists.push(PostingListBuilder::new(false)); + builder.posting_lists[0].add(0, PositionRecorder::Count(1)); + builder.docs.append(999_999, 1); + builder.write(store.as_ref()).await.unwrap(); + + write_test_metadata( + &store, + (0..=unmatched_partition).collect(), + InvertedIndexParams::default(), + ) + .await; + let cache = Arc::new(LanceCache::with_capacity(64 * 1024 * 1024)); + let index = InvertedIndex::load(store, None, cache.as_ref()) + .await + .unwrap(); + let tokens = Arc::new(Tokens::new(vec!["pipeline".to_owned()], DocType::Text)); + let params = + Arc::new(FtsSearchParams::new().with_limit(Some(matching_partitions as usize))); + + let search = || { + index.bm25_search( + tokens.clone(), + params.clone(), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + }; + let (mut cold_row_ids, cold_scores) = search().await.unwrap(); + cold_row_ids.sort_unstable(); + let expected = (0..matching_partitions) + .map(|partition_id| partition_id * 1_000 + 7) + .collect::>(); + assert_eq!(cold_row_ids, expected); + assert_eq!(cold_scores.len(), expected.len()); + + index + .prewarm_with_options(&FtsPrewarmOptions::default()) + .await + .unwrap(); + let (mut prewarmed_row_ids, prewarmed_scores) = search().await.unwrap(); + prewarmed_row_ids.sort_unstable(); + assert_eq!(prewarmed_row_ids, expected); + assert_eq!(prewarmed_scores, cold_scores); + } + + #[tokio::test] + async fn test_prewarmed_modern_search_uses_resident_address_projection() { + let (_tmpdir, cache, index) = load_global_scoring_test_index(true, 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()); + + 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.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() + && partition.inverted_list.modern_posting_validation_ready() + })); + assert!(index.prewarm_state.lock().await.satisfies(false)); + + let resident = index + .bm25_search( + tokens.clone(), + params.clone(), + 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)); + + cache.clear().await; + assert!(index.document_projections_resident.load(Ordering::Acquire)); + assert_eq!(cache.size().await, 0); + let resident_address_owners = index + .partitions + .iter() + .map(|partition| { + partition + .docs + .modern() + .unwrap() + .address_buffer_handle() + .strong_count() + }) + .collect::>(); + assert_eq!(resident_address_owners, vec![0, 0]); + assert!( + index + .partitions + .iter() + .all(|partition| { !partition.docs.modern().unwrap().projection_resident() }) + ); + + let after_eviction = index + .bm25_search( + tokens.clone(), + params.clone(), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + assert_eq!(after_eviction, deferred); + assert!(!index.document_projections_resident.load(Ordering::Acquire)); + + cache.clear().await; + index.prewarm_with_options(&prewarm_options).await.unwrap(); + assert!(index.document_projections_resident_now()); + assert!(index.document_projections_resident.load(Ordering::Acquire)); + + let re_prewarms_after_eviction = index + .bm25_search( + tokens, + params, + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + assert_eq!(re_prewarms_after_eviction, deferred); + } + + #[tokio::test] + async fn test_resident_modern_search_loads_partition_stats_without_global_stats() { + let (_tmpdir, _cache, index) = load_global_scoring_test_index(true, 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() + .address_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( @@ -11035,7 +11341,7 @@ mod tests { params: &FtsSearchParams, scorer: Arc, shared_threshold: Arc, - ) -> Vec { + ) -> Vec> { let LoadedPostings { postings, grouped_expansions, @@ -11055,28 +11361,21 @@ mod tests { assert!(!exact_scoring_required); assert!(grouped_expansions.is_empty()); - let mask = NoFilter.mask(); - let docs_for_wand = partition - .docs - .docs_for_wand(Operator::Or, mask.as_ref()) - .await - .unwrap(); - let mut candidates = partition - .bm25_search( - docs_for_wand.as_ref(), + let documents = partition.docs.modern().unwrap(); + let lengths = documents.lengths().await.unwrap(); + let visibility = documents.visibility(NoFilter.mask(), false).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] @@ -11084,7 +11383,7 @@ mod tests { // Partition 0 wins under its local corpus statistics but loses under // the global statistics. If its local score escapes into the shared // floor, partition 1 will incorrectly prune the real global winner. - let (_tmpdir, index) = load_global_scoring_test_index(true, true).await; + let (_tmpdir, _cache, index) = load_global_scoring_test_index(true, true).await; let first_partition = index .partitions .iter() @@ -11134,10 +11433,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)); @@ -11155,10 +11451,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!( @@ -11187,7 +11480,7 @@ mod tests { #[tokio::test] async fn test_mixed_impact_and_legacy_partitions_use_global_final_scores() { - let (_tmpdir, index) = load_global_scoring_test_index(true, false).await; + let (_tmpdir, _cache, index) = load_global_scoring_test_index(true, false).await; let impact_partition = index .partitions @@ -11246,15 +11539,9 @@ mod tests { #[tokio::test] async fn test_two_legacy_partitions_keep_private_thresholds() { - // Both partitions are legacy (no impacts), so their BM25 statistics - // are partition-local scale. Partition 0's matching doc scores high - // under its own statistics while partition 1's matching doc (the true - // global winner, row 200) scores low under partition 1's local - // statistics. The chunked search runs both partitions sequentially in - // one chunk: if partition 0's local k-th score leaked into a shared - // threshold, partition 1 would prune row 200 before global rescoring - // and return the wrong winner. - let (_tmpdir, index) = load_global_scoring_test_index(false, false).await; + // Legacy BM25 scores use partition-local statistics, so sharing one + // pruning floor across partitions can discard the global winner. + let (_tmpdir, _cache, index) = load_global_scoring_test_index(false, false).await; for partition in index.partitions.iter() { let posting = partition .inverted_list @@ -12962,6 +13249,65 @@ 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.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)); + 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) + .await + .unwrap_err(); + assert!(matches!(error, Error::Index { .. })); + 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) + .await + .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) + .await + .unwrap(); + assert_eq!(posting_entries(&second), vec![(0, 1)]); + assert!(validation.get().is_some()); + } + /// Runtime synthetic grouping must return correct posting lists for every /// token, including across synthetic group boundaries. #[tokio::test] @@ -12986,7 +13332,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, @@ -13039,7 +13386,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, @@ -13052,6 +13400,12 @@ mod tests { .prewarm_posting_lists(false, 2) .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(); 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 faf87dd6bac..00000000000 --- a/rust/lance-index/src/scalar/inverted/lazy_docset.rs +++ /dev/null @@ -1,510 +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`. -//! -//! For modern partitions the `doc_id -> row_id` column has exactly one -//! home: the [`DocRowIdsKey`] index-cache entry. No `DocSet` keeps a -//! long-lived copy or reference, so evicting the entry really frees the -//! memory; borrowers reload it single-flight. Legacy and frag-reuse -//! partitions rewrite the mapping at load and keep a private owned copy. - -use std::borrow::Cow; -use std::sync::Arc; - -use arrow::array::AsArray; -use arrow::datatypes::{UInt32Type, UInt64Type}; -use arrow_array::{Array, UInt32Array, UInt64Array}; -use lance_core::ROW_ID; -use lance_core::Result; -use lance_core::cache::{CacheKey, WeakLanceCache}; -use lance_core::deepsize::DeepSizeOf; -use tokio::sync::OnceCell; - -use crate::scalar::RowIdRemapper; -use crate::scalar::inverted::index::{DocSet, NUM_TOKEN_COL}; -use crate::scalar::inverted::query::Operator; -use crate::scalar::inverted::wand::should_flat_search; -use crate::scalar::{IndexReader, IndexStore}; -use lance_select::mask::RowAddrMask; - -/// 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, -} - -/// A partition's full `doc_id -> row_id` column, stored as its own -/// index-cache entry so the cache weighs it at insert time and can -/// evict it independently of the partition object. Keeping it out of -/// the partition's lazily-populated fields keeps the memory visible -/// to capacity accounting. -#[derive(Debug)] -pub struct CachedDocRowIds { - pub row_ids: Arc, -} - -impl DeepSizeOf for CachedDocRowIds { - fn deep_size_of_children(&self, _ctx: &mut lance_core::deepsize::Context) -> usize { - self.row_ids.len() * std::mem::size_of::() - } -} - -/// Cache key for [`CachedDocRowIds`], scoped per partition. -#[derive(Debug, Clone)] -pub struct DocRowIdsKey { - pub partition_id: u64, -} - -impl CacheKey for DocRowIdsKey { - type ValueType = CachedDocRowIds; - - fn key(&self) -> Cow<'_, str> { - format!("doc-row-ids-{}", self.partition_id).into() - } - - fn type_name() -> &'static str { - "DocRowIds" - } -} - -/// Store-backed DocSet view that loads on demand and caches. -/// -/// Holds the [`IndexStore`] and docs-file path rather than an open -/// [`IndexReader`], so a cached partition does not pin a docs-file -/// handle for its whole lifetime. The reader is re-opened on demand -/// inside each column accessor and dropped when that read completes; -/// because the resulting buffers are cached (num_tokens and the resident -/// scoring set in the `OnceCell`s below, the row_id column as its own -/// index-cache entry), a contributing partition re-opens only on a -/// cold miss, and a partition that never scores never opens the docs -/// file at all after construction. -pub struct DeferredDocSet { - store: Arc, - docs_path: String, - /// Scopes this partition's [`DocRowIdsKey`] in the index cache. - partition_id: u64, - /// The index cache holding the row_id column entry. - index_cache: WeakLanceCache, - is_legacy: bool, - frag_reuse_index: Option>, - /// 256-document-block partitions score with quantized document lengths; the - /// 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, - /// Modern partitions: num_tokens + `inv`, no row_ids (the column stays - /// in the cache entry). Legacy / frag reuse: the full rewritten DocSet. - resident: OnceCell>, -} - -impl std::fmt::Debug for LazyDocSet { - 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.resident.initialized()), - ) - .field("num_tokens_loaded", &d.num_tokens.initialized()) - .field("resident_loaded", &d.resident.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.resident - .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) - // The row_id column is not a field here: it lives in the - // index cache as its own weighed entry (`DocRowIdsKey`). - } - } - } -} - -impl LazyDocSet { - #[allow(clippy::too_many_arguments)] - pub fn new( - store: Arc, - docs_path: String, - partition_id: u64, - index_cache: WeakLanceCache, - num_rows: usize, - is_legacy: bool, - frag_reuse_index: Option>, - quantized_scoring: bool, - ) -> Self { - Self::Deferred(Box::new(DeferredDocSet { - store, - docs_path, - partition_id, - index_cache, - is_legacy, - frag_reuse_index, - quantized_scoring, - num_rows, - num_tokens: OnceCell::new(), - resident: 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 - .resident - .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, - } - } - - /// Make the scoring state resident: num_tokens, `inv`, and the row-ids - /// column loaded into its [`DocRowIdsKey`] cache entry. The returned - /// DocSet carries no row_ids for modern partitions, so nothing pins the - /// entry. Prewarm calls this. - pub async fn ensure_loaded(&self) -> Result> { - match self { - Self::Loaded(l) => Ok(l.docs.clone()), - Self::Deferred(d) => d.ensure_loaded().await, - } - } - - /// A fully-owned DocSet, row_ids included, for rebuild paths that - /// mutate it. Never stashed, so it does not pin the cache entry. - pub async fn owned_docset(&self) -> Result { - match self { - Self::Loaded(l) => Ok((*l.docs).clone()), - Self::Deferred(d) => { - if d.is_legacy || d.frag_reuse_index.is_some() { - return Ok((*d.ensure_loaded().await?).clone()); - } - let (snapshot, row_ids) = - futures::try_join!(d.num_tokens_snapshot(), d.row_ids_column())?; - let mut docs = - DocSet::from_columns(row_ids.as_ref(), snapshot.column.as_ref(), false, None)?; - docs.set_quantized_scoring(d.quantized_scoring); - Ok(docs) - } - } - } - - /// Materialize a DocSet that carries num_tokens but no row_ids. - /// Used by the deferred-row_id scoring path; the per-partition - /// caller resolves surviving doc_ids -> row_ids post-wand via - /// [`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 DocSet shape for a wand walk: trivial mask → num_tokens - /// only; legacy / frag reuse → owned rewritten DocSet; flat-shaped mask - /// (same [`should_flat_search`] predicate wand uses) → resident set with - /// `inv`; any other mask → per-query view borrowing the cache entry. - pub async fn docs_for_wand( - &self, - operator: Operator, - mask: &RowAddrMask, - ) -> Result> { - match self { - Self::Loaded(l) => Ok(l.docs.clone()), - Self::Deferred(d) => { - if mask.is_select_all() && !self.has_frag_reuse_remap() { - return self.ensure_num_tokens_loaded().await; - } - if d.is_legacy - || d.frag_reuse_index.is_some() - || should_flat_search(operator, mask, d.num_rows as u64) - { - return self.ensure_loaded().await; - } - let (snapshot_docs, row_ids) = { - let (snapshot, row_ids) = - futures::try_join!(d.num_tokens_snapshot(), d.row_ids_column())?; - (snapshot.docs.clone(), row_ids) - }; - Ok(Arc::new( - snapshot_docs.with_shared_row_ids(row_ids.values().clone()), - )) - } - } - } - - /// 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(resident) = self.resident.get() { - return Ok(resident.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 - } - - /// The full `ROW_ID` column, loaded through the index cache as its own - /// weighed entry ([`DocRowIdsKey`]) rather than a field on this struct, - /// so the memory is accounted for at insert time and evictable under - /// pressure. Concurrent loads are deduped by the cache backend. - async fn row_ids_column(&self) -> Result> { - let store = self.store.clone(); - let docs_path = self.docs_path.clone(); - let num_rows = self.num_rows; - let cached = self - .index_cache - .get_or_insert_with_key( - DocRowIdsKey { - partition_id: self.partition_id, - }, - || async move { - let reader = store.open_index_file(&docs_path).await?; - let batch = reader.read_range(0..num_rows, Some(&[ROW_ID])).await?; - Result::Ok(CachedDocRowIds { - row_ids: Arc::new(batch[ROW_ID].as_primitive::().clone()), - }) - }, - ) - .await?; - Ok(cached.row_ids.clone()) - } - - async fn ensure_loaded(&self) -> Result> { - let docs = self - .resident - .get_or_try_init(|| async { - let mut docs = if self.is_legacy || self.frag_reuse_index.is_some() { - // Both rewrite the mapping, so keep a private owned copy. - DocSet::load( - self.reader().await?, - self.is_legacy, - self.frag_reuse_index.clone(), - ) - .await? - } else { - let (snapshot, row_ids) = - futures::try_join!(self.num_tokens_snapshot(), self.row_ids_column())?; - DocSet::from_cached_num_tokens_with_inv( - row_ids.as_ref(), - snapshot.column.as_ref(), - snapshot.docs.total_tokens_num(), - ) - }; - docs.set_quantized_scoring(self.quantized_scoring); - Result::Ok(Arc::new(docs)) - }) - .await? - .clone(); - Ok(docs) - } - - async fn ensure_num_tokens_loaded(&self) -> Result> { - if let Some(resident) = self.resident.get() { - return Ok(resident.clone()); - } - Ok(self.num_tokens_snapshot().await?.docs.clone()) - } - - async fn resolve_row_ids(&self, doc_ids: &[u32]) -> Result> { - // Only legacy / frag-reuse residents carry (rewritten) row_ids; - // modern partitions read the cache entry, keeping it LRU-hot. - if let Some(resident) = self.resident.get() - && resident.has_row_ids() - { - return Ok(doc_ids.iter().map(|&d| resident.row_id(d)).collect()); - } - let arr = self.row_ids_column().await?; - Ok(doc_ids.iter().map(|&d| arr.value(d as usize)).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_resident_docset_serves_num_tokens_reads() { - let temp_dir = TempObjDir::default(); - let cache = Arc::new(LanceCache::no_cache()); - let store = Arc::new(LanceIndexStore::new( - ObjectStore::local().into(), - temp_dir.clone(), - cache.clone(), - )); - let docs = LazyDocSet::new( - store, - "unused".to_owned(), - 0, - WeakLanceCache::from(&cache), - 3, - false, - None, - false, - ); - assert_eq!(docs.total_tokens_cached(), None); - - let row_ids = UInt64Array::from(vec![10, 20, 30]); - 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.resident.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 f730d6f9d4b..dcd8f5ac0d0 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -21,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, @@ -167,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, @@ -180,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, @@ -202,17 +203,6 @@ pub static FLAT_SEARCH_PERCENT_THRESHOLD: LazyLock = LazyLock::new(|| { .parse::() .unwrap_or(10) }); - -/// Single shared definition of the flat-search predicate: [`Wand::search`] -/// and `LazyDocSet::docs_for_wand` MUST agree — a masked query scored -/// without row_ids silently skips the `mask.selected` filter. -pub(super) fn should_flat_search(operator: Operator, mask: &RowAddrMask, num_docs: u64) -> bool { - operator == Operator::Or - && mask.iter_addrs().is_some() - && mask.max_len().is_some_and(|num_rows_matched| { - num_rows_matched * 100 <= *FLAT_SEARCH_PERCENT_THRESHOLD * num_docs - }) -} // Bulk MAXSCORE path for top-k disjunctions (Lucene MaxScoreBulkScorer // style). Default on: with right-sized partitions it wins by a wide margin // (Lucene-parity latency) and its results are score-identical to the classic @@ -1328,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, @@ -1402,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()); @@ -1427,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 @@ -1496,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, @@ -1517,6 +1543,274 @@ 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) 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: V, +} + +impl<'a> ModernWandDocuments<'a, AllModernDocuments> { + pub(crate) fn all(lengths: &'a DocLengths) -> Self { + Self { + lengths, + visibility: AllModernDocuments, + } + } +} + +impl<'a> ModernWandDocuments<'a, &'a DocVisibility> { + pub(crate) fn filtered(lengths: &'a DocLengths, visibility: &'a DocVisibility) -> Self { + Self { + lengths, + visibility, + } + } +} + +impl WandDocuments for ModernWandDocuments<'_, V> { + 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 @@ -1636,7 +1930,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, @@ -1673,7 +1967,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 @@ -1696,11 +1990,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(); @@ -1740,7 +2034,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, } @@ -1767,8 +2061,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))?; @@ -1806,19 +2099,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![]); } - if should_flat_search(self.operator, &mask, self.docs.len() as u64) { - let row_ids = mask - .iter_addrs() - .expect("should_flat_search guarantees an iterable mask"); - return self.flat_search(params, row_ids, metrics); + 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 @@ -1833,7 +2126,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 @@ -1855,15 +2148,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 { @@ -1880,41 +2167,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() @@ -1938,7 +2200,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(), @@ -1975,50 +2237,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![]); } - debug_assert!( - self.docs.is_empty() || self.docs.supports_reverse_lookup(), - "flat_search needs reverse lookups (inv, or sorted legacy row_ids); \ - the caller picked the wrong DocSet shape for this mask" - ); - // we need to map the row ids to doc ids, and sort them, - // because WAND PostingIterator can't go back to the previous doc id. - // 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() @@ -2032,7 +2267,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); @@ -2062,10 +2297,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 @@ -2079,7 +2313,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(), @@ -2092,9 +2326,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. @@ -2109,9 +2341,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, @@ -2119,7 +2350,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() @@ -2295,9 +2525,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 @@ -2308,14 +2539,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() { @@ -2348,7 +2574,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), ), }; } @@ -2359,9 +2585,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( @@ -2467,7 +2693,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, @@ -2502,17 +2728,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 @@ -2548,7 +2768,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) } @@ -2558,9 +2778,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)| { @@ -2597,14 +2817,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 @@ -2703,10 +2916,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); @@ -2785,10 +2995,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)); @@ -2822,14 +3029,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; @@ -3239,7 +3444,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])); @@ -3247,7 +3452,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)); } } }, @@ -3282,16 +3487,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 @@ -3331,7 +3529,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( @@ -3373,14 +3571,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) { @@ -4274,15 +4465,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::>(); @@ -4663,11 +4851,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!( @@ -4677,13 +4861,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 @@ -4724,11 +4905,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 } @@ -4781,7 +4958,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) @@ -4853,11 +5030,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()); } @@ -4905,20 +5078,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)) }; @@ -4966,13 +5127,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); @@ -5262,7 +5417,6 @@ mod tests { let result = wand .search( &FtsSearchParams::new().with_limit(Some(1)), - Arc::new(RowAddrMask::default()), &NoOpMetricsCollector, ) .unwrap(); @@ -5483,14 +5637,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); @@ -5786,13 +5933,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 @@ -5841,15 +5990,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 @@ -5905,13 +6053,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] @@ -5983,7 +6133,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(), @@ -5994,10 +6144,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]); } @@ -6050,7 +6197,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(), @@ -6059,13 +6209,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:?}" ); } @@ -6089,6 +6240,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_256_document_blocks_without_impacts_use_conservative_quantized_score_bound() { let exact_doc_length = 300; @@ -6391,7 +6559,7 @@ mod tests { params.phrase_slop = Some(slop); } - let normalize = |result: Vec| { + let normalize = |result: Vec>| { let mut rows = result .into_iter() .map(|candidate| { @@ -6399,10 +6567,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::>(); @@ -6418,14 +6583,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) };