diff --git a/docs/src/format/index/scalar/fts.md b/docs/src/format/index/scalar/fts.md index d5c75158011..563d7ebf994 100644 --- a/docs/src/format/index/scalar/fts.md +++ b/docs/src/format/index/scalar/fts.md @@ -34,6 +34,15 @@ An FTS index may contain multiple partitions. Each partition has its own set of | `_rowid` | UInt64 | false | Document row ID | | `_num_tokens` | UInt32 | false | Number of tokens in the document | +Partitioned `docs.lance` files may include the optional schema metadata key +`total_tokens`. Its decimal `UInt64` value is the sum of `_num_tokens` in that +file. `_num_tokens` remains the canonical per-document data. Readers use the +metadata value to construct exact corpus statistics without scanning the column; +when the key is absent, they compute the sum from `_num_tokens`. Writers produce +the key from the same document table in the same file commit. A present value +that cannot be parsed, or that differs when `_num_tokens` is subsequently +loaded, is file corruption. + ### FTS List File Schema | Column | Type | Nullable | Description | diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 72604243c1c..fd7207388b2 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 c90f96053ef..3619eeae59d 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..c415111e131 --- /dev/null +++ b/rust/lance-index/src/scalar/inverted/documents.rs @@ -0,0 +1,1921 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Typed document-side state for partitioned FTS indices. +//! +//! Posting lists use dense, partition-local document identifiers. This module +//! keeps that identity separate from dataset-version row addresses so scoring +//! never has to infer which value a numeric slot represents. + +use std::ops::Range; +use std::sync::{Arc, OnceLock}; + +use arrow::buffer::ScalarBuffer; +use arrow_array::{Array, RecordBatch, UInt32Array, UInt64Array}; +use lance_core::deepsize::DeepSizeOf; +use lance_core::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"; + +/// Maximum number of disjoint point ranges used to resolve final addresses. +const MAX_POINT_ADDRESS_RANGES: usize = 256; +/// Maximum projected address bytes used by the point-read plan. +const MAX_POINT_ADDRESS_BYTES: usize = 64 * 1024; + +/// Dense, immutable document identity inside one FTS partition. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub(super) struct DocId(u32); + +impl DocId { + pub(crate) fn new(value: u32) -> Self { + Self(value) + } + + pub(crate) fn get(self) -> u32 { + self.0 + } + + fn as_usize(self) -> usize { + self.0 as usize + } +} + +/// Immutable corpus statistics for one partition. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct PartitionStats { + pub(crate) num_docs: usize, + pub(crate) total_tokens: u64, +} + +/// Exact document lengths plus the optional quantized scoring representation. +#[derive(Debug)] +pub(super) struct DocLengths { + values: ScalarBuffer, + total_tokens: u64, + quantized_scoring: bool, + norms: OnceLock>, +} + +impl DeepSizeOf for DocLengths { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.values.deep_size_of_children(context) + + self + .norms + .get() + .map(|norms| std::mem::size_of_val(norms.as_ref())) + .unwrap_or(0) + } +} + +impl DocLengths { + fn try_new( + values: ScalarBuffer, + expected_num_docs: usize, + persisted_total_tokens: Option, + quantized_scoring: bool, + path: &str, + ) -> Result { + if values.len() != expected_num_docs { + return Err(corrupt_docs( + path, + format!( + "{NUM_TOKEN_COL} has {} rows but the file footer reports {expected_num_docs}", + values.len() + ), + )); + } + let total_tokens = values.iter().try_fold(0_u64, |total, &value| { + total + .checked_add(u64::from(value)) + .ok_or_else(|| corrupt_docs(path, format!("{NUM_TOKEN_COL} sum overflows u64"))) + })?; + if let Some(expected) = persisted_total_tokens + && expected != total_tokens + { + return Err(corrupt_docs( + path, + format!( + "{TOTAL_TOKENS_KEY} metadata is {expected}, but {NUM_TOKEN_COL} sums to {total_tokens}" + ), + )); + } + Ok(Self { + values, + total_tokens, + quantized_scoring, + norms: OnceLock::new(), + }) + } + + pub(crate) fn len(&self) -> usize { + self.values.len() + } + + pub(crate) fn total_tokens(&self) -> u64 { + self.total_tokens + } + + #[inline] + pub(crate) fn exact(&self, doc_id: DocId) -> u32 { + self.values[doc_id.as_usize()] + } + + pub(crate) fn scoring_norms(&self) -> Option<&[u8]> { + if !self.quantized_scoring { + return None; + } + Some( + self.norms + .get_or_init(|| { + self.values + .iter() + .map(|&length| quantize_doc_length(length)) + .collect() + }) + .as_ref(), + ) + } + + 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(ScalarBuffer), + Owned(Vec), +} + +impl AddressValues { + fn len(&self) -> usize { + match self { + Self::Shared(values) => values.len(), + Self::Owned(values) => values.len(), + } + } + + fn get(&self, index: usize) -> u64 { + match self { + Self::Shared(values) => values[index], + Self::Owned(values) => values[index], + } + } +} + +impl DeepSizeOf for AddressValues { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + match self { + Self::Shared(values) => values.deep_size_of_children(context), + Self::Owned(values) => values.deep_size_of_children(context), + } + } +} + +/// A compact address-sorted view of the live DocIds in a projection. +/// +/// Most newly built partitions already store addresses in DocId order, so the +/// identity variant adds no per-document memory. Remapped or otherwise +/// unsorted projections keep only a u32 permutation instead of duplicating +/// the u64 address column. +#[derive(Debug)] +enum AddressDocIdLookup { + Identity, + Sorted(Box<[u32]>), +} + +impl DeepSizeOf for AddressDocIdLookup { + fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize { + match self { + Self::Identity => 0, + Self::Sorted(doc_ids) => std::mem::size_of_val(doc_ids.as_ref()), + } + } +} + +impl AddressDocIdLookup { + fn build(projection: &VersionAddressProjection) -> Self { + if projection.live_docs.is_none() + && (1..projection.addresses.len()) + .all(|index| projection.addresses.get(index - 1) <= projection.addresses.get(index)) + { + return Self::Identity; + } + + let mut doc_ids = match projection.live_docs.as_ref() { + Some(live_docs) => live_docs.iter().collect::>(), + None => (0..projection.addresses.len() as u32).collect::>(), + }; + doc_ids.sort_unstable_by_key(|&doc_id| projection.addresses.get(doc_id as usize)); + Self::Sorted(doc_ids.into_boxed_slice()) + } + + fn len(&self, projection: &VersionAddressProjection) -> usize { + match self { + Self::Identity => projection.addresses.len(), + Self::Sorted(doc_ids) => doc_ids.len(), + } + } + + fn doc_id_at(&self, position: usize) -> u32 { + match self { + Self::Identity => position as u32, + Self::Sorted(doc_ids) => doc_ids[position], + } + } + + fn address_at(&self, projection: &VersionAddressProjection, position: usize) -> u64 { + projection.addresses.get(self.doc_id_at(position) as usize) + } + + fn partition_point( + &self, + projection: &VersionAddressProjection, + mut predicate: impl FnMut(u64) -> bool, + ) -> usize { + let mut left = 0; + let mut right = self.len(projection); + while left < right { + let middle = left + (right - left) / 2; + if predicate(self.address_at(projection, middle)) { + left = middle + 1; + } else { + right = middle; + } + } + left + } + + fn insert_address_range( + &self, + projection: &VersionAddressProjection, + start: u64, + end: u64, + selected: &mut RoaringBitmap, + ) { + let first = self.partition_point(projection, |address| address < start); + let after_last = self.partition_point(projection, |address| address <= end); + for position in first..after_last { + selected.insert(self.doc_id_at(position)); + } + } + + fn matching_doc_ids( + &self, + projection: &VersionAddressProjection, + addresses: &RowAddrTreeMap, + ) -> RoaringBitmap { + let mut selected = RoaringBitmap::new(); + for (&fragment_id, selection) in addresses.iter() { + match selection { + RowAddrSelection::Full => { + let start = u64::from(RowAddress::new_from_parts(fragment_id, 0)); + let end = u64::from(RowAddress::new_from_parts(fragment_id, u32::MAX)); + self.insert_address_range(projection, start, end, &mut selected); + } + RowAddrSelection::Partial(offsets) => { + let mut offsets = offsets.iter(); + while let Some(range) = offsets.next_range() { + let start = + u64::from(RowAddress::new_from_parts(fragment_id, *range.start())); + let end = u64::from(RowAddress::new_from_parts(fragment_id, *range.end())); + self.insert_address_range(projection, start, end, &mut selected); + } + } + } + } + selected + } + + fn visibility( + &self, + projection: &VersionAddressProjection, + mask: &RowAddrMask, + ) -> DocVisibility { + match mask { + RowAddrMask::AllowList(allowed) => { + DocVisibility::Selected(self.matching_doc_ids(projection, allowed)) + } + RowAddrMask::BlockList(blocked) => { + let blocked = self.matching_doc_ids(projection, blocked); + let mut selected = projection.live_doc_ids(); + selected -= &blocked; + DocVisibility::Selected(selected) + } + } + } +} + +/// Addresses projected into the dataset version that opened the index. +#[derive(Debug)] +pub(super) struct VersionAddressProjection { + 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>, +} + +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(raw.values().clone()), + 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(addresses), + live_docs: Some(live_docs), + doc_ids_by_address: OnceCell::new(), + }) + } + + fn address(&self, doc_id: DocId) -> Option { + if self + .live_docs + .as_ref() + .is_some_and(|live| !live.contains(doc_id.get())) + { + None + } else { + Some(self.addresses.get(doc_id.as_usize())) + } + } + + fn live_doc_ids(&self) -> RoaringBitmap { + self.live_docs + .clone() + .unwrap_or_else(|| (0..self.addresses.len() as u32).collect()) + } + + async fn doc_ids_by_address(self: &Arc) -> Result> { + self.doc_ids_by_address + .get_or_try_init(|| { + let projection = self.clone(); + async move { + spawn_cpu(move || Result::Ok(Arc::new(AddressDocIdLookup::build(&projection)))) + .await + } + }) + .await + .cloned() + } + + async fn materialize_visibility( + self: &Arc, + mask: Arc, + ) -> Result { + let lookup = self.doc_ids_by_address().await?; + let projection = self.clone(); + spawn_cpu(move || Result::Ok(lookup.visibility(&projection, &mask))).await + } +} + +/// Query-local selection in the partition-local DocId domain. +#[derive(Debug, Clone)] +pub(super) enum DocVisibility { + All, + Selected(RoaringBitmap), + Filtered { + projection: Arc, + 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, + num_docs: usize, + persisted_total_tokens: Option, + quantized_scoring: bool, + remapper: Option>, + lengths: OnceCell>, + projection: OnceCell>, + 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, + reader: &dyn IndexReader, + remapper: Option>, + quantized_scoring: bool, + ) -> Result { + let num_docs = reader.num_rows(); + if num_docs > u32::MAX as usize { + return Err(corrupt_docs( + &path, + format!("document count {num_docs} exceeds dense DocId capacity"), + )); + } + let persisted_total_tokens = reader + .schema() + .metadata + .get(TOTAL_TOKENS_KEY) + .map(|value| { + value.parse::().map_err(|error| { + corrupt_docs( + &path, + format!("invalid {TOTAL_TOKENS_KEY} metadata value {value:?}: {error}"), + ) + }) + }) + .transpose()?; + Ok(Self { + store, + path, + num_docs, + persisted_total_tokens, + quantized_scoring, + remapper, + lengths: OnceCell::new(), + projection: OnceCell::new(), + 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() + } + + pub(crate) fn projection_loaded(&self) -> bool { + self.projection.initialized() + } + + pub(crate) fn query_ready(&self) -> bool { + self.prewarm_complete.initialized() + && self + .lengths + .get() + .is_some_and(|lengths| lengths.scoring_ready()) + && self + .projection + .get() + .is_some_and(|projection| projection.doc_ids_by_address.initialized()) + } + + async fn reader(&self) -> Result> { + self.store.open_index_file(&self.path).await + } + + fn lengths_from_batch(&self, batch: &RecordBatch) -> Result> { + let column = required_u32_column(batch, NUM_TOKEN_COL, &self.path)?; + if column.null_count() != 0 { + return Err(corrupt_docs( + &self.path, + format!("{NUM_TOKEN_COL} contains null values"), + )); + } + Ok(Arc::new(DocLengths::try_new( + column.values().clone(), + self.num_docs, + self.persisted_total_tokens, + self.quantized_scoring, + &self.path, + )?)) + } + + fn projection_from_batch(&self, batch: &RecordBatch) -> Result> { + let column = required_u64_column(batch, ROW_ID, &self.path)?; + Ok(Arc::new(VersionAddressProjection::try_new( + column, + self.num_docs, + self.remapper.as_deref(), + &self.path, + )?)) + } + + pub(crate) async fn stats(&self) -> Result { + let total_tokens = match self.persisted_total_tokens { + Some(total_tokens) => total_tokens, + 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() + } + + pub(crate) async fn projection(&self) -> Result> { + self.projection + .get_or_try_init(|| async { + let reader = self.reader().await?; + let batch = reader.read_range(0..self.num_docs, Some(&[ROW_ID])).await?; + self.projection_from_batch(&batch) + }) + .await + .cloned() + } + + 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.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.projection.get()?.clone(); + if mask.is_select_all() { + return Some(DocVisibility::Selected(projection.live_doc_ids())); + } + if materialize_selected { + None + } else { + Some(DocVisibility::Filtered { projection, mask }) + } + } + + /// Resolve final global top-k DocIds to current row addresses. + pub(crate) async fn resolve_addresses(&self, doc_ids: &[DocId]) -> Result> { + if doc_ids.is_empty() { + return Ok(Vec::new()); + } + for doc_id in doc_ids { + if doc_id.as_usize() >= self.num_docs { + return Err(corrupt_docs( + &self.path, + format!( + "candidate DocId {} is outside [0, {})", + doc_id.get(), + self.num_docs + ), + )); + } + } + if let Some(projection) = self.projection.get() { + return doc_ids + .iter() + .map(|&doc_id| { + projection.address(doc_id).ok_or_else(|| { + corrupt_docs( + &self.path, + format!("candidate DocId {} is not live", doc_id.get()), + ) + }) + }) + .collect(); + } + if self.remapper.is_some() { + let projection = self.projection().await?; + return doc_ids + .iter() + .map(|&doc_id| { + projection.address(doc_id).ok_or_else(|| { + corrupt_docs( + &self.path, + format!("candidate DocId {} is not live", doc_id.get()), + ) + }) + }) + .collect(); + } + + let mut unique = doc_ids + .iter() + .map(|doc_id| doc_id.get()) + .collect::>(); + unique.sort_unstable(); + unique.dedup(); + let ranges = coalesce_doc_ids(&unique); + let point_bytes = unique.len().saturating_mul(std::mem::size_of::()); + let use_bulk = ranges.len() > MAX_POINT_ADDRESS_RANGES + || point_bytes > MAX_POINT_ADDRESS_BYTES + || unique.len() == self.num_docs; + let reader = self.reader().await?; + let address_column = if use_bulk { + let batch = reader.read_range(0..self.num_docs, Some(&[ROW_ID])).await?; + required_u64_column(&batch, ROW_ID, &self.path)?.clone() + } else { + let batch = reader.read_ranges(&ranges, Some(&[ROW_ID])).await?; + required_u64_column(&batch, ROW_ID, &self.path)?.clone() + }; + if address_column.null_count() != 0 { + return Err(corrupt_docs( + &self.path, + format!("{ROW_ID} contains null values"), + )); + } + let addresses = address_column.values(); + if use_bulk && addresses.len() != self.num_docs { + return Err(corrupt_docs( + &self.path, + format!( + "{ROW_ID} bulk read returned {} rows, expected {}", + addresses.len(), + self.num_docs + ), + )); + } + if !use_bulk && addresses.len() != unique.len() { + return Err(corrupt_docs( + &self.path, + format!( + "{ROW_ID} point read returned {} rows, expected {}", + addresses.len(), + unique.len() + ), + )); + } + + if use_bulk { + Ok(doc_ids + .iter() + .map(|doc_id| addresses[doc_id.as_usize()]) + .collect()) + } else { + let by_doc_id = unique + .into_iter() + .zip(addresses.iter().copied()) + .collect::>(); + doc_ids + .iter() + .map(|doc_id| { + by_doc_id.get(&doc_id.get()).copied().ok_or_else(|| { + Error::internal(format!( + "resolved address missing for DocId {}", + doc_id.get() + )) + }) + }) + .collect() + } + } + + /// Estimated Arrow payload retained while resolving these final candidates. + /// The estimate is used to cap cross-partition read concurrency; a single + /// oversized partition is still allowed to make progress. + pub(crate) fn estimated_address_read_bytes(&self, doc_ids: &[DocId]) -> usize { + if doc_ids.is_empty() || self.projection.initialized() { + return 0; + } + if self.remapper.is_some() { + return self.num_docs.saturating_mul(std::mem::size_of::()); + } + + // Avoid repeating the sort/dedup performed by `resolve_addresses`. + // Treating every requested DocId as distinct and disjoint can only + // overestimate the payload, which makes the concurrency bound safer. + let point_bytes = doc_ids.len().saturating_mul(std::mem::size_of::()); + let use_bulk = doc_ids.len() > MAX_POINT_ADDRESS_RANGES + || point_bytes > MAX_POINT_ADDRESS_BYTES + || doc_ids.len() >= self.num_docs; + if use_bulk { + self.num_docs.saturating_mul(std::mem::size_of::()) + } else { + point_bytes + } + } + + /// Materialize the build-side table for rewrite/update operations. + pub(crate) async fn load_build_docset(&self) -> Result { + DocSet::load(self.reader().await?, false, self.remapper.clone()).await + } + + /// Resolve one DocId synchronously when its current-version projection is resident. + pub(crate) fn cached_row_address(&self, doc_id: DocId) -> Result> { + let Some(projection) = self.projection.get() else { + return Ok(None); + }; + if doc_id.as_usize() >= self.num_docs { + return Err(corrupt_docs( + &self.path, + format!( + "candidate DocId {} is outside [0, {})", + doc_id.get(), + self.num_docs + ), + )); + } + projection + .address(doc_id) + .map(RowAddress::new_from_u64) + .map(Some) + .ok_or_else(|| { + corrupt_docs( + &self.path, + format!("candidate DocId {} is not live", doc_id.get()), + ) + }) + } + + pub(crate) async fn prewarm(&self) -> Result<()> { + self.prewarm_complete + .get_or_try_init(|| async { + if self.lengths.get().is_none() && self.projection.get().is_none() { + let reader = self.reader().await?; + let batch = reader + .read_range(0..self.num_docs, Some(&[ROW_ID, NUM_TOKEN_COL])) + .await?; + let lengths = self.lengths_from_batch(&batch)?; + let projection = self.projection_from_batch(&batch)?; + + // A concurrent single-column request may win either OnceCell. + // Awaiting the accessors below joins that initialization without + // replacing the already-published value. + let _ = self.lengths.set(lengths); + let _ = self.projection.set(projection); + } + + let lengths = self.lengths().await?; + spawn_cpu(move || { + let _ = lengths.scoring_norms(); + Result::Ok(()) + }) + .await?; + self.projection().await?.doc_ids_by_address().await?; + Result::Ok(()) + }) + .await?; + Ok(()) + } +} + +fn required_u32_column<'a>( + batch: &'a RecordBatch, + name: &str, + path: &str, +) -> Result<&'a UInt32Array> { + let column = batch + .column_by_name(name) + .ok_or_else(|| corrupt_docs(path, format!("required column {name} is missing")))?; + column + .as_any() + .downcast_ref::() + .ok_or_else(|| { + corrupt_docs( + path, + format!( + "column {name} has type {}, expected UInt32", + column.data_type() + ), + ) + }) +} + +fn required_u64_column<'a>( + batch: &'a RecordBatch, + name: &str, + path: &str, +) -> Result<&'a UInt64Array> { + let column = batch + .column_by_name(name) + .ok_or_else(|| corrupt_docs(path, format!("required column {name} is missing")))?; + column + .as_any() + .downcast_ref::() + .ok_or_else(|| { + corrupt_docs( + path, + format!( + "column {name} has type {}, expected UInt64", + column.data_type() + ), + ) + }) +} + +fn coalesce_doc_ids(doc_ids: &[u32]) -> Vec> { + let mut ranges = Vec::new(); + let Some((&first, rest)) = doc_ids.split_first() else { + return ranges; + }; + let mut start = first; + let mut end = first + 1; + for &doc_id in rest { + if doc_id == end { + end += 1; + } else { + ranges.push(start as usize..end as usize); + start = doc_id; + end = doc_id + 1; + } + } + ranges.push(start as usize..end as usize); + ranges +} + +fn corrupt_docs(path: &str, message: impl Into) -> Error { + Error::corrupt_file(Path::from(path), message) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + use arrow_array::{ArrayRef, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use async_trait::async_trait; + use lance_core::cache::LanceCache; + use lance_core::utils::tempfile::TempObjDir; + use lance_io::object_store::ObjectStore; + use lance_select::RowAddrTreeMap; + use roaring::RoaringTreemap; + use 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) { + let directory = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + ObjectStore::local().into(), + directory.clone(), + Arc::new(LanceCache::no_cache()), + )); + (directory, store) + } + + async fn write_documents( + store: &dyn IndexStore, + path: &str, + addresses: UInt64Array, + lengths: UInt32Array, + total_tokens: Option<&str>, + ) { + let schema = Arc::new(Schema::new(vec![ + Field::new(ROW_ID, DataType::UInt64, addresses.null_count() != 0), + Field::new(NUM_TOKEN_COL, DataType::UInt32, lengths.null_count() != 0), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(addresses) as ArrayRef, + Arc::new(lengths) as ArrayRef, + ], + ) + .unwrap(); + let mut writer = store.new_index_file(path, schema).await.unwrap(); + writer.write_record_batch(batch).await.unwrap(); + if let Some(total_tokens) = total_tokens { + writer + .finish_with_metadata(HashMap::from([( + TOTAL_TOKENS_KEY.to_owned(), + total_tokens.to_owned(), + )])) + .await + .unwrap(); + } else { + writer.finish().await.unwrap(); + } + } + + async fn open_documents( + store: Arc, + path: &str, + remapper: Option>, + ) -> Result { + let reader = store.open_index_file(path).await?; + PartitionDocuments::try_new(store, path.to_owned(), reader.as_ref(), remapper, false) + } + + fn counted_store( + inner: Arc, + target: &str, + ) -> (Arc, Arc) { + let counts = Arc::new(DocumentReadCounts::default()); + ( + Arc::new(CountingStore { + inner, + target: target.to_owned(), + counts: counts.clone(), + 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 coalesces_sorted_doc_ids() { + assert_eq!( + coalesce_doc_ids(&[1, 2, 3, 7, 9, 10]), + vec![1..4, 7..8, 9..11] + ); + } + + #[test] + fn required_document_columns_validate_name_and_type() { + let schema = Arc::new(Schema::new(vec![Field::new( + NUM_TOKEN_COL, + DataType::UInt64, + false, + )])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(UInt64Array::from(vec![1])) as ArrayRef], + ) + .unwrap(); + + let missing = required_u64_column(&batch, ROW_ID, "docs.lance").unwrap_err(); + assert!( + missing + .to_string() + .contains("required column _rowid is missing") + ); + let wrong_type = required_u32_column(&batch, NUM_TOKEN_COL, "docs.lance").unwrap_err(); + assert!(wrong_type.to_string().contains("expected UInt32")); + } + + #[test] + fn live_documents_use_doc_ids() { + let projection = VersionAddressProjection { + addresses: AddressValues::Owned(vec![10, 20, 30]), + live_docs: Some(RoaringBitmap::from_iter([0, 2])), + doc_ids_by_address: OnceCell::new(), + }; + let 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"), + ); + + 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 + .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 + .materialize_visibility(blocked) + .await + .expect("valid block-list") + else { + panic!("block-list must compile to DocIds") + }; + assert_eq!(selected.iter().collect::>(), vec![0, 3]); + let second_lookup = projection + .doc_ids_by_address() + .await + .expect("cached address lookup"); + assert!(Arc::ptr_eq(&first_lookup, &second_lookup)); + } + + #[tokio::test] + async fn materialized_visibility_handles_unsorted_duplicate_addresses_and_full_fragments() { + let row_address = |fragment_id, row_offset| { + u64::from(RowAddress::new_from_parts(fragment_id, row_offset)) + }; + let projection = Arc::new(VersionAddressProjection { + addresses: AddressValues::Owned(vec![ + row_address(2, 5), + row_address(1, 2), + row_address(1, 2), + row_address(1, 4), + ]), + live_docs: None, + doc_ids_by_address: OnceCell::new(), + }); + + let allowed = Arc::new(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([ + row_address(1, 2), + row_address(1, 3), + row_address(1, 4), + ]))); + let DocVisibility::Selected(selected) = projection + .materialize_visibility(allowed) + .await + .expect("valid allow-list") + else { + panic!("allow-list must compile to DocIds") + }; + assert_eq!(selected.iter().collect::>(), vec![1, 2, 3]); + + let mut full_fragment = RowAddrTreeMap::new(); + full_fragment.insert_fragment(2); + let DocVisibility::Selected(selected) = projection + .materialize_visibility(Arc::new(RowAddrMask::from_allowed(full_fragment))) + .await + .expect("valid full-fragment allow-list") + else { + panic!("allow-list must compile to DocIds") + }; + assert_eq!(selected.iter().collect::>(), vec![0]); + } + + #[test] + fn lazy_visibility_projects_only_candidate_doc_ids() { + let projection = Arc::new(VersionAddressProjection { + addresses: AddressValues::Owned(vec![10, 20, 30]), + live_docs: Some(RoaringBitmap::from_iter([0, 2])), + doc_ids_by_address: OnceCell::new(), + }); + let visibility = DocVisibility::Filtered { + projection: projection.clone(), + mask: Arc::new(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([ + 20, 30, + ]))), + }; + + assert!(!visibility.selected(DocId::new(0))); + assert!(!visibility.selected(DocId::new(1))); + assert!(visibility.selected(DocId::new(2))); + assert!(!projection.doc_ids_by_address.initialized()); + } + + #[test] + fn doc_lengths_validate_shape_total_and_memory() { + let mismatch = DocLengths::try_new(ScalarBuffer::from(vec![2, 3]), 3, None, false, "docs") + .unwrap_err(); + assert!(mismatch.to_string().contains("2 rows")); + + let mismatch = DocLengths::try_new( + ScalarBuffer::from(vec![2, 3, 5]), + 3, + Some(11), + false, + "docs", + ) + .unwrap_err(); + assert!(mismatch.to_string().contains("sums to 10")); + + let lengths = + DocLengths::try_new(ScalarBuffer::from(vec![2, 3, 5]), 3, Some(10), true, "docs") + .unwrap(); + let before_norms = lengths.deep_size_of(); + assert_eq!(lengths.total_tokens(), 10); + assert_eq!(lengths.scoring_norms().unwrap().len(), 3); + assert_eq!(lengths.deep_size_of() - before_norms, 3); + } + + #[tokio::test] + async fn persisted_stats_are_footer_only_and_compatible_with_full_docset_reader() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + Some("10"), + ) + .await; + + let reader = store.open_index_file(path).await.unwrap(); + assert_eq!( + reader.schema().metadata.get(TOTAL_TOKENS_KEY), + Some(&"10".to_owned()) + ); + let complete = DocSet::load(reader, false, None).await.unwrap(); + assert_eq!(complete.len(), 3); + assert_eq!(complete.row_id(1), 20); + assert_eq!(complete.total_tokens_num(), 10); + + let (counting, counts) = counted_store(store, path); + let documents = open_documents(counting, path, None).await.unwrap(); + assert_eq!( + documents.stats().await.unwrap(), + PartitionStats { + num_docs: 3, + total_tokens: 10, + } + ); + assert_eq!(counts.rows.load(Ordering::Relaxed), 0); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); + } + + #[tokio::test] + async fn missing_stats_fall_back_once_to_lengths() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + None, + ) + .await; + let (counting, counts) = counted_store(store, path); + let documents = open_documents(counting, path, None).await.unwrap(); + + assert_eq!(documents.stats().await.unwrap().total_tokens, 10); + assert_eq!(documents.stats().await.unwrap().total_tokens, 10); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 1); + assert_eq!(counts.length_rows.load(Ordering::Relaxed), 3); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 0); + assert!(documents.lengths_loaded()); + assert!(!documents.projection_loaded()); + } + + #[tokio::test] + async fn prewarm_loads_document_columns_and_address_lookup_once() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + Some("10"), + ) + .await; + let (counting, counts) = counted_store(store, path); + let documents = open_documents(counting, path, None).await.unwrap(); + + futures::future::join_all((0..8).map(|_| documents.prewarm())) + .await + .into_iter() + .collect::>>() + .unwrap(); + let projection = documents.projection().await.unwrap(); + let first_lookup = projection.doc_ids_by_address().await.unwrap(); + documents.prewarm().await.unwrap(); + let second_lookup = projection.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_address(DocId::new(1)).unwrap(), + Some(RowAddress::new_from_u64(20)) + ); + assert_eq!(counts.open_calls.load(Ordering::Relaxed), 2); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 1); + assert_eq!(counts.length_rows.load(Ordering::Relaxed), 3); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 3); + } + + #[tokio::test] + async fn prewarm_materializes_quantized_norms_before_becoming_query_ready() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 300, 5]), + Some("307"), + ) + .await; + let reader = store.open_index_file(path).await.unwrap(); + let documents = + PartitionDocuments::try_new(store, path.to_owned(), reader.as_ref(), None, true) + .unwrap(); + + assert!(!documents.query_ready()); + documents.prewarm().await.unwrap(); + let lengths = documents.lengths().await.unwrap(); + assert!(lengths.scoring_ready()); + assert_eq!(lengths.scoring_norms().unwrap().len(), 3); + assert!(documents.query_ready()); + } + + #[tokio::test] + async fn cancelled_or_failed_prewarm_can_retry_without_partial_publication() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + Some("10"), + ) + .await; + + let (pausing, pause) = faulting_store(store.clone(), path, PAUSE_ONCE); + let documents = Arc::new(open_documents(pausing, path, None).await.unwrap()); + let task = tokio::spawn({ + let documents = documents.clone(); + async move { documents.prewarm().await } + }); + tokio::time::timeout(Duration::from_secs(5), pause.started.notified()) + .await + .expect("prewarm should reach the injected pending read"); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); + assert!(!documents.query_ready()); + documents.prewarm().await.unwrap(); + assert!(documents.query_ready()); + + let (failing, _fault) = faulting_store(store, path, FAIL_ONCE); + let documents = open_documents(failing, path, None).await.unwrap(); + let error = documents.prewarm().await.unwrap_err(); + assert!(error.to_string().contains("injected document read failure")); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); + assert!(!documents.query_ready()); + documents.prewarm().await.unwrap(); + assert!(documents.query_ready()); + } + + #[tokio::test] + async fn prewarm_reuses_an_already_loaded_document_column() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + write_documents( + store.as_ref(), + path, + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + Some("10"), + ) + .await; + let (counting, counts) = counted_store(store, path); + let documents = open_documents(counting, path, None).await.unwrap(); + + documents.lengths().await.unwrap(); + documents.prewarm().await.unwrap(); + documents.prewarm().await.unwrap(); + + assert_eq!(counts.open_calls.load(Ordering::Relaxed), 3); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 2); + assert_eq!(counts.length_rows.load(Ordering::Relaxed), 3); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 3); + } + + #[tokio::test] + async fn invalid_or_mismatched_stats_are_corruption_not_fallback() { + let (_directory, store) = test_store(); + write_documents( + store.as_ref(), + "invalid.lance", + UInt64Array::from(vec![10]), + UInt32Array::from(vec![2]), + Some("not-a-u64"), + ) + .await; + let error = open_documents(store.clone(), "invalid.lance", None) + .await + .unwrap_err(); + assert!(error.to_string().contains("invalid total_tokens")); + + write_documents( + store.as_ref(), + "mismatch.lance", + UInt64Array::from(vec![10, 20, 30]), + UInt32Array::from(vec![2, 3, 5]), + Some("11"), + ) + .await; + let (counting, counts) = counted_store(store, "mismatch.lance"); + let documents = open_documents(counting, "mismatch.lance", None) + .await + .unwrap(); + assert_eq!(documents.stats().await.unwrap().total_tokens, 11); + let error = documents.lengths().await.unwrap_err(); + assert!(error.to_string().contains("sums to 10")); + assert_eq!(counts.length_rows.load(Ordering::Relaxed), 3); + assert!(!documents.lengths_loaded()); + } + + #[tokio::test] + async fn final_address_resolution_obeys_point_and_bulk_budgets() { + let (_directory, store) = test_store(); + let path = "docs.lance"; + let num_docs = 600_u64; + write_documents( + store.as_ref(), + path, + UInt64Array::from_iter_values((0..num_docs).map(|id| id + 1_000)), + UInt32Array::from_iter_values((0..num_docs).map(|_| 1)), + Some("600"), + ) + .await; + let (counting, counts) = counted_store(store, path); + let documents = open_documents(counting, path, None).await.unwrap(); + + assert!(documents.resolve_addresses(&[]).await.unwrap().is_empty()); + assert_eq!(counts.rows.load(Ordering::Relaxed), 0); + + let point_ids = [DocId::new(5), DocId::new(6), DocId::new(10), DocId::new(5)]; + assert_eq!(documents.estimated_address_read_bytes(&point_ids), 32); + assert_eq!( + documents.resolve_addresses(&point_ids).await.unwrap(), + vec![1005, 1006, 1010, 1005] + ); + assert_eq!(counts.ranges_calls.load(Ordering::Relaxed), 1); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 0); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 3); + + let bulk_ids = (0..=512).step_by(2).map(DocId::new).collect::>(); + assert_eq!(bulk_ids.len(), MAX_POINT_ADDRESS_RANGES + 1); + 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), 1); + assert_eq!(counts.range_calls.load(Ordering::Relaxed), 1); + assert_eq!(counts.address_rows.load(Ordering::Relaxed), 603); + assert!(!documents.projection_loaded()); + } + + #[tokio::test] + async fn document_column_nulls_are_reported_as_corruption() { + let (_directory, store) = test_store(); + write_documents( + store.as_ref(), + "null-length.lance", + UInt64Array::from(vec![Some(10), Some(20)]), + UInt32Array::from(vec![Some(2), None]), + Some("2"), + ) + .await; + let documents = open_documents(store.clone(), "null-length.lance", None) + .await + .unwrap(); + assert!( + documents + .lengths() + .await + .unwrap_err() + .to_string() + .contains("_num_tokens contains null") + ); + + write_documents( + store.as_ref(), + "null-address.lance", + UInt64Array::from(vec![Some(10), None]), + UInt32Array::from(vec![2, 3]), + Some("5"), + ) + .await; + let documents = open_documents(store, "null-address.lance", None) + .await + .unwrap(); + assert!( + documents + .resolve_addresses(&[DocId::new(1)]) + .await + .unwrap_err() + .to_string() + .contains("_rowid contains null") + ); + 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 9b1496f7140..bde7c7aa8fe 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,13 +51,18 @@ 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::{ @@ -267,13 +273,13 @@ pub fn validate_format_version_block_size( } #[derive(Debug)] -struct PartitionCandidates { +struct PartitionCandidates { tokens_by_position: Vec, grouped_expansions: Vec, - candidates: Vec, + candidates: Vec>, } -impl PartitionCandidates { +impl PartitionCandidates { fn empty() -> Self { Self { tokens_by_position: Vec::new(), @@ -283,6 +289,180 @@ impl PartitionCandidates { } } +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)] struct LoadedPostings { postings: Vec, @@ -291,6 +471,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 { @@ -472,6 +673,18 @@ pub(super) fn parse_format_version_from_metadata( } } +#[derive(Debug, Default)] +struct InvertedPrewarmState { + query_ready: bool, + positions_ready: bool, +} + +impl InvertedPrewarmState { + fn satisfies(&self, with_position: bool) -> bool { + self.query_ready && (!with_position || self.positions_ready) + } +} + #[derive(Clone)] pub struct InvertedIndex { params: InvertedIndexParams, @@ -481,6 +694,9 @@ pub struct InvertedIndex { format_version: InvertedListFormatVersion, pub(crate) partitions: Vec>, corpus_stats: Arc>, + prewarm_state: Arc>, + /// One-way publication: document projections never leave memory after loading. + document_projections_resident: Arc, // Fragments which are contained in the index, but no longer in the dataset. // These should be pruned at search time since we don't prune them at update time. deleted_fragments: RoaringBitmap, @@ -504,35 +720,6 @@ impl DeepSizeOf for InvertedIndex { } } -/// Resolve any `Pending` candidates that wand emitted via the -/// deferred-row_id path. After this returns, every entry in -/// `candidates` carries a real row_id. -async fn resolve_deferred_candidates( - docs: &LazyDocSet, - candidates: &mut [DocCandidate], -) -> Result<()> { - let pending: Vec = candidates - .iter() - .filter_map(|c| match c.addr { - CandidateAddr::Pending(d) => Some(d), - CandidateAddr::RowId(_) => None, - }) - .collect(); - if pending.is_empty() { - return Ok(()); - } - let mut iter = docs.resolve_row_ids(&pending).await?.into_iter(); - for c in candidates { - if matches!(c.addr, CandidateAddr::Pending(_)) { - let r = iter.next().ok_or_else(|| { - Error::internal("resolve_row_ids returned fewer items than requested") - })?; - c.addr = CandidateAddr::RowId(r); - } - } - Ok(()) -} - impl InvertedIndex { fn format_version(&self) -> InvertedListFormatVersion { self.format_version @@ -721,29 +908,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() @@ -879,42 +1073,48 @@ impl InvertedIndex { if limit == 0 { return Ok((Vec::new(), Vec::new())); } - - fn push_scored_candidate( - candidates: &mut BinaryHeap>, - limit: usize, - addr: CandidateAddr, - score: f32, - ) -> Result<()> { - // resolve_deferred_candidates ran upstream, so every candidate - // carries a real row_id at this point. - let row_id = match addr { - CandidateAddr::RowId(r) => r, - CandidateAddr::Pending(_) => { - return Err(Error::internal( - "bm25_search post-condition: deferred candidate left unresolved", - )); - } - }; - - if candidates.len() < limit { - candidates.push(Reverse(ScoredDoc::new(row_id, score))); - } else if candidates.peek().unwrap().0.score.0 < score { - candidates.pop(); - candidates.push(Reverse(ScoredDoc::new(row_id, score))); - } - Ok(()) - } - let mask = prefilter.mask(); + if self.is_legacy() { + self.bm25_search_legacy( + tokens, + params, + operator, + mask, + metrics, + scorer, + impact_scorer, + limit, + ) + .await + } else { + self.bm25_search_modern(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`). + #[allow(clippy::too_many_arguments)] + async fn bm25_search_legacy( + &self, + tokens: Arc, + params: Arc, + operator: Operator, + mask: Arc, + metrics: Arc, + scorer: &MemBM25Scorer, + impact_scorer: Arc, + limit: usize, + ) -> Result<(Vec, Vec)> { let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); - let legacy_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + let local_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); let parts = self .partitions .iter() @@ -926,9 +1126,14 @@ impl InvertedIndex { let metrics = metrics.clone(); let impact_scorer = impact_scorer.clone(); let impact_shared_threshold = impact_shared_threshold.clone(); - let legacy_shared_threshold = legacy_shared_threshold.clone(); + let local_shared_threshold = local_shared_threshold.clone(); async move { - let loaded_postings = part + let LoadedPostings { + postings, + grouped_expansions, + impact_safe, + exact_scoring_required, + } = part .load_posting_lists( tokens.as_ref(), params.as_ref(), @@ -937,19 +1142,9 @@ impl InvertedIndex { metrics.as_ref(), ) .await?; - let LoadedPostings { - postings, - grouped_expansions, - impact_safe, - exact_scoring_required, - } = loaded_postings; if postings.is_empty() { - // No hits in this partition; its DocSet stays - // unloaded, so we never pay the per-doc - // row_id/num_tokens download for it. return Result::Ok(PartitionCandidates::empty()); } - let docs_for_wand = part.docs.docs_for_wand(mask.as_ref()).await?; let max_position = postings .iter() .map(|posting| posting.term_index() as usize) @@ -957,124 +1152,331 @@ impl InvertedIndex { .unwrap_or_default(); let mut tokens_by_position = vec![String::new(); max_position + 1]; for posting in &postings { - let idx = posting.term_index() as usize; - tokens_by_position[idx] = posting.token().to_owned(); + tokens_by_position[posting.term_index() as usize] = + posting.token().to_owned(); } - let params = params.clone(); - let mask = mask.clone(); - let metrics = metrics.clone(); - let part_for_wand = part.clone(); + let docs = part.docs.legacy().cloned().ok_or_else(|| { + Error::internal("legacy index contains modern partition documents") + })?; let use_global_scorer = impact_safe || exact_scoring_required; - let partition_threshold = if use_global_scorer { + let threshold = if use_global_scorer { impact_shared_threshold } else { - legacy_shared_threshold + local_shared_threshold }; let wand_scorer = use_global_scorer.then(|| impact_scorer.clone()); + let part_for_wand = part.clone(); let candidates = spawn_cpu(move || { - let candidates = part_for_wand.bm25_search( - docs_for_wand.as_ref(), + part_for_wand.bm25_search_legacy( + docs.as_ref(), params.as_ref(), operator, - mask, + mask.as_ref(), postings, wand_scorer, metrics.as_ref(), - partition_threshold, - )?; - std::result::Result::<_, Error>::Ok(candidates) + threshold, + ) }) .await?; - let mut partition_result = PartitionCandidates { + Result::Ok(PartitionCandidates { tokens_by_position, grouped_expansions, candidates, - }; - resolve_deferred_candidates(&part.docs, &mut partition_result.candidates) - .await?; - Result::Ok(partition_result) + }) } }) .collect::>(); + + let mut ranked = BinaryHeap::new(); + let mut idf_cache = HashMap::new(); let mut parts = stream::iter(parts).buffer_unordered(get_num_compute_intensive_cpus()); - let mut idf_cache: HashMap = HashMap::new(); - while let Some(res) = parts.try_next().await? { - if res.candidates.is_empty() { - continue; - } - let PartitionCandidates { - tokens_by_position, - grouped_expansions, - candidates: part_candidates, - } = res; - let mut idf_by_position = Vec::with_capacity(tokens_by_position.len()); - for token in &tokens_by_position { - let idf_weight = match idf_cache.get(token) { - Some(weight) => *weight, - None => { - let weight = scorer.query_weight(token); - idf_cache.insert(token.clone(), weight); - weight - } - }; - idf_by_position.push(idf_weight); + while let Some(partition) = parts.try_next().await? { + for (row_id, score) in rescore_partition_candidates(partition, scorer, &mut idf_cache) { + push_scored_key(&mut ranked, limit, row_id, score); } + } + Ok(ranked + .into_sorted_vec() + .into_iter() + .map(|Reverse(doc)| (doc.row_id, doc.score.0)) + .unzip()) + } - if grouped_expansions.is_empty() { - for DocCandidate { - addr, - freqs, - doc_length, - .. - } in part_candidates - { - let mut score = 0.0; - for (term_index, freq) in freqs.into_iter() { - debug_assert!((term_index as usize) < idf_by_position.len()); - score += idf_by_position[term_index as usize] - * scorer.doc_weight(freq, doc_length); + 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 + } + } + + fn has_resident_document_projections(&self) -> bool { + if self.document_projections_resident.load(Ordering::Acquire) { + return true; + } + let resident = self.partitions.iter().all(|partition| { + partition + .docs + .modern() + .is_some_and(|documents| documents.projection_loaded()) + }); + if resident { + self.document_projections_resident + .store(true, Ordering::Release); + } + resident + } + + async fn bm25_search_modern_resident( + &self, + request: ModernSearchRequest<'_>, + ) -> Result<(Vec, Vec)> { + let ranked = self.bm25_search_modern_candidates(request).await?; + self.resolve_resident_modern_candidates(ranked) + } + + async fn bm25_search_modern_deferred( + &self, + request: ModernSearchRequest<'_>, + ) -> Result<(Vec, Vec)> { + // 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 + } + + 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 local_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits())); + let parts = self + .partitions + .iter() + .enumerate() + .map(|(partition_ordinal, part)| { + let part = part.clone(); + let tokens = tokens.clone(); + let params = params.clone(); + let mask = mask.clone(); + let metrics = metrics.clone(); + let impact_scorer = impact_scorer.clone(); + let impact_shared_threshold = impact_shared_threshold.clone(); + let local_shared_threshold = local_shared_threshold.clone(); + async move { + let LoadedPostings { + postings, + grouped_expansions, + impact_safe, + exact_scoring_required, + } = part + .load_posting_lists( + tokens.as_ref(), + params.as_ref(), + operator, + impact_scorer.as_ref(), + metrics.as_ref(), + ) + .await?; + if postings.is_empty() { + return Result::Ok((partition_ordinal, PartitionCandidates::empty())); } - push_scored_candidate(&mut candidates, limit, addr, score)?; - } - } else { - let grouped_positions = grouped_expansions - .iter() - .map(|group| group.position) - .collect::>(); - for DocCandidate { - addr, - posting_doc_id, - freqs, - doc_length, - } in part_candidates - { - let mut score = 0.0; - for (term_index, freq) in freqs.into_iter() { - if grouped_positions.contains(&term_index) { - continue; - } - debug_assert!((term_index as usize) < idf_by_position.len()); - score += idf_by_position[term_index as usize] - * scorer.doc_weight(freq, doc_length); + let documents = part.docs.modern().cloned().ok_or_else(|| { + Error::internal("modern index contains legacy partition documents") + })?; + let 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((partition_ordinal, PartitionCandidates::empty())); } - for group in &grouped_expansions { - for term in group.terms.iter() { - let Some(freq) = term.frequency(posting_doc_id) else { - continue; - }; - score += term.query_weight() * scorer.doc_weight(freq, doc_length); - } + let lengths = 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(); } - push_scored_candidate(&mut candidates, limit, addr, score)?; + let use_global_scorer = impact_safe || exact_scoring_required; + let threshold = if use_global_scorer { + impact_shared_threshold + } else { + local_shared_threshold + }; + let wand_scorer = use_global_scorer.then(|| impact_scorer.clone()); + let part_for_wand = part.clone(); + let candidates = spawn_cpu(move || { + part_for_wand.bm25_search_modern( + lengths.as_ref(), + &visibility, + params.as_ref(), + operator, + postings, + wand_scorer, + metrics.as_ref(), + threshold, + ) + }) + .await?; + Result::Ok(( + partition_ordinal, + PartitionCandidates { + tokens_by_position, + grouped_expansions, + candidates, + }, + )) } + }) + .collect::>(); + + let mut ranked = BinaryHeap::new(); + let mut idf_cache = HashMap::new(); + let mut parts = stream::iter(parts).buffer_unordered(get_num_compute_intensive_cpus()); + while let Some((partition_ordinal, partition)) = parts.try_next().await? { + for (doc_id, score) in rescore_partition_candidates(partition, scorer, &mut idf_cache) { + push_scored_partition_doc( + &mut ranked, + limit, + PartitionDocId::try_new(partition_ordinal, doc_id)?, + score, + ); } } - Ok(candidates - .into_sorted_vec() + Ok(ranked.into_sorted_vec()) + } + + fn resolve_resident_modern_candidates( + &self, + ranked: Vec>, + ) -> Result<(Vec, Vec)> { + let mut addresses = Vec::with_capacity(ranked.len()); + let mut scores = Vec::with_capacity(ranked.len()); + for Reverse(candidate) in ranked { + let partition_ordinal = candidate.document.partition_ordinal(); + let doc_id = candidate.document.doc_id; + let documents = self + .partitions + .get(partition_ordinal) + .and_then(|partition| partition.docs.modern()) + .ok_or_else(|| { + Error::internal(format!( + "resident FTS candidate DocId {} references missing modern partition ordinal {partition_ordinal}", + doc_id.get() + )) + })?; + let address = documents.cached_row_address(doc_id)?.ok_or_else(|| { + Error::internal(format!( + "resident FTS candidate DocId {} in partition ordinal {partition_ordinal} lost its cached address projection", + doc_id.get() + )) + })?; + addresses.push(address.into()); + scores.push(candidate.score.0); + } + Ok((addresses, scores)) + } + + async fn resolve_deferred_modern_candidates( + &self, + ranked: Vec>, + ) -> Result<(Vec, Vec)> { + let mut addresses = vec![0_u64; ranked.len()]; + let mut by_partition = BTreeMap::>::new(); + for (rank, Reverse(candidate)) in ranked.iter().enumerate() { + 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(doc)| (doc.row_id, doc.score.0)) - .unzip()) + .map(|Reverse(candidate)| candidate.score.0) + .collect(); + Ok((addresses, scores)) } async fn load_legacy_index( @@ -1134,16 +1536,18 @@ impl InvertedIndex { store, tokens, inverted_list, - docs: Arc::new(LazyDocSet::from_loaded(docs)), + docs: PartitionDocumentStore::Legacy(Arc::new(docs)), token_set_format: TokenSetFormat::Arrow, })], corpus_stats: Arc::new(OnceCell::new()), + 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`], @@ -1259,6 +1663,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, })) } @@ -1477,7 +1883,17 @@ fn prewarm_chunk_ranges( impl InvertedIndex { pub async fn prewarm_with_options(&self, options: &FtsPrewarmOptions) -> Result<()> { - let with_position = options.with_position; + let mut state = self.prewarm_state.lock().await; + if state.satisfies(options.with_position) { + return Ok(()); + } + self.prewarm_query_state(options.with_position).await?; + state.query_ready = true; + state.positions_ready |= options.with_position; + Ok(()) + } + + async fn prewarm_query_state(&self, with_position: bool) -> Result<()> { let chunk_concurrency = self.store.io_parallelism().max(1); let prewarm_started = Instant::now(); info!( @@ -1511,12 +1927,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, @@ -1533,8 +1945,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" ); @@ -1666,10 +2092,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, } @@ -1710,31 +2135,25 @@ 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, - 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, }) } @@ -1838,7 +2257,7 @@ impl InvertedPartition { doc_ids: &[u32], frequencies: &[u32], block_size: usize, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Vec { @@ -1864,7 +2283,7 @@ impl InvertedPartition { fn union_plain_posting_lists( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Result { @@ -1900,7 +2319,7 @@ impl InvertedPartition { fn union_plain_posting_lists_with_positions( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Result { @@ -1956,7 +2375,7 @@ impl InvertedPartition { fn union_compressed_posting_lists( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Result { @@ -2015,7 +2434,7 @@ impl InvertedPartition { fn union_compressed_posting_lists_with_positions( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, query_weight: f32, scorer: &MemBM25Scorer, ) -> Result { @@ -2079,7 +2498,7 @@ impl InvertedPartition { fn union_posting_lists( postings: Vec, - docs: &DocSet, + docs: &LoadedDocLengths, with_positions: bool, query_weight: f32, scorer: &MemBM25Scorer, @@ -2235,7 +2654,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 }; @@ -2274,7 +2698,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( @@ -2332,38 +2756,92 @@ impl InvertedPartition { }) } - #[instrument(level = "debug", skip_all)] - // Deferred-DocSet adds the `docs` param (caller materializes it) on top of - // the cross-partition `shared_threshold`, tipping this hot-path search fn - // one over the limit. Bundling args isn't worth the churn here. #[allow(clippy::too_many_arguments)] - pub fn bm25_search( + fn bm25_search_legacy( &self, docs: &DocSet, params: &FtsSearchParams, operator: Operator, - mask: Arc, + mask: &RowAddrMask, + postings: Vec, + impact_scorer: Option>, + metrics: &dyn MetricsCollector, + shared_threshold: Arc, + ) -> Result>> { + let documents = LegacyWandDocuments::new(docs, mask); + self.bm25_search_with_documents( + &documents, + params, + operator, + postings, + impact_scorer, + metrics, + shared_threshold, + ) + } + + #[allow(clippy::too_many_arguments)] + fn bm25_search_modern( + &self, + lengths: &DocLengths, + visibility: &DocVisibility, + params: &FtsSearchParams, + operator: Operator, + postings: Vec, + impact_scorer: Option>, + metrics: &dyn MetricsCollector, + shared_threshold: Arc, + ) -> Result>> { + if visibility.is_all() { + let documents = ModernWandDocuments::all(lengths); + self.bm25_search_with_documents( + &documents, + params, + operator, + postings, + impact_scorer, + metrics, + shared_threshold, + ) + } else { + let documents = ModernWandDocuments::filtered(lengths, visibility); + self.bm25_search_with_documents( + &documents, + params, + operator, + postings, + impact_scorer, + metrics, + shared_threshold, + ) + } + } + + #[instrument(level = "debug", skip_all)] + #[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> { + ) -> Result>> { if postings.is_empty() { return Ok(Vec::new()); } - // Caller selects the DocSet shape via `LazyDocSet::docs_for_wand` - // and passes it in here; wand uses `docs.has_row_ids()` to - // handle the num_tokens-only case. let hits = if let Some(scorer) = impact_scorer { - let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) + let mut wand = Wand::new(operator, postings.into_iter(), documents, scorer) .with_shared_threshold(shared_threshold); - wand.search(params, mask, metrics)? + wand.search(params, metrics)? } else { let scorer = IndexBM25Scorer::new(std::iter::once(self)); - let mut wand = Wand::new(operator, postings.into_iter(), docs, scorer) + let mut wand = Wand::new(operator, postings.into_iter(), documents, scorer) .with_shared_threshold(shared_threshold); - wand.search(params, mask, metrics)? + wand.search(params, metrics)? }; Ok(hits) } @@ -2377,10 +2855,7 @@ impl InvertedPartition { self.inverted_list.block_size(), ); builder.tokens = self.tokens.into_mutable(); - // into_builder rewrites every doc, so materialize the full - // DocSet now and clone it out of the Arc. - let docs_arc = self.docs.ensure_loaded().await?; - builder.docs = (*docs_arc).clone(); + builder.docs = self.docs.load_build_docset().await?; builder .posting_lists @@ -2793,6 +3268,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, } @@ -2866,7 +3349,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 } } @@ -2900,6 +3392,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, @@ -2910,6 +3408,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), }) } @@ -3195,6 +3696,11 @@ impl PostingListReader { .clone(), }; + 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).await?; @@ -3204,6 +3710,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 @@ -3491,6 +4078,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()), @@ -3505,7 +4093,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| { @@ -3513,6 +4107,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); @@ -3542,8 +4139,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; @@ -3551,15 +4165,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) }) @@ -3568,7 +4196,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 @@ -6558,9 +7192,7 @@ impl DocSet { } /// Build a `DocSet` from already-loaded `row_id` and `num_tokens` - /// arrow columns. Lets callers that have one column already in hand - /// (e.g. `LazyDocSet` after `total_tokens_num` pre-fetched - /// `num_tokens`) skip re-reading that column. + /// Arrow columns without re-reading either column. pub fn from_columns( row_id_col: &UInt64Array, num_tokens_col: &arrow_array::UInt32Array, @@ -7431,6 +8063,17 @@ mod tests { use super::*; + #[test] + fn address_read_concurrency_respects_payload_budget() { + assert_eq!(address_read_concurrency(64, 0), 64); + assert_eq!(address_read_concurrency(64, 8 * 1024 * 1024), 8); + assert_eq!(address_read_concurrency(64, 16 * 1024 * 1024), 4); + assert_eq!( + address_read_concurrency(64, 2 * MAX_CONCURRENT_ADDRESS_READ_BYTES), + 1 + ); + } + #[derive(Debug)] struct MetadataAccessDeniedStore { inner: Arc, @@ -9364,56 +10007,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(&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()); let filtered = RowAddrMask::allow_nothing(); - let full = partition.docs.docs_for_wand(&filtered).await.unwrap(); - assert!(full.has_row_ids()); - assert!(matches!(&full.num_tokens, NumTokens::Owned(_))); - assert_eq!(full.total_tokens_num(), 100); + let visibility = documents + .visibility(Arc::new(filtered), true) + .await + .unwrap(); + assert!(visibility.is_empty()); + assert!(!documents.projection_loaded()); assert_eq!( - partition.docs.resolve_row_ids(&[0, 99]).await.unwrap(), + documents + .resolve_addresses(&[DocId::new(0), DocId::new(99)]) + .await + .unwrap(), [0, 99] ); } #[tokio::test] - async fn test_concurrent_total_and_num_tokens_view_initialization() { + async fn test_no_hit_partition_does_not_load_document_columns() { + let (index, _counter, _tmpdir) = load_counted_v2_index(100, LanceCache::no_cache()).await; + let 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!(row_ids.is_empty()); + assert!(scores.is_empty()); + assert!(!documents.lengths_loaded()); + assert!(!documents.projection_loaded()); + } + + #[tokio::test] + async fn test_concurrent_stats_and_lengths_initialization() { let (index, _counter, _tmpdir) = load_counted_v2_index(100, LanceCache::no_cache()).await; - let docs = index.partitions[0].docs.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] @@ -10162,13 +10843,128 @@ mod tests { (tmpdir, index) } + #[tokio::test] + async fn test_prewarmed_modern_search_uses_resident_address_projection() { + let (_tmpdir, index) = load_global_scoring_test_index(true).await; + let tokens = Arc::new(Tokens::new(vec!["alpha".to_owned()], DocType::Text)); + let params = Arc::new(FtsSearchParams::new().with_limit(Some(2))); + + assert!(!index.has_resident_document_projections()); + let deferred = index + .bm25_search( + tokens.clone(), + params.clone(), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + assert!(!index.has_resident_document_projections()); + + 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, + params, + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + None, + ) + .await + .unwrap(); + assert_eq!(resident, deferred); + assert_eq!(resident.0.len(), 2); + assert!(resident.0.contains(&100)); + assert!(resident.0.contains(&200)); + } + + #[tokio::test] + async fn test_resident_modern_search_loads_partition_stats_without_global_stats() { + let (_tmpdir, index) = load_global_scoring_test_index(false).await; + assert!(index.corpus_stats.get().is_none()); + assert!( + index + .partitions + .iter() + .all(|partition| partition.docs.cached_stats().is_none()) + ); + + for partition in &index.partitions { + partition.docs.modern().unwrap().projection().await.unwrap(); + } + assert!(index.has_resident_document_projections()); + + let scorer = MemBM25Scorer::new(56_100, 112, HashMap::from([("alpha".to_owned(), 2)])); + let result = index + .bm25_search( + Arc::new(Tokens::new(vec!["alpha".to_owned()], DocType::Text)), + Arc::new(FtsSearchParams::new().with_limit(Some(2))), + Operator::Or, + Arc::new(NoFilter), + Arc::new(NoOpMetricsCollector), + Some(&scorer), + ) + .await + .unwrap(); + + assert_eq!(result.0.len(), 2); + assert!(result.0.contains(&100)); + assert!(result.0.contains(&200)); + assert!(index.corpus_stats.get().is_none()); + assert!( + index + .partitions + .iter() + .all(|partition| partition.docs.cached_stats().is_some()) + ); + } + async fn search_test_impact_partition( partition: &InvertedPartition, tokens: &Tokens, params: &FtsSearchParams, scorer: Arc, shared_threshold: Arc, - ) -> Vec { + ) -> Vec> { let LoadedPostings { postings, grouped_expansions, @@ -10188,24 +10984,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(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] @@ -10263,10 +11056,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)); @@ -10284,10 +11074,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!( @@ -12042,6 +12829,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] @@ -12066,7 +12912,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, @@ -12119,7 +12966,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, @@ -12132,6 +12980,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 025c91cde21..00000000000 --- a/rust/lance-index/src/scalar/inverted/lazy_docset.rs +++ /dev/null @@ -1,405 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors - -//! Deferred-load wrapper around [`DocSet`]. -//! -//! The inverted-index `DocSet` holds the per-doc `row_id` and `num_tokens` -//! arrays for a partition. Eager loading on partition open pulls roughly -//! 12 bytes × num_docs per partition; across thousands of partitions on -//! cold object storage that's tens of GiB of IO before a query has even -//! checked whether a partition contains the term it's looking for. -//! -//! [`LazyDocSet`] defers the load. Cheap sync getters (`len`, -//! `total_tokens_cached`) work without IO; async getters fetch on -//! demand and cache. Wand scoring still needs per-doc num_tokens, but -//! only partitions that actually contribute hits pay -//! `ensure_num_tokens_loaded`/`ensure_loaded`. - -use std::sync::Arc; - -use arrow::array::AsArray; -use arrow::datatypes::{UInt32Type, UInt64Type}; -use arrow_array::{Array, UInt32Array, UInt64Array}; -use lance_core::ROW_ID; -use lance_core::Result; -use lance_core::deepsize::DeepSizeOf; -use tokio::sync::OnceCell; - -use crate::scalar::RowIdRemapper; -use crate::scalar::inverted::index::{DocSet, NUM_TOKEN_COL}; -use crate::scalar::{IndexReader, IndexStore}; -use lance_select::mask::RowAddrMask; - -/// Lazy view over an inverted-index partition's `DocSet`. -/// -/// Two variants: -/// - `Loaded`: a pre-materialized DocSet (legacy paths, tests). -/// Sync accessors return cached values; async accessors return -/// the same DocSet. -/// - `Deferred`: backed by an [`IndexReader`]; columns are read and -/// cached on first request. -pub enum LazyDocSet { - Loaded(LoadedDocSet), - Deferred(Box), -} - -/// Pre-materialized DocSet view -- no reader, no IO. -pub struct LoadedDocSet { - docs: Arc, - num_rows: usize, - total_tokens: u64, -} - -/// Atomically published num-tokens state for deferred scoring. -/// -/// Keeping the Arrow column and the zero-copy `DocSet` view that carries its -/// total in one `OnceCell` prevents cancellation from exposing a partially -/// initialized scoring cache. -struct NumTokensSnapshot { - column: Arc, - docs: Arc, -} - -/// Store-backed DocSet view that loads on demand and caches. -/// -/// Holds the [`IndexStore`] and docs-file path rather than an open -/// [`IndexReader`], so a cached partition does not pin a docs-file -/// handle for its whole lifetime. The reader is re-opened on demand -/// inside each column accessor and dropped when that read completes; -/// because the resulting buffers are cached in the `OnceCell`s below, -/// a contributing partition re-opens only on a cold miss, and a -/// partition that never scores never opens the docs file at all after -/// construction. -pub struct DeferredDocSet { - store: Arc, - docs_path: String, - is_legacy: bool, - frag_reuse_index: Option>, - /// 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, - /// `ROW_ID` arrow buffer cached on first read. - row_ids_col: OnceCell>, - /// Full DocSet, materialized on first `ensure_loaded`. - full: OnceCell>, -} - -impl std::fmt::Debug for LazyDocSet { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Loaded(l) => f - .debug_struct("LazyDocSet::Loaded") - .field("num_rows", &l.num_rows) - .field("total_tokens", &l.total_tokens) - .finish(), - Self::Deferred(d) => f - .debug_struct("LazyDocSet::Deferred") - .field("num_rows", &d.num_rows) - .field( - "total_tokens_loaded", - &(d.num_tokens.initialized() || d.full.initialized()), - ) - .field("num_tokens_loaded", &d.num_tokens.initialized()) - .field("full_loaded", &d.full.initialized()) - .finish(), - } - } -} - -impl DeepSizeOf for LazyDocSet { - fn deep_size_of_children(&self, ctx: &mut lance_core::deepsize::Context) -> usize { - match self { - Self::Loaded(l) => l.docs.deep_size_of_children(ctx), - Self::Deferred(d) => { - d.full - .get() - .map(|d| d.deep_size_of_children(ctx)) - .unwrap_or(0) - + d.num_tokens - .get() - .map(|snapshot| { - let arr: &dyn Array = snapshot.column.as_ref(); - snapshot.docs.deep_size_of_children(ctx) - + arr.deep_size_of_children(ctx) - }) - .unwrap_or(0) - + d.row_ids_col - .get() - .map(|arr| { - let arr: &dyn Array = arr.as_ref(); - arr.deep_size_of_children(ctx) - }) - .unwrap_or(0) - } - } - } -} - -impl LazyDocSet { - #[allow(clippy::too_many_arguments)] - pub fn new( - store: Arc, - docs_path: String, - num_rows: usize, - is_legacy: bool, - frag_reuse_index: Option>, - quantized_scoring: bool, - ) -> Self { - Self::Deferred(Box::new(DeferredDocSet { - store, - docs_path, - is_legacy, - frag_reuse_index, - quantized_scoring, - num_rows, - num_tokens: OnceCell::new(), - row_ids_col: OnceCell::new(), - full: OnceCell::new(), - })) - } - - /// Wrap an already-materialized [`DocSet`]. Used by legacy paths - /// and tests that need to seed a partition without a reader. - pub fn from_loaded(docs: DocSet) -> Self { - let num_rows = docs.len(); - let total_tokens = docs.total_tokens_num(); - Self::Loaded(LoadedDocSet { - docs: Arc::new(docs), - num_rows, - total_tokens, - }) - } - - pub fn len(&self) -> usize { - match self { - Self::Loaded(l) => l.num_rows, - Self::Deferred(d) => d.num_rows, - } - } - - /// Sync read of cached `total_tokens`. Returns `None` for a - /// `Deferred` LazyDocSet that hasn't yet had any of - /// `total_tokens_num` / `ensure_num_tokens_loaded` / `ensure_loaded` - /// run. Used by sync scoring code that has already paid for one - /// of those async calls. - pub fn total_tokens_cached(&self) -> Option { - match self { - Self::Loaded(l) => Some(l.total_tokens), - Self::Deferred(d) => d - .full - .get() - .map(|docs| docs.total_tokens_num()) - .or_else(|| { - d.num_tokens - .get() - .map(|snapshot| snapshot.docs.total_tokens_num()) - }), - } - } - - /// True if this DocSet carries a FragReuseIndex. Callers MUST - /// avoid the deferred-row_id path when this is set: targeted - /// row_id reads return raw stored ids, bypassing the per-id - /// `remap_row_id` filter that `DocSet::from_columns` applies. - pub fn has_frag_reuse_remap(&self) -> bool { - match self { - Self::Loaded(_) => false, - Self::Deferred(d) => d.frag_reuse_index.is_some(), - } - } - - /// Sum of `num_tokens` across all docs. - pub async fn total_tokens_num(&self) -> Result { - match self { - Self::Loaded(l) => Ok(l.total_tokens), - Self::Deferred(d) => d.total_tokens_num().await, - } - } - - /// Materialize the full DocSet, including row_ids. - pub async fn ensure_loaded(&self) -> Result> { - match self { - Self::Loaded(l) => Ok(l.docs.clone()), - Self::Deferred(d) => d.ensure_loaded().await, - } - } - - /// Materialize a DocSet that carries num_tokens but no row_ids. - /// Used by the deferred-row_id scoring path; the per-partition - /// caller resolves surviving doc_ids -> row_ids post-wand via - /// [`Self::resolve_row_ids`]. The tokens-only result is cached separately; - /// a later `ensure_loaded` must still produce a full DocSet. - pub async fn ensure_num_tokens_loaded(&self) -> Result> { - match self { - Self::Loaded(l) => Ok(l.docs.clone()), - Self::Deferred(d) => d.ensure_num_tokens_loaded().await, - } - } - - /// Pick the right DocSet shape for a wand walk under `mask`: - /// the num_tokens-only deferred form when the mask is trivial - /// AND no FragReuseIndex needs to filter row_ids; otherwise the - /// full DocSet. Encapsulates the policy so callers don't have to - /// rederive the conditions for the targeted-read fast path. - pub async fn docs_for_wand(&self, mask: &RowAddrMask) -> Result> { - if mask.is_select_all() && !self.has_frag_reuse_remap() { - self.ensure_num_tokens_loaded().await - } else { - self.ensure_loaded().await - } - } - - /// Resolve a batch of `doc_id`s to their `row_id`s. Used by the - /// deferred-row_id scoring path to map post-wand top-K candidates - /// without going through a full DocSet build. - /// - /// Not safe with a FragReuseIndex (see - /// [`Self::has_frag_reuse_remap`]): the targeted reads return - /// raw stored ids without applying the remap/skip. - pub async fn resolve_row_ids(&self, doc_ids: &[u32]) -> Result> { - match self { - Self::Loaded(l) => Ok(doc_ids.iter().map(|&d| l.docs.row_id(d)).collect()), - Self::Deferred(d) => d.resolve_row_ids(doc_ids).await, - } - } -} - -impl DeferredDocSet { - /// Open a fresh docs-file reader. Dropped by the caller once its read - /// completes, so no handle is pinned across the partition's lifetime. - async fn reader(&self) -> Result> { - self.store.open_index_file(&self.docs_path).await - } - - async fn total_tokens_num(&self) -> Result { - if let Some(full) = self.full.get() { - return Ok(full.total_tokens_num()); - } - Ok(self.num_tokens_snapshot().await?.docs.total_tokens_num()) - } - - async fn num_tokens_snapshot(&self) -> Result<&NumTokensSnapshot> { - self.num_tokens - .get_or_try_init(|| async { - let reader = self.reader().await?; - let batch = reader - .read_range(0..self.num_rows, Some(&[NUM_TOKEN_COL])) - .await?; - let column = Arc::new(batch[NUM_TOKEN_COL].as_primitive::().clone()); - let total_tokens = column.values().iter().map(|&n| n as u64).sum(); - let mut docs = DocSet::from_cached_num_tokens(column.as_ref(), total_tokens); - docs.set_quantized_scoring(self.quantized_scoring); - Result::Ok(NumTokensSnapshot { - column, - docs: Arc::new(docs), - }) - }) - .await - } - - async fn row_ids_column(&self) -> Result> { - self.row_ids_col - .get_or_try_init(|| async { - let reader = self.reader().await?; - let batch = reader.read_range(0..self.num_rows, Some(&[ROW_ID])).await?; - Result::Ok(Arc::new(batch[ROW_ID].as_primitive::().clone())) - }) - .await - .cloned() - } - - async fn ensure_loaded(&self) -> Result> { - let docs = self - .full - .get_or_try_init(|| async { - // If the stats path already pulled NUM_TOKEN_COL, - // read only ROW_ID and rebuild from the two columns. - let mut docs = if let Some(num_tokens) = self.num_tokens.get() { - let row_ids = self.row_ids_column().await?; - DocSet::from_columns( - row_ids.as_ref(), - num_tokens.column.as_ref(), - self.is_legacy, - self.frag_reuse_index.clone(), - )? - } else { - DocSet::load( - self.reader().await?, - self.is_legacy, - self.frag_reuse_index.clone(), - ) - .await? - }; - docs.set_quantized_scoring(self.quantized_scoring); - Result::Ok(Arc::new(docs)) - }) - .await? - .clone(); - Ok(docs) - } - - async fn ensure_num_tokens_loaded(&self) -> Result> { - if let Some(full) = self.full.get() { - return Ok(full.clone()); - } - Ok(self.num_tokens_snapshot().await?.docs.clone()) - } - - async fn resolve_row_ids(&self, doc_ids: &[u32]) -> Result> { - if let Some(full) = self.full.get() - && full.has_row_ids() - { - return Ok(doc_ids.iter().map(|&d| full.row_id(d)).collect()); - } - if let Some(arr) = self.row_ids_col.get() { - return Ok(doc_ids.iter().map(|&d| arr.value(d as usize)).collect()); - } - let ranges: Vec> = doc_ids - .iter() - .map(|&d| d as usize..d as usize + 1) - .collect(); - let reader = self.reader().await?; - let batch = reader.read_ranges(&ranges, Some(&[ROW_ID])).await?; - let arr = batch[ROW_ID].as_primitive::(); - Ok((0..arr.len()).map(|i| arr.value(i)).collect()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::scalar::lance_format::LanceIndexStore; - use lance_core::cache::LanceCache; - use lance_core::utils::tempfile::TempObjDir; - use lance_io::object_store::ObjectStore; - - #[tokio::test] - async fn test_full_docset_is_a_complete_cached_snapshot() { - let temp_dir = TempObjDir::default(); - let store = Arc::new(LanceIndexStore::new( - ObjectStore::local().into(), - temp_dir.clone(), - Arc::new(LanceCache::no_cache()), - )); - let docs = LazyDocSet::new(store, "unused".to_owned(), 3, false, None, false); - assert_eq!(docs.total_tokens_cached(), None); - - let row_ids = UInt64Array::from(vec![10, 20, 30]); - let num_tokens = UInt32Array::from(vec![3, 5, 8]); - let full = Arc::new(DocSet::from_columns(&row_ids, &num_tokens, false, None).unwrap()); - let LazyDocSet::Deferred(deferred) = &docs else { - panic!("expected a deferred DocSet"); - }; - deferred.full.set(full.clone()).unwrap(); - - let wand_docs = docs.ensure_num_tokens_loaded().await.unwrap(); - assert!(Arc::ptr_eq(&wand_docs, &full)); - assert_eq!(wand_docs.total_tokens_num(), 16); - assert_eq!(docs.total_tokens_cached(), Some(16)); - } -} diff --git a/rust/lance-index/src/scalar/inverted/scorer.rs b/rust/lance-index/src/scalar/inverted/scorer.rs index 3a33a67ff7a..37b3dc5381c 100644 --- a/rust/lance-index/src/scalar/inverted/scorer.rs +++ b/rust/lance-index/src/scalar/inverted/scorer.rs @@ -164,24 +164,20 @@ pub struct IndexBM25Scorer<'a> { } impl<'a> IndexBM25Scorer<'a> { - /// Sync constructor. Reads each partition's cached `total_tokens` via - /// `LazyDocSet::total_tokens_cached()`; callers must have already - /// populated it (via `ensure_loaded`, `ensure_num_tokens_loaded`, or - /// `total_tokens_num`). Panics with a clear message otherwise — this - /// is the wand-scoring path where the contract is statically known. + /// Sync constructor. Query setup populates immutable partition stats + /// before entering the CPU-only WAND executor. pub fn new(partitions: impl Iterator) -> Self { let partitions = partitions.collect::>(); - let num_docs = partitions.iter().map(|p| p.docs.len()).sum(); - let total_tokens: u64 = partitions + let stats = partitions .iter() - .map(|p| { - p.docs.total_tokens_cached().expect( - "IndexBM25Scorer::new requires each partition's total_tokens to be \ - cached; call `ensure_loaded` / `ensure_num_tokens_loaded` / \ - `total_tokens_num` first", + .map(|partition| { + partition.docs.cached_stats().expect( + "IndexBM25Scorer::new requires partition stats to be loaded before WAND", ) }) - .sum(); + .collect::>(); + let num_docs = stats.iter().map(|stats| stats.num_docs).sum(); + let total_tokens: u64 = stats.iter().map(|stats| stats.total_tokens).sum(); let avgdl = total_tokens as f32 / num_docs as f32; Self { partitions, diff --git a/rust/lance-index/src/scalar/inverted/wand.rs b/rust/lance-index/src/scalar/inverted/wand.rs index b6e2972bdbe..dcd8f5ac0d0 100644 --- a/rust/lance-index/src/scalar/inverted/wand.rs +++ b/rust/lance-index/src/scalar/inverted/wand.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::ops::Deref; use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::{Arc, LazyLock}; use std::{ @@ -22,6 +21,7 @@ use crate::metrics::MetricsCollector; use super::{ CompressedPositionStorage, + documents::{DocId, DocLengths, DocVisibility}, impact::{IMPACT_LEVEL1_BLOCKS, ImpactScoreCache, ImpactSkipData}, index::{PositionStreamCodec, dequantize_doc_length}, query::Operator, @@ -168,10 +168,10 @@ impl TopKCollector { } } - fn into_candidates( + fn into_candidates( self, - mut to_addr: impl FnMut(u64) -> CandidateAddr, - ) -> Result> { + mut to_candidate: impl FnMut(u64) -> C, + ) -> Result>> { let Self { heap, mut frequency_slots, @@ -181,7 +181,7 @@ impl TopKCollector { .map( |Reverse((doc, doc_length, posting_doc_id, frequency_slot))| { Ok(DocCandidate { - addr: to_addr(doc.row_id), + document: to_candidate(doc.row_id), posting_doc_id, freqs: frequency_slots.take(frequency_slot)?, doc_length, @@ -1318,12 +1318,12 @@ impl PostingIterator { // The norm cache arg tips this hot-path fn over the limit; bundling the // scoring inputs isn't worth the churn here. #[allow(clippy::too_many_arguments)] - fn collect_window_scores( + fn collect_window_scores( &mut self, window_min: u64, up_to: u64, clause_idx: usize, - docs: &DocSet, + docs: &D, scorer: &S, norm_k: Option<(&[u8], &[f32; 256])>, acc: &mut WindowAccumulator, @@ -1392,10 +1392,7 @@ impl PostingIterator { if doc_id > up_to { break; } - let doc_length = match &doc { - DocInfo::Raw(raw) => docs.scoring_num_tokens(raw.doc_id), - DocInfo::Located(located) => docs.num_tokens_by_row_id(located.row_id), - }; + let doc_length = docs.doc_length(&doc); let score = self.score(scorer, doc.frequency(), doc_length); let slot = (doc_id - window_min) as usize; acc.add(clause_idx, slot, score, doc.frequency()); @@ -1417,6 +1414,55 @@ impl PostingIterator { PostingList::Plain(ref plain) => plain.row_ids.get(self.index + 1).cloned(), } } + + #[cfg(test)] + fn validate_modern_doc_ids(&self, num_docs: usize) -> Result<()> { + validate_modern_posting_doc_ids(&self.list, &self.token, num_docs) + } +} + +/// Validate the dense-DocId boundary before any document-length lookup. +/// Modern postings are sorted, so decoding only the final physical block is +/// sufficient to validate their maximum DocId. +pub(super) fn validate_modern_posting_doc_ids( + posting: &PostingList, + token: &str, + num_docs: usize, +) -> Result<()> { + let PostingList::Compressed(list) = posting else { + return Err(Error::index(format!( + "modern FTS posting for token {token:?} uses a legacy row-address layout" + ))); + }; + if list.length == 0 { + return Ok(()); + } + let last_block_idx = list.blocks.len().checked_sub(1).ok_or_else(|| { + Error::index(format!( + "modern FTS posting for token {token:?} has length {} but no blocks", + list.length + )) + })?; + let mut state = CompressedState::new(list.block_size); + state.decompress( + list.blocks.value(last_block_idx), + last_block_idx, + list.blocks.len(), + list.length, + list.posting_tail_codec, + list.block_size, + ); + let max_doc_id = state.doc_ids.last().copied().ok_or_else(|| { + Error::index(format!( + "modern FTS posting for token {token:?} has an empty final block" + )) + })?; + if max_doc_id as usize >= num_docs { + return Err(Error::index(format!( + "modern FTS posting for token {token:?} contains DocId {max_doc_id}, outside [0, {num_docs})" + ))); + } + Ok(()) } /// Inner window span (in doc ids) of the bulk MAXSCORE path. Same as Lucene's @@ -1486,19 +1532,9 @@ impl WindowAccumulator { } } -/// How wand identified a candidate: either it already had the real -/// row_id (DocSet carried row_ids), or only the partition-local -/// doc_id (deferred-row_id path; the caller must resolve via -/// [`super::lazy_docset::LazyDocSet::resolve_row_ids`]). -#[derive(Debug, Clone, Copy)] -pub enum CandidateAddr { - RowId(u64), - Pending(u32), -} - #[derive(Debug)] -pub struct DocCandidate { - pub addr: CandidateAddr, +pub struct DocCandidate { + pub document: C, /// The document key used by the posting lists: doc_id for compressed /// postings, row_id for legacy plain postings. pub posting_doc_id: u64, @@ -1507,6 +1543,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 @@ -1626,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, @@ -1663,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 @@ -1686,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(); @@ -1730,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, } @@ -1757,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))?; @@ -1796,23 +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![]); } - match (mask.max_len(), mask.iter_addrs()) { - (Some(num_rows_matched), Some(row_ids)) - if self.operator == Operator::Or - && num_rows_matched * 100 - <= FLAT_SEARCH_PERCENT_THRESHOLD.deref() * self.docs.len() as u64 => - { - return self.flat_search(params, row_ids, metrics); - } - _ => {} + if self.operator == Operator::Or + && let Some((num_docs_selected, documents)) = self.documents.flat_documents() + && num_docs_selected.saturating_mul(100) + <= (*FLAT_SEARCH_PERCENT_THRESHOLD as usize).saturating_mul(self.documents.len()) + { + return self.flat_search(params, documents, metrics); } // Top-k disjunctions over compressed lists can opt into the bulk @@ -1827,7 +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 @@ -1849,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 { @@ -1874,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 { + let Some(document_key) = self.documents.document_key(&doc) else { if self.operator == Operator::Or { self.push_back_leads(doc.doc_id() + 1); } continue; - } - if docs_has_row_ids && !mask.selected(row_id) { - if self.operator == Operator::Or { - self.push_back_leads(doc.doc_id() + 1); - } - continue; - } - - let doc_length = match &doc { - DocInfo::Raw(doc) => self.docs.scoring_num_tokens(doc.doc_id), - DocInfo::Located(doc) => self.docs.num_tokens_by_row_id(doc.row_id), }; + let doc_length = self.documents.doc_length(&doc); + let score = if self.operator == Operator::Or { self.advance_all_tail(doc.doc_id(), Some(doc_length), Some(&mut score)); if params.phrase_slop.is_some() @@ -1932,7 +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(), @@ -1969,45 +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![]); } - // we need to map the row ids to doc ids, and sort them, - // because WAND PostingIterator can't go back to the previous doc id. - // A list column maps one row id to several doc ids, so expand every - // document the row owns — keying on a single doc id would drop matches - // at non-last list positions (lancedb#3352). - let doc_ids = row_ids - .flat_map(|row_addr| { - let row_id: u64 = row_addr.into(); - self.docs - .doc_ids(row_id) - .map(move |doc_id| (doc_id, row_id)) - }) - .sorted_unstable() - .collect::>(); + // Posting iterators are forward-only, so selected DocIds are sorted + // before driving the sparse executor. + let documents = documents.sorted_unstable().collect::>(); let is_compressed = self .head .peek() @@ -2021,7 +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); @@ -2051,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 @@ -2068,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(), @@ -2081,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. @@ -2098,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, @@ -2108,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() @@ -2284,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 @@ -2297,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() { @@ -2337,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), ), }; } @@ -2348,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( @@ -2456,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, @@ -2491,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 @@ -2537,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) } @@ -2547,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)| { @@ -2586,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 @@ -2692,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); @@ -2774,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)); @@ -2811,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; @@ -3228,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])); @@ -3236,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)); } } }, @@ -3271,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 @@ -3320,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( @@ -3362,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) { @@ -4263,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::>(); @@ -4652,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!( @@ -4666,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 @@ -4713,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 } @@ -4770,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) @@ -4842,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()); } @@ -4894,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)) }; @@ -4955,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); @@ -5251,7 +5417,6 @@ mod tests { let result = wand .search( &FtsSearchParams::new().with_limit(Some(1)), - Arc::new(RowAddrMask::default()), &NoOpMetricsCollector, ) .unwrap(); @@ -5472,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); @@ -5775,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 @@ -5830,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 @@ -5894,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] @@ -5972,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(), @@ -5983,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]); } @@ -6039,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(), @@ -6048,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:?}" ); } @@ -6078,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; @@ -6380,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| { @@ -6388,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::>(); @@ -6407,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) };