diff --git a/rust/lance/src/dataset/index.rs b/rust/lance/src/dataset/index.rs index 7cac7815125..c2c24e6f19a 100644 --- a/rust/lance/src/dataset/index.rs +++ b/rust/lance/src/dataset/index.rs @@ -73,7 +73,12 @@ impl IndexRemapper for DatasetIndexRemapper { .any(|frag_idx| affected_frag_ids.contains(&(frag_idx as u64))), }; if needs_remapped { - let remap_result = self.remap_index(index, &mapping).await?; + // Box the remap future at the call site: inlining `remap_index` into this + // loop's async layout otherwise exceeds rustc's depth limit. It has to be + // boxed here, not inside `remap_index` — boxing internally turns the + // future's `Send` check into a `Box: Send` trait obligation that + // overflows the solver through the cache types (E0275 downstream). + let remap_result = Box::pin(self.remap_index(index, &mapping)).await?; match remap_result { RemapResult::Drop => continue, RemapResult::Keep(id) => { diff --git a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs index d80e0e26795..c39af5145b0 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs @@ -654,7 +654,11 @@ impl LsmPointLookupPlanner { scanner.with_row_address(); } scanner.filter_expr(filter.clone()); - scanner.create_plan().await? + // Box at the call site: `create_plan`'s inlined async layout exceeds + // rustc's depth limit up this point-lookup chain, and boxing inside + // `create_plan` instead triggers a `Box: Send` solver overflow + // (E0275 downstream). Same for the other arms below. + Box::pin(scanner.create_plan()).await? } LsmDataSource::FlushedMemTable { path, .. } => { let dataset = open_flushed_dataset( @@ -672,7 +676,7 @@ impl LsmPointLookupPlanner { let cols = cols_with_tombstone(&cols, dataset.schema().field(TOMBSTONE).is_some()); scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; scanner.filter_expr(filter.clone()); - scanner.create_plan().await? + Box::pin(scanner.create_plan()).await? } LsmDataSource::ActiveMemTable { batch_store, @@ -695,7 +699,7 @@ impl LsmPointLookupPlanner { // over insert-ordered scan would return the *oldest* of // multiple rows sharing the target primary key. scanner.with_row_id(); - let raw = scanner.create_plan().await?; + let raw = Box::pin(scanner.create_plan()).await?; // The filter already restricts to the exact PK value, so the // scan yields that key's insert history. Within the active // memtable larger `_rowid` = newer insert, so sorting `_rowid` diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs index d73329d8a41..6afd8d56eca 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -44,11 +44,129 @@ use lance_core::datatypes::{Field, Schema}; use lance_core::{Error, Result}; use roaring::RoaringBitmap; -use lance_table::format::DataFile; +use lance_table::format::overlay::DataOverlayFile; +use lance_table::format::{DataFile, Fragment, IndexMetadata}; use lance_table::utils::stream::ReadBatchFut; use crate::dataset::fragment::{FileFragment, FragReadConfig, GenericFileReader}; +/// The physical offsets within a fragment whose value for an indexed field may be +/// stale relative to an index built at `index_version`, and so must be excluded +/// from that index's results and re-evaluated against current values on the flat +/// path. +/// +/// The set is the union, over every overlay whose `committed_version` is newer +/// than `index_version`, of that overlay's coverage **restricted to the indexed +/// fields**. The restriction makes exclusion field-aware: an overlay that touches +/// only non-indexed fields contributes nothing. An overlay whose +/// `committed_version <= index_version` is already incorporated by the index and +/// is ignored. +pub fn overlay_exclusion_offsets( + overlays: &[DataOverlayFile], + indexed_field_ids: &[i32], + index_version: u64, +) -> Result { + let mut excluded = RoaringBitmap::new(); + for overlay in overlays { + if overlay.committed_version <= index_version { + continue; + } + for (field_pos, field_id) in overlay.data_file.fields.iter().enumerate() { + if indexed_field_ids.contains(field_id) { + excluded |= &*overlay.coverage_for_field(field_pos)?; + } + } + } + Ok(excluded) +} + +// Stale row offsets contributed by one fragment's overlays for a given index version. +// Applies a cheap version gate first: if every overlay predates the segment it is already +// incorporated by the index, so there is nothing stale and the field/bitmap work is skipped. +fn stale_offsets_for_fragment( + fragment: &Fragment, + fields: &[i32], + index_version: u64, +) -> Result { + if fragment + .overlays + .iter() + .all(|o| o.committed_version <= index_version) + { + return Ok(RoaringBitmap::new()); + } + overlay_exclusion_offsets(&fragment.overlays, fields, index_version) +} + +// A missing `fragment_bitmap` means the index predates fragment-bitmap tracking; treat it as +// covering every fragment (matching `DatasetPreFilter::new`) so overlay-stale rows can't slip +// through unmasked. Only skip fragments explicitly absent from a present bitmap. +fn covers_fragment(coverage: Option<&RoaringBitmap>, frag_id: u32) -> bool { + coverage.is_none_or(|c| c.contains(frag_id)) +} + +/// Index by fragment id the fragments that carry at least one overlay. Overlays are rare, so +/// this is empty on the common path, letting callers skip index loading entirely; when non-empty +/// it bounds the stale-collection loops to `O(overlaid fragments)`. +pub fn overlaid_fragments(fragments: &[Fragment]) -> HashMap { + fragments + .iter() + .filter(|f| !f.overlays.is_empty()) + .map(|f| (f.id as u32, f)) + .collect() +} + +/// Insert into `stale` the ids of fragments covered by `segment` whose index entries may be +/// stale because an overlay committed after the segment was built touches a field the segment +/// indexes. Field-aware and version-gated via [`overlay_exclusion_offsets`]. +/// +/// `overlaid_frags` holds only the fragments that actually carry overlays (rare), so the loop is +/// `O(overlaid_frags)` rather than `O(fragments the segment covers)`. +pub fn collect_overlay_stale_frags( + segment: &IndexMetadata, + overlaid_frags: &HashMap, + stale: &mut RoaringBitmap, +) -> Result<()> { + let coverage = segment.fragment_bitmap.as_ref(); + for (&frag_id, fragment) in overlaid_frags { + if stale.contains(frag_id) || !covers_fragment(coverage, frag_id) { + continue; + } + if !stale_offsets_for_fragment(fragment, &segment.fields, segment.dataset_version)? + .is_empty() + { + stale.insert(frag_id); + } + } + Ok(()) +} + +/// Like [`collect_overlay_stale_frags`] but with row-level granularity: instead of marking the +/// whole fragment stale, it computes exactly which row offsets within each covered fragment are +/// stale and accumulates them into `stale` (fragment_id → stale row offsets). +/// +/// Used by the scalar and vector paths to block only the affected rows from index results and +/// re-evaluate only those rows on the flat path, keeping overhead proportional to the number of +/// overlaid rows rather than the whole fragment size. +pub fn collect_overlay_stale_rows_for_segment( + segment: &IndexMetadata, + overlaid_frags: &HashMap, + stale: &mut HashMap, +) -> Result<()> { + let coverage = segment.fragment_bitmap.as_ref(); + for (&frag_id, fragment) in overlaid_frags { + if !covers_fragment(coverage, frag_id) { + continue; + } + let excluded = + stale_offsets_for_fragment(fragment, &segment.fields, segment.dataset_version)?; + if !excluded.is_empty() { + *stale.entry(frag_id).or_default() |= &excluded; + } + } + Ok(()) +} + /// The plan for merging one field's overlays into one batch: which source (base or /// a particular overlay) supplies each output row, and which overlay values must be /// fetched to do it. @@ -759,6 +877,7 @@ async fn fetch_overlay_values( mod tests { use super::*; use arrow_array::{Int32Array, StringArray, UInt32Array}; + use lance_table::format::overlay::OverlayCoverage; use std::sync::Arc; fn i32_array(values: impl IntoIterator>) -> ArrayRef { @@ -1061,4 +1180,175 @@ mod tests { assert!(spliced.is_null(1)); assert!(!spliced.is_null(2)); } + + /// A dense overlay covering `offsets` for `field_ids`, committed at `version`. + fn dense_overlay( + field_ids: Vec, + offsets: impl IntoIterator, + version: u64, + ) -> lance_table::format::overlay::DataOverlayFile { + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", field_ids, None), + coverage: OverlayCoverage::dense(bitmap(offsets)), + committed_version: version, + } + } + + #[test] + fn test_exclusion_offsets_version_gate() { + // index built at version 5; only overlays committed > 5 are excluded. + let overlays = vec![ + dense_overlay(vec![3], [0, 1], 4), + dense_overlay(vec![3], [2, 7], 6), + ]; + let excluded = overlay_exclusion_offsets(&overlays, &[3], 5).unwrap(); + assert_eq!(excluded, bitmap([2, 7])); + // An overlay exactly at the index version is already incorporated. + let overlays = vec![dense_overlay(vec![3], [9], 5)]; + assert!( + overlay_exclusion_offsets(&overlays, &[3], 5) + .unwrap() + .is_empty() + ); + } + + #[test] + fn test_exclusion_offsets_is_field_aware() { + // An overlay touching only an unrelated field excludes nothing. + let overlays = vec![dense_overlay(vec![2], [0, 1, 2], 9)]; + assert!( + overlay_exclusion_offsets(&overlays, &[3], 1) + .unwrap() + .is_empty() + ); + // The union spans only the indexed fields the overlay actually carries. + let overlays = vec![dense_overlay(vec![2, 3], [4], 9)]; + assert_eq!( + overlay_exclusion_offsets(&overlays, &[3], 1).unwrap(), + bitmap([4]) + ); + } + + #[test] + fn test_exclusion_offsets_sparse_per_field() { + // Sparse overlay: field 2 covers {2,3}, field 4 covers {1}. + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![2, 4], None), + coverage: OverlayCoverage::sparse(vec![bitmap([2, 3]), bitmap([1])]), + committed_version: 9, + }; + let overlays = vec![overlay]; + // Only the bitmap for the indexed field (4) contributes. + assert_eq!( + overlay_exclusion_offsets(&overlays, &[4], 1).unwrap(), + bitmap([1]) + ); + assert_eq!( + overlay_exclusion_offsets(&overlays, &[2], 1).unwrap(), + bitmap([2, 3]) + ); + } + + #[test] + fn test_exclusion_offsets_unions_multiple_overlays() { + let overlays = vec![ + dense_overlay(vec![3], [1], 6), + dense_overlay(vec![3], [4, 5], 7), + ]; + assert_eq!( + overlay_exclusion_offsets(&overlays, &[3], 1).unwrap(), + bitmap([1, 4, 5]) + ); + } + + /// An index segment covering `fields`, built at `dataset_version`, with the given + /// fragment coverage (`None` = legacy index predating fragment-bitmap tracking). + fn segment( + fields: Vec, + dataset_version: u64, + fragment_bitmap: Option, + ) -> IndexMetadata { + IndexMetadata { + uuid: uuid::Uuid::new_v4(), + name: "idx".into(), + fields, + dataset_version, + fragment_bitmap, + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + } + } + + fn fragment_with_overlay(id: u64, overlay: DataOverlayFile) -> Fragment { + let mut fragment = Fragment::new(id); + fragment.overlays.push(overlay); + fragment + } + + #[test] + fn test_collect_frags_missing_bitmap_covers_all() { + // A segment with no fragment_bitmap (legacy index predating bitmap tracking) must treat + // every overlaid fragment as covered so stale rows can't leak past the index unmasked. + let fragment = fragment_with_overlay(3, dense_overlay(vec![3], [1, 2], 9)); + let overlaid: HashMap = HashMap::from([(3u32, &fragment)]); + + let mut stale = RoaringBitmap::new(); + collect_overlay_stale_frags(&segment(vec![3], 1, None), &overlaid, &mut stale).unwrap(); + assert_eq!(stale, bitmap([3]), "missing bitmap must cover fragment 3"); + + // A present bitmap that excludes fragment 3 leaves it untouched. + let mut stale = RoaringBitmap::new(); + collect_overlay_stale_frags( + &segment(vec![3], 1, Some(bitmap([0]))), + &overlaid, + &mut stale, + ) + .unwrap(); + assert!( + stale.is_empty(), + "fragment absent from bitmap is not covered" + ); + + // A present bitmap that includes fragment 3 marks it stale. + let mut stale = RoaringBitmap::new(); + collect_overlay_stale_frags( + &segment(vec![3], 1, Some(bitmap([3]))), + &overlaid, + &mut stale, + ) + .unwrap(); + assert_eq!(stale, bitmap([3])); + } + + #[test] + fn test_collect_rows_missing_bitmap_covers_all() { + // Same covers-all guarantee at row-level granularity. + let fragment = fragment_with_overlay(3, dense_overlay(vec![3], [1, 2], 9)); + let overlaid: HashMap = HashMap::from([(3u32, &fragment)]); + + let mut stale = HashMap::new(); + collect_overlay_stale_rows_for_segment(&segment(vec![3], 1, None), &overlaid, &mut stale) + .unwrap(); + assert_eq!( + stale.get(&3), + Some(&bitmap([1, 2])), + "missing bitmap must cover fragment 3" + ); + + // A present bitmap that excludes fragment 3 yields no stale rows. + let mut stale = HashMap::new(); + collect_overlay_stale_rows_for_segment( + &segment(vec![3], 1, Some(bitmap([0]))), + &overlaid, + &mut stale, + ) + .unwrap(); + assert!( + stale.is_empty(), + "fragment absent from bitmap contributes no rows" + ); + } } diff --git a/rust/lance/src/dataset/rowids.rs b/rust/lance/src/dataset/rowids.rs index 59c1b4e40cc..dbb28bd8305 100644 --- a/rust/lance/src/dataset/rowids.rs +++ b/rust/lance/src/dataset/rowids.rs @@ -6,6 +6,7 @@ use crate::session::caches::{RowIdIndexKey, RowIdSequenceKey}; use crate::{Error, Result}; use futures::{Stream, StreamExt, TryFutureExt, TryStreamExt}; use lance_core::utils::deletion::DeletionVector; +use lance_select::{RowAddrSelection, RowAddrTreeMap}; use lance_table::{ format::{Fragment, RowIdMeta}, rowids::{FragmentRowIdIndex, RowIdIndex, RowIdSequence, read_row_ids}, @@ -87,6 +88,90 @@ pub async fn get_row_id_index( } } +/// Map a set of physical row addresses to their stable row ids +/// +/// For each fragment present in `addrs`, the live rows in physical order carry +/// the stable ids yielded by the fragment's [`RowIdSequence`] in the same +/// order. Zipping the two (skipping deleted physical offsets) gives the +/// `physical offset -> stable id` mapping. Addresses that point at deleted rows +/// have no live counterpart and are dropped, which is correct: those rows are +/// not part of the answer. +pub(crate) async fn translate_addr_treemap_to_row_ids( + dataset: &Dataset, + addrs: &RowAddrTreeMap, +) -> Result { + let mut row_ids = RowAddrTreeMap::new(); + for (fragment_id, selection) in addrs.iter() { + let file_fragment = dataset.get_fragment(*fragment_id as usize).ok_or_else(|| { + Error::internal(format!( + "fragment {fragment_id} referenced by an address-domain index result \ + was not found in the dataset" + )) + })?; + let sequence = load_row_id_sequence(dataset, file_fragment.metadata()).await?; + + match selection { + RowAddrSelection::Full => { + // The whole fragment is selected: every live row's id qualifies. + row_ids |= RowAddrTreeMap::from(sequence.as_ref()); + } + RowAddrSelection::Partial(offsets) => { + let Some(max_offset) = offsets.max() else { + continue; + }; + let (deletion_vector, num_physical_rows) = futures::try_join!( + file_fragment.get_deletion_vector(), + file_fragment.physical_rows() + )?; + let num_physical_rows = u32::try_from(num_physical_rows).map_err(|_| { + Error::internal(format!( + "fragment_id={fragment_id} has num_physical_rows={num_physical_rows}, \ + which exceeds the maximum representable physical offset" + )) + })?; + if max_offset >= num_physical_rows { + return Err(Error::internal(format!( + "fragment_id={fragment_id} selection has max_offset={max_offset}, \ + but num_physical_rows={num_physical_rows}" + ))); + } + let mut ids = sequence.iter(); + for physical_offset in 0..num_physical_rows { + if physical_offset > max_offset { + break; + } + let deleted = deletion_vector + .as_ref() + .is_some_and(|dv| dv.contains(physical_offset)); + if deleted { + continue; + } + match ids.next() { + Some(id) => { + if offsets.contains(physical_offset) { + row_ids.insert(id); + } + } + // The sequence yields one id per live row, so it can only + // run dry before `max_offset` if it holds fewer ids than the + // fragment has live rows. Breaking would silently drop the + // remaining selected offsets and let stale index results + // escape masking, so treat the mismatch as corruption. + None => { + return Err(Error::internal(format!( + "fragment_id={fragment_id} row-id sequence exhausted at \ + physical_offset={physical_offset} before reaching \ + max_offset={max_offset} (num_physical_rows={num_physical_rows})" + ))); + } + } + } + } + } + } + Ok(row_ids) +} + /// Resolve row addresses to row ids, positionally: `addr = (fragment << 32) | offset` /// looks up `offset` in the fragment's [`RowIdSequence`]. The inverse companion of /// [`get_row_id_index`]. diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index af9316d3a6f..fa965d6b7f2 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use datafusion::config::ConfigOptions; use lance_select::result::IndexExprResultWireFormat; @@ -17,7 +17,7 @@ use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaR use arrow_select::concat::concat_batches; use async_recursion::async_recursion; use chrono::Utc; -use datafusion::common::{DFSchema, JoinType, NullEquality, SchemaExt, exec_datafusion_err}; +use datafusion::common::{DFSchema, JoinType, NullEquality, exec_datafusion_err}; use datafusion::functions_aggregate; use datafusion::logical_expr::{Expr, ScalarUDF, col, lit}; use datafusion::physical_expr::PhysicalSortExpr; @@ -83,11 +83,15 @@ use tracing::{Span, info_span, instrument}; use uuid::Uuid; use super::Dataset; +use crate::dataset::overlay::{ + collect_overlay_stale_frags, collect_overlay_stale_rows_for_segment, overlaid_fragments, +}; use crate::dataset::row_offsets_to_row_addresses; +use crate::dataset::rowids::translate_addr_treemap_to_row_ids; use crate::dataset::utils::SchemaAdapter; use crate::index::DatasetIndexInternalExt; use crate::index::scalar::inverted::{load_segment_details, load_segments}; -use crate::index::scalar_logical::scalar_index_fragment_bitmap; +use crate::index::scalar_logical::{load_named_scalar_segments, scalar_index_fragment_bitmap}; use crate::index::vector::utils::{ default_distance_type_for, get_vector_dim, get_vector_type, validate_distance_type_for, }; @@ -2908,6 +2912,8 @@ impl Scanner { fragments: Option>>, scan_range: Option>, ) -> Result> { + // Kept for the overlay stale-Take path below, which re-evaluates blocked stale rows. + let user_projection = projection.clone(); let mut read_options = FilteredReadOptions::basic_full_read(&self.dataset) .with_filter_plan(filter_plan.clone()) .with_projection(projection); @@ -2949,6 +2955,24 @@ impl Scanner { read_options = read_options.with_only_indexed_fragments(); } + // Mask data overlay files: a row with an overlay committed after an index it relies on + // touched an indexed field can no longer be trusted to that index. Block just those rows + // from the index result (their fragments stay indexed, so non-stale rows keep the index) + // and re-evaluate them on a targeted take path below — O(stale_rows), not O(fragment). + let mut overlay_stale_rows: HashMap = HashMap::new(); + if let Some(index_query) = filter_plan.index_query.as_ref() { + let candidate_frags = read_options + .fragments + .clone() + .unwrap_or_else(|| self.dataset.fragments().clone()); + overlay_stale_rows = self + .overlay_stale_index_rows(index_query, &candidate_frags) + .await?; + if let Some(block) = self.stale_rows_block_mask(&overlay_stale_rows).await? { + read_options = read_options.with_overlay_block(block); + } + } + let result_format = self.index_expr_result_format(); let index_input = filter_plan.index_query.clone().map(|index_query| { Arc::new(ScalarIndexExec::new( @@ -2958,10 +2982,36 @@ impl Scanner { )) as Arc }); - Ok(Arc::new(FilteredReadExec::try_new( + let plan: Arc = Arc::new(FilteredReadExec::try_new( self.dataset.clone(), read_options, index_input, + )?); + + if overlay_stale_rows.is_empty() { + return Ok(plan); + } + + // Stale-Take path: take the stale rows' current (overlay-merged) values and re-apply the + // full filter, then union with the indexed read. These rows were blocked from the index + // result above, so this is the only path that can surface them. + let filter = filter_plan.full_expr.as_ref().expect_ok()?; + let filter_cols = Planner::column_names_in_expr(filter); + let take_projection = user_projection.union_columns(filter_cols, OnMissing::Error)?; + + let stale_node = self + .stale_rows_take(&overlay_stale_rows, take_projection) + .await?; + let planner = Planner::new(stale_node.schema()); + let optimized_filter = planner.optimize_expr(filter.clone())?; + let filtered = Arc::new(LanceFilterExec::try_new(optimized_filter, stale_node)?); + let stale_path: Arc = + Arc::new(project(filtered, plan.schema().as_ref())?); + + let unioned = UnionExec::try_new(vec![plan, stale_path])?; + Ok(Arc::new(RepartitionExec::try_new( + unioned, + datafusion::physical_plan::Partitioning::RoundRobinBatch(1), )?)) } @@ -3329,6 +3379,10 @@ impl Scanner { self.fragments_covered_by_fts_query(&query).await?, ) .await?; + // Data overlay masking: match queries drop stale segments and re-evaluate the affected + // fragments on the flat-text path (`plan_match_query` / `fts_stale_frags_and_fresh_segments`), + // and phrase queries exclude stale segments (`plan_phrase_query`). Both keep stale index + // hits out of the result. let fts_exec = self .plan_fts(&query, ¶ms, filter_plan, &prefilter_source) .await?; @@ -3521,12 +3575,40 @@ impl Scanner { .to_string())); } - Ok(Arc::new(PhraseQueryExec::new( - self.dataset.clone(), - query.clone(), - params.clone(), - prefilter_source.clone(), - ))) + // Mask data overlay files: a fragment with an overlay committed after this FTS index can + // no longer be trusted to its inverted-index positions, so a stale phrase could still + // match. Exclude any segment covering such a fragment. Unlike match queries, phrase + // queries have no flat re-evaluation path, so — exactly as for unindexed fragments, which + // phrase queries already do not search — overlaid fragments are simply dropped from the + // phrase result; new phrase matches there surface once compaction folds the overlay into + // the base. This removes stale hits without introducing wrong ones. + let target_fragments = self + .fragments + .clone() + .unwrap_or_else(|| self.dataset.fragments().to_vec()); + let (_flat_frag_ids, fresh_segments) = self + .fts_stale_frags_and_fresh_segments(&column, &target_fragments) + .await?; + + let exec: Arc = match fresh_segments { + // Every segment covers a stale fragment: with no flat phrase path there is nothing + // trustworthy left to search, so the phrase query returns no rows. + Some(segs) if segs.is_empty() => Arc::new(EmptyExec::new(FTS_SCHEMA.clone())), + Some(segs) => Arc::new(PhraseQueryExec::new_with_segments( + self.dataset.clone(), + query.clone(), + params.clone(), + prefilter_source.clone(), + segs, + )), + None => Arc::new(PhraseQueryExec::new( + self.dataset.clone(), + query.clone(), + params.clone(), + prefilter_source.clone(), + )), + }; + Ok(exec) } async fn plan_match_query( @@ -3561,30 +3643,62 @@ impl Scanner { let unindexed_fragments = self .retain_target_fragments(self.dataset.unindexed_fragments(&index.name).await?); - // If all target fragments are unindexed, skip index entirely - if unindexed_fragments.len() == target_fragments.len() { + // Fragments whose FTS index entries may be stale due to a newer data overlay. + // These are excluded from the indexed path and re-evaluated on the flat path. + let (stale_flat_frag_ids, fresh_segments) = self + .fts_stale_frags_and_fresh_segments(&column, &target_fragments) + .await?; + + // Fragments that need flat evaluation: unindexed + stale (deduplicated). + let flat_fragments: Vec = { + let mut seen = RoaringBitmap::new(); + let mut frags = Vec::new(); + for f in unindexed_fragments.iter().chain( + target_fragments + .iter() + .filter(|f| stale_flat_frag_ids.contains(f.id as u32)), + ) { + if seen.insert(f.id as u32) { + frags.push(f.clone()); + } + } + frags + }; + + // If all target fragments need flat evaluation, skip the indexed path. + if flat_fragments.len() == target_fragments.len() { if self.fast_search { return Ok(Arc::new(EmptyExec::new(FTS_SCHEMA.clone()))); } let flat_match_plan = self - .plan_flat_match_query(unindexed_fragments, query, params, filter_plan) + .plan_flat_match_query(flat_fragments, query, params, filter_plan) .await?; return Ok(flat_match_plan); } - // Mixed case: use index + flat search for unindexed - let match_plan: Arc = Arc::new(MatchQueryExec::new( - self.dataset.clone(), - query.clone(), - params.clone(), - prefilter_source.clone(), - )); + // Build the indexed path. When overlays made some segments stale we use + // `new_with_segments` to restrict the search to fresh segments only. + let match_plan: Arc = match fresh_segments { + Some(segs) => Arc::new(MatchQueryExec::new_with_segments( + self.dataset.clone(), + query.clone(), + params.clone(), + prefilter_source.clone(), + segs, + )), + None => Arc::new(MatchQueryExec::new( + self.dataset.clone(), + query.clone(), + params.clone(), + prefilter_source.clone(), + )), + }; - if self.fast_search || unindexed_fragments.is_empty() { + if self.fast_search || flat_fragments.is_empty() { (Some(match_plan), None) } else { let flat_match_plan = self - .plan_flat_match_query(unindexed_fragments, query, params, filter_plan) + .plan_flat_match_query(flat_fragments, query, params, filter_plan) .await?; (Some(match_plan), Some(flat_match_plan)) } @@ -3840,9 +3954,24 @@ impl Scanner { "Refine factor cannot be zero".to_string(), )); } + // Mask data overlay files: compute which row addresses within each segment have + // been updated by a newer overlay so their ANN entries may be stale. + // These stale rows are blocked from ANN results via the prefilter and re-scored + // on the targeted flat path below — only the specific stale rows, not the whole + // fragment, so sparse overlays incur near-zero overhead. + let stale_rows = self.overlay_stale_vector_rows(&index_segments)?; + // Build a prefilter block mask for stale rows (empty = no-op fast path). + let overlay_block = self.stale_rows_block_mask(&stale_rows).await?; + let ann_node = match vector_type { - DataType::FixedSizeList(_, _) => self.ann(&q, &index_segments, filter_plan).await?, - DataType::List(_) => self.multivec_ann(&q, &index_segments, filter_plan).await?, + DataType::FixedSizeList(_, _) => { + self.ann(&q, &index_segments, filter_plan, overlay_block.clone()) + .await? + } + DataType::List(_) => { + self.multivec_ann(&q, &index_segments, filter_plan, overlay_block.clone()) + .await? + } _ => unreachable!(), }; @@ -3860,7 +3989,14 @@ impl Scanner { if !self.fast_search { knn_node = self - .knn_combined(&q, &index_name, &index_segments, knn_node, filter_plan) + .knn_combined( + &q, + &index_name, + &index_segments, + &stale_rows, + knn_node, + filter_plan, + ) .await?; } @@ -3995,6 +4131,7 @@ impl Scanner { q: &Query, index_name: &str, indexed_segments: &[IndexMetadata], + stale_rows: &HashMap, mut knn_node: Arc, filter_plan: &ExprFilterPlan, ) -> Result> { @@ -4011,30 +4148,41 @@ impl Scanner { self.dataset.unindexed_fragments(index_name).await? }; - if !fallback_fragments.is_empty() { - let q = q.clone(); - debug_assert!(q.metric_type.is_some()); + let has_fallback = !fallback_fragments.is_empty(); + let has_stale = !stale_rows.is_empty(); - // If the vector column is not present, we need to take the vector column, so - // that the distance value is comparable with the flat search ones. - if knn_node.schema().column_with_name(&q.column).is_none() { - let vector_projection = self - .dataset - .empty_projection() - .union_column(&q.column, OnMissing::Error) - .unwrap(); - knn_node = self.take(knn_node, vector_projection)?; - } + if !has_fallback && !has_stale { + return Ok(knn_node); + } - let mut columns = vec![q.column.clone()]; - if let Some(expr) = filter_plan.full_expr.as_ref() { - let filter_columns = Planner::column_names_in_expr(expr); - columns.extend(filter_columns); - } - let vector_scan_projection = Arc::new(self.dataset.schema().project(&columns).unwrap()); - // Note: we could try and use the scalar indices here to reduce the scope of this scan but the - // most common case is that fragments that are newer than the vector index are going to be newer - // than the scalar indices anyways + let q = q.clone(); + debug_assert!(q.metric_type.is_some()); + + // Ensure the vector column is present for distance computation. + if knn_node.schema().column_with_name(&q.column).is_none() { + let vector_projection = self + .dataset + .empty_projection() + .union_column(&q.column, OnMissing::Error)?; + knn_node = self.take(knn_node, vector_projection)?; + } + + let mut columns = vec![q.column.clone()]; + if let Some(expr) = filter_plan.full_expr.as_ref() { + let filter_columns = Planner::column_names_in_expr(expr); + columns.extend(filter_columns); + } + + // Collect flat-path plans; union order matches original (flat before ANN) so test snapshots + // and downstream plan analyses remain stable. + let mut flat_inputs: Vec> = Vec::new(); + + // Flat KNN for unindexed (new-data) fragments. + if has_fallback { + let vector_scan_projection = Arc::new(self.dataset.schema().project(&columns)?); + // Note: we could try and use the scalar indices here to reduce the scope of this scan + // but the most common case is that fragments newer than the vector index are also + // newer than the scalar indices. let mut scan_node = self.scan_fragments( true, false, @@ -4045,39 +4193,51 @@ impl Scanner { Arc::new(fallback_fragments), // Can't pushdown limit/offset in an ANN search None, - // We are re-ordering anyways, so no need to get data in data - // in a deterministic order. + // We are re-ordering anyways, so no need to get data in a deterministic order. false, ); - if let Some(expr) = filter_plan.full_expr.as_ref() { - // If there is a prefilter we need to manually apply it to the new data scan_node = Arc::new(LanceFilterExec::try_new(expr.clone(), scan_node)?); } - // first we do flat search on just the new data - let topk_appended = self.flat_knn(scan_node, &q)?; - - // To do a union, we need to make the schemas match. Right now - // knn_node: _distance, _rowid, vector - // topk_appended: vector, , _rowid, _distance - let topk_appended = project(topk_appended, knn_node.schema().as_ref())?; - assert!( - topk_appended - .schema() - .equivalent_names_and_types(&knn_node.schema()) - ); - // union - let unioned = UnionExec::try_new(vec![Arc::new(topk_appended), knn_node])?; - // Enforce only 1 partition. - let unioned = RepartitionExec::try_new( - unioned, - datafusion::physical_plan::Partitioning::RoundRobinBatch(1), - )?; - // then we do a flat search on KNN(new data) + ANN(indexed data) - return self.flat_knn(Arc::new(unioned), &q); + let topk_fallback = self.flat_knn(scan_node, &q)?; + let topk_fallback: Arc = + Arc::new(project(topk_fallback, knn_node.schema().as_ref())?); + flat_inputs.push(topk_fallback); + } + + // Flat KNN for stale rows only (row-level precision). + // Only specific row addresses need re-scoring, not the whole fragment, so sparse overlays + // incur near-zero overhead. + if has_stale { + // Fetch vector + filter columns for the stale rows. `flat_knn` sorts by row id, so the + // take must carry it (the fallback scan above gets it via `scan_fragments`). + let mut take_proj = self + .dataset + .empty_projection() + .with_row_id() + .union_column(&q.column, OnMissing::Error)?; + if let Some(expr) = filter_plan.full_expr.as_ref() { + let filter_columns = Planner::column_names_in_expr(expr); + take_proj = take_proj.union_columns(filter_columns, OnMissing::Error)?; + } + let mut stale_node = self.stale_rows_take(stale_rows, take_proj).await?; + if let Some(expr) = filter_plan.full_expr.as_ref() { + stale_node = Arc::new(LanceFilterExec::try_new(expr.clone(), stale_node)?); + } + let topk_stale = self.flat_knn(stale_node, &q)?; + let topk_stale: Arc = + Arc::new(project(topk_stale, knn_node.schema().as_ref())?); + flat_inputs.push(topk_stale); } - Ok(knn_node) + // Union: flat paths first (matching original order), then ANN results. + flat_inputs.push(knn_node); + let unioned = UnionExec::try_new(flat_inputs)?; + let unioned = RepartitionExec::try_new( + unioned, + datafusion::physical_plan::Partitioning::RoundRobinBatch(1), + )?; + self.flat_knn(Arc::new(unioned), &q) } #[async_recursion] @@ -4108,29 +4268,237 @@ impl Scanner { } } - /// Given an index query, split the fragments into two sets + /// Given an index query, split the fragments into two groups and collect per-row stale data. /// - /// The first set is the relevant fragments, which are covered by ALL indices in the query - /// The second set is the missing fragments, which are missed by at least one index + /// - `relevant_frags`: covered by ALL indices. Stale rows within them are returned separately + /// so callers can block them from `MaterializeIndexExec` and re-score via a targeted take. + /// - `missing_frags`: not covered by at least one index; fall back to full scan + filter. + /// - `stale_rows`: per-fragment row offsets whose indexed values are stale due to a data + /// overlay committed after the index was built (field-aware, version-gated). Empty when no + /// overlays are present. /// - /// There is no point in handling the case where a fragment is covered by some (but not all) - /// of the indices. If we have to do a full scan of the fragment then we do it + /// There is no point in partially indexing a fragment (some indices cover it, others do not). + /// If we have to do a full scan of a fragment for any reason, we do it entirely. async fn partition_frags_by_coverage( &self, index_expr: &ScalarIndexExpr, fragments: Arc>, - ) -> Result<(Vec, Vec)> { + ) -> Result<(Vec, Vec, HashMap)> { let covered_frags = self.fragments_covered_by_index_query(index_expr).await?; + let stale_rows = self + .overlay_stale_index_rows(index_expr, &fragments) + .await?; let mut relevant_frags = Vec::with_capacity(fragments.len()); let mut missing_frags = Vec::with_capacity(fragments.len()); for fragment in fragments.iter() { if covered_frags.contains(fragment.id as u32) { + // Indexed fragments stay on the indexed path. Stale rows within them are blocked + // from the index result and re-evaluated separately via a targeted take. relevant_frags.push(fragment.clone()); } else { missing_frags.push(fragment.clone()); } } - Ok((relevant_frags, missing_frags)) + Ok((relevant_frags, missing_frags, stale_rows)) + } + + /// Per-row stale offsets for each fragment whose indexed values may be stale because an + /// overlay committed *after* an index was built touches a field that index covers. + /// + /// The check is field-aware (an overlay touching only unindexed fields excludes nothing) and + /// version-gated (an overlay with `committed_version <= index.dataset_version` is already + /// incorporated by the index), via [`overlay_exclusion_offsets`]. + async fn overlay_stale_index_rows( + &self, + index_expr: &ScalarIndexExpr, + fragments: &[Fragment], + ) -> Result> { + // Overlays are rare; skip all index loading when none of the candidate fragments has one. + let overlaid_frags = overlaid_fragments(fragments); + if overlaid_frags.is_empty() { + return Ok(HashMap::new()); + } + + // Walk the (boolean) index expression tree to collect leaf searches. + let mut searches = Vec::new(); + let mut stack = vec![index_expr]; + while let Some(expr) = stack.pop() { + match expr { + ScalarIndexExpr::Not(inner) => stack.push(inner), + ScalarIndexExpr::And(lhs, rhs) | ScalarIndexExpr::Or(lhs, rhs) => { + stack.push(lhs); + stack.push(rhs); + } + ScalarIndexExpr::Query(search) => searches.push(search), + } + } + + // `load_named_scalar_segments` returns cached index metadata — no disk I/O on the hot + // path. Even without the cache, this code is only reached when at least one fragment has + // overlays (rare), so the per-leaf cost is acceptable. + let mut stale: HashMap = HashMap::new(); + for search in searches { + let segments = load_named_scalar_segments( + self.dataset.as_ref(), + &search.column, + &search.index_name, + ) + .await?; + for segment in &segments { + collect_overlay_stale_rows_for_segment(segment, &overlaid_frags, &mut stale)?; + } + } + Ok(stale) + } + + /// Compute per-row stale data for a vector index's segments. + /// + /// Returns a map from fragment_id to the set of row offsets within that fragment that are stale + /// (their vector values have been updated by a newer overlay since the index was built). An + /// empty map means no stale rows — the fast path where no masking is needed. + fn overlay_stale_vector_rows( + &self, + segments: &[IndexMetadata], + ) -> Result> { + // Scope to the query's target fragments (all dataset fragments if unscoped). + let dataset_frags = self.dataset.fragments(); + let fragments: &[Fragment] = match self.fragments.as_ref() { + Some(f) => f.as_slice(), + None => dataset_frags.as_slice(), + }; + let overlaid_frags = overlaid_fragments(fragments); + if overlaid_frags.is_empty() { + return Ok(HashMap::new()); + } + let mut stale: HashMap = HashMap::new(); + for segment in segments { + collect_overlay_stale_rows_for_segment(segment, &overlaid_frags, &mut stale)?; + } + Ok(stale) + } + + /// Compute which FTS segments are stale due to data overlay files committed after the + /// index was built, and which fragments must therefore fall back to the flat text path. + /// + /// Returns `(flat_frag_ids, Some(fresh_segments))` when overlays are present: + /// - `flat_frag_ids`: fragment IDs that must be scanned flat (stale fragments, plus any + /// other fragments co-located in a segment that covers a stale one — the whole segment is + /// excluded, so all fragments it covered must move to flat). + /// - `fresh_segments`: the subset of FTS segments that cover no stale fragment; safe to + /// pass to `MatchQueryExec::new_with_segments`. + /// + /// Returns `(empty, None)` on the fast path (no overlays, or no segments load). + async fn fts_stale_frags_and_fresh_segments( + &self, + column: &str, + target_fragments: &[Fragment], + ) -> Result<(RoaringBitmap, Option>)> { + // Fast path: no overlays on any target fragment. + if target_fragments.iter().all(|f| f.overlays.is_empty()) { + return Ok((RoaringBitmap::new(), None)); + } + + let Some(segments) = load_segments(&self.dataset, column).await? else { + return Ok((RoaringBitmap::new(), None)); + }; + + let overlaid_frags = overlaid_fragments(target_fragments); + let mut stale_frag_ids = RoaringBitmap::new(); + for seg in &segments { + collect_overlay_stale_frags(seg, &overlaid_frags, &mut stale_frag_ids)?; + } + + if stale_frag_ids.is_empty() { + // Overlays exist but none are on this FTS column or predate the index. + return Ok((stale_frag_ids, None)); + } + + // Any segment covering a stale fragment is excluded from the indexed path. + // All fragments covered by that segment (stale + co-located fresh ones) must + // fall to the flat path, since the indexed path no longer covers them. + let mut flat_frag_ids = stale_frag_ids.clone(); + let mut fresh_segments = Vec::with_capacity(segments.len()); + for seg in segments { + match &seg.fragment_bitmap { + Some(bm) if !bm.is_disjoint(&stale_frag_ids) => { + flat_frag_ids |= bm; + // exclude this segment from the indexed path + } + Some(_) => fresh_segments.push(seg), + None => { + // Coverage unknown (legacy segment without a fragment bitmap): we can neither + // trust it to exclude overlay-stale rows nor tell which fragments it indexes. + // Exclude it from the indexed path and route every target fragment to flat. + flat_frag_ids.extend(target_fragments.iter().map(|f| f.id as u32)); + } + } + } + + Ok((flat_frag_ids, Some(fresh_segments))) + } + + /// Collect the stale rows into a [`RowAddrTreeMap`] in the domain the index results use. + /// + /// Index results are in the row-id domain (see `ScalarQuery::evaluate_nullable`), and a + /// physical row address equals its row id only when the dataset does not use stable row ids. + /// Under stable row ids the addresses are translated to their row ids so the result lines up + /// with the index output it is combined with (block mask) or taken against. + async fn stale_rows_in_id_domain( + &self, + stale_rows: &HashMap, + ) -> Result { + let mut tree_map = RowAddrTreeMap::new(); + for (&frag_id, offsets) in stale_rows { + tree_map.insert_bitmap(frag_id, offsets.clone()); + } + if self.dataset.manifest.uses_stable_row_ids() { + tree_map = translate_addr_treemap_to_row_ids(&self.dataset, &tree_map).await?; + } + Ok(tree_map) + } + + /// Build a block-list mask over stale rows, or `None` when there are none. + /// + /// The mask removes these rows from an index result so the index never emits them; they + /// are re-evaluated against their current (overlay-merged) values on a targeted take path + /// (see [`Self::stale_rows_take`]). + async fn stale_rows_block_mask( + &self, + stale_rows: &HashMap, + ) -> Result> { + if stale_rows.is_empty() { + return Ok(None); + } + let tree_map = self.stale_rows_in_id_domain(stale_rows).await?; + Ok(Some(RowAddrMask::from_block(tree_map))) + } + + /// Take the stale rows by physical address, projecting `projection`, to re-evaluate only + /// those rows (rather than their whole fragments) against their current overlay-merged values. + /// + /// The rows are identified by an address allow list routed through `FilteredReadExec`, not a + /// `_rowid` column (`_rowid` and row address diverge under stable row ids — see + /// [`Self::stale_rows_in_id_domain`]). + async fn stale_rows_take( + &self, + stale_rows: &HashMap, + projection: Projection, + ) -> Result> { + let take_id_map = self.stale_rows_in_id_domain(stale_rows).await?; + let take_ids: Vec = take_id_map + .row_addrs() + .map(|it| it.map(u64::from).collect()) + .unwrap_or_default(); + let index_input = self.u64s_as_take_input(take_ids)?; + let mut read_options = FilteredReadOptions::new(projection); + if let Some(fragments) = self.fragments.as_ref() { + read_options = read_options.with_fragments(Arc::new(fragments.clone())); + } + Ok(Arc::new(FilteredReadExec::try_new( + self.dataset.clone(), + read_options, + Some(index_input), + )?)) } // First perform a lookup in a scalar index for ids and then perform a take on the @@ -4152,16 +4520,24 @@ impl Scanner { let needs_recheck = index_expr.needs_recheck(); - // Figure out which fragments are covered by ALL indices - let (relevant_frags, missing_frags) = self + // Figure out which fragments are covered by ALL indices, and which rows within + // covered fragments are stale due to data overlay files. + let (relevant_frags, missing_frags, stale_rows) = self .partition_frags_by_coverage(index_expr, fragments) .await?; - let mut plan: Arc = Arc::new(MaterializeIndexExec::new( + // Build the MaterializeIndexExec, blocking stale row addresses so the index never + // emits them. Stale rows are re-scored separately via a targeted take below. + let mat_exec = MaterializeIndexExec::new( self.dataset.clone(), index_expr.clone(), Arc::new(relevant_frags), - )); + ); + let mat_exec = match self.stale_rows_block_mask(&stale_rows).await? { + Some(block) => mat_exec.with_overlay_block(block), + None => mat_exec, + }; + let mut plan: Arc = Arc::new(mat_exec); let refine_expr = filter_plan.refine_expr.as_ref(); @@ -4207,6 +4583,21 @@ impl Scanner { plan = Arc::new(AddRowAddrExec::try_new(plan, self.dataset.clone(), 0)?); } + // Both the missing-fragments path (full scan) and the stale-rows path (targeted take) + // need the user's projection extended with any filter columns. Compute it once. + let fallback_projection: Option = + if !missing_frags.is_empty() || !stale_rows.is_empty() { + let filter = filter_plan.full_expr.as_ref().expect_ok()?; + let filter_cols = Planner::column_names_in_expr(filter); + Some( + projection + .clone() + .union_columns(filter_cols, OnMissing::Error)?, + ) + } else { + None + }; + let new_data_path: Option> = if !missing_frags.is_empty() { log::trace!( "scalar_indexed_scan will need full scan of {} missing fragments", @@ -4225,10 +4616,8 @@ impl Scanner { // If there were no extra columns then we still need the project // because Materialize -> Take puts the row id at the left and // Scan puts the row id at the right - let filter = filter_plan.full_expr.as_ref().unwrap(); - let filter_cols = Planner::column_names_in_expr(filter); - let scan_projection = projection.union_columns(filter_cols, OnMissing::Error)?; - + let scan_projection = fallback_projection.clone().expect_ok()?; + let filter = filter_plan.full_expr.as_ref().expect_ok()?; let scan_schema = Arc::new(scan_projection.to_bare_schema()); let scan_arrow_schema = Arc::new(scan_schema.as_ref().into()); let planner = Planner::new(scan_arrow_schema); @@ -4257,16 +4646,37 @@ impl Scanner { None }; - if let Some(new_data_path) = new_data_path { - let unioned = UnionExec::try_new(vec![plan, new_data_path])?; - // Enforce only 1 partition. - let unioned = Arc::new(RepartitionExec::try_new( - unioned, - datafusion::physical_plan::Partitioning::RoundRobinBatch(1), - )?); - Ok(unioned) + // Stale-Take path: re-evaluate only the stale row addresses against the full filter + // (row-level optimization). These rows were blocked from the index result above; + // here we take their current (overlay-merged) values and re-apply the predicate. + // The schema matches `plan` via `project(…, plan.schema())`. + let stale_take_path: Option> = if stale_rows.is_empty() { + None } else { + let filter = filter_plan.full_expr.as_ref().expect_ok()?; + let take_projection = fallback_projection.expect_ok()?; + + let stale_node = self.stale_rows_take(&stale_rows, take_projection).await?; + + let planner = Planner::new(stale_node.schema()); + let optimized_filter = planner.optimize_expr(filter.clone())?; + let filtered = Arc::new(LanceFilterExec::try_new(optimized_filter, stale_node)?); + Some(Arc::new(project(filtered, plan.schema().as_ref())?)) + }; + + let extra_paths: Vec> = [new_data_path, stale_take_path] + .into_iter() + .flatten() + .collect(); + if extra_paths.is_empty() { Ok(plan) + } else { + let all_paths = std::iter::once(plan).chain(extra_paths).collect(); + let unioned = UnionExec::try_new(all_paths)?; + Ok(Arc::new(RepartitionExec::try_new( + unioned, + datafusion::physical_plan::Partitioning::RoundRobinBatch(1), + )?)) } } @@ -4700,11 +5110,18 @@ impl Scanner { q: &Query, index: &[IndexMetadata], filter_plan: &ExprFilterPlan, + overlay_block: Option, ) -> Result> { let prefilter_source = self .prefilter_source(filter_plan, self.get_indexed_frags(index)) .await?; - let inner_fanout_search = new_knn_exec(self.dataset.clone(), index, q, prefilter_source)?; + let inner_fanout_search = new_knn_exec( + self.dataset.clone(), + index, + q, + prefilter_source, + overlay_block, + )?; let sort_expr = PhysicalSortExpr { expr: expressions::col(DIST_COL, inner_fanout_search.schema().as_ref())?, options: SortOptions { @@ -4731,6 +5148,7 @@ impl Scanner { q: &Query, index: &[IndexMetadata], filter_plan: &ExprFilterPlan, + overlay_block: Option, ) -> Result> { // we split the query procedure into two steps: // 1. collect the candidates by vector searching on each query vector @@ -4763,6 +5181,7 @@ impl Scanner { index, &query, prefilter_source.clone(), + overlay_block.clone(), )?; let sort_expr = PhysicalSortExpr { expr: expressions::col(DIST_COL, ann_node.schema().as_ref())?, @@ -4843,11 +5262,14 @@ impl Scanner { // are not in the fragments we are scanning. if filter_plan.is_exact_index_search() && self.fragments.is_none() { let index_query = filter_plan.index_query.as_ref().expect_ok()?; - let (_, missing_frags) = self + let (_, missing_frags, stale_rows) = self .partition_frags_by_coverage(index_query, fragments.clone()) .await?; - if missing_frags.is_empty() || self.fast_search { + // Overlay-stale rows must never reach the direct ScalarIndexExec path: it would hand + // ANN/FTS a selection vector containing rows whose indexed values are now stale. When + // any exist, fall through to the filtered-read prefilter, which masks them. + if stale_rows.is_empty() && (missing_frags.is_empty() || self.fast_search) { log::trace!("prefilter entirely satisfied by exact index search"); let result_format = self.index_expr_result_format(); // We can only avoid materializing the index for a prefilter if: @@ -4971,7 +5393,12 @@ impl Scanner { #[instrument(level = "info", skip(self))] pub async fn explain_plan(&self, verbose: bool) -> Result { - let plan = self.create_plan().await?; + // Box the plan-building future at the call site: `create_plan`'s inlined async + // layout otherwise exceeds rustc's depth limit here. It has to be boxed at the + // call site rather than inside `create_plan` — boxing internally turns the + // future's `Send` check into a `Box: Send` trait obligation that + // overflows the solver through the cache types (E0275 in downstream crates). + let plan = Box::pin(self.create_plan()).await?; let display = DisplayableExecutionPlan::new(plan.as_ref()); Ok(format!("{}", display.indent(verbose))) @@ -10368,10 +10795,7 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") // 8KB stays under the 64KB inline threshold, so the blob is a normal column in // the data file rather than a dedicated blob file. - let blob_meta = std::collections::HashMap::from([( - "lance-encoding:blob".to_string(), - "true".to_string(), - )]); + let blob_meta = HashMap::from([("lance-encoding:blob".to_string(), "true".to_string())]); let blobs = array::rand_fixedbin(ByteCount::from(8 * 1024), true).with_metadata(blob_meta); let data = gen_batch() .col("filterme", array::step::()) @@ -10430,10 +10854,7 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") use lance_io::assert_io_lt; use lance_table::io::commit::RenameCommitHandler; - let blob_meta = std::collections::HashMap::from([( - "lance-encoding:blob".to_string(), - "true".to_string(), - )]); + let blob_meta = HashMap::from([("lance-encoding:blob".to_string(), "true".to_string())]); let a_field = ArrowField::new("a", DataType::Int32, false); let blob_field = ArrowField::new("blob", DataType::LargeBinary, false).with_metadata(blob_meta); diff --git a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs new file mode 100644 index 00000000000..cfabdc13266 --- /dev/null +++ b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs @@ -0,0 +1,1057 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! End-to-end tests for data-overlay index masking: a scalar index masks data overlay files so that +//! queries stay correct while overlays remain (stale index hits are dropped and new +//! matches are added by re-evaluating overlay-covered rows on the flat path). + +use std::sync::Arc; + +use futures::TryStreamExt; + +use arrow_array::cast::AsArray; +use arrow_array::types::Int32Type; +use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchIterator, StringArray}; +use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; +use lance_index::IndexType; +use lance_index::scalar::FullTextSearchQuery; +use lance_index::scalar::ScalarIndexParams; +use lance_index::scalar::inverted::InvertedIndexParams; +use lance_io::utils::CachedFileSize; +use lance_linalg::distance::MetricType; +use lance_table::format::DataFile; +use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; +use roaring::RoaringBitmap; +use rstest::rstest; + +use lance_file::writer::{FileWriter, FileWriterOptions}; + +use crate::Dataset; +use crate::dataset::transaction::{DataOverlayGroup, Operation}; +use crate::dataset::{WriteDestination, WriteParams}; +use crate::index::DatasetIndexExt; +use crate::index::vector::VectorIndexParams; + +/// Two-fragment Int32 dataset: `id` (field 0) = 0..12 and `age` (field 1) = id * 10, +/// six rows per file (fragments 0 and 1). In-memory store so overlay files can be written +/// with a store-relative `data/.lance` path and committed against the dataset. +async fn create_base_dataset() -> Dataset { + create_base_dataset_with(false).await +} + +async fn create_base_dataset_with(stable_row_ids: bool) -> Dataset { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("age", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..12)), + Arc::new(Int32Array::from_iter_values((0..12).map(|v| v * 10))), + ], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: 6, + enable_stable_row_ids: stable_row_ids, + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap() +} + +async fn build_age_index(dataset: &mut Dataset) { + dataset + .create_index( + &["age"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); +} + +/// Write an overlay file covering `fields` of `fragment_id` with `coverage` and the given +/// per-field value columns, then commit it as a `DataOverlay` transaction. `name` makes +/// the overlay file unique. +async fn commit_overlay( + dataset: Dataset, + name: &str, + fragment_id: u64, + fields: &[i32], + coverage: OverlayCoverage, + columns: Vec, +) -> Dataset { + let read_version = dataset.version().version; + let overlay_schema = dataset.schema().project_by_ids(fields, true); + + let filename = format!("{name}.lance"); + // Use dataset.base so the path is absolute for file:// stores. + // to_local_path() prepends '/' to the object_store path, so a bare + // "data/foo.lance" would resolve to /data/foo.lance (root fs). With + // base we get e.g. tmp/lance-bench/data/foo.lance → /tmp/lance-bench/data/foo.lance. + // For memory:// stores base is empty so the result is the same as before. + let path = dataset.base.clone().join("data").join(filename.as_str()); + let obj_writer = dataset.object_store.create(&path).await.unwrap(); + let mut writer = + FileWriter::try_new(obj_writer, overlay_schema, FileWriterOptions::default()).unwrap(); + let (major, minor) = writer.version().to_numbers(); + for (i, array) in columns.into_iter().enumerate() { + writer.write_column(i, array).await.unwrap(); + } + let summary = writer.finish().await.unwrap(); + + let mut data_file = DataFile::new_unstarted(filename, major, minor); + data_file.fields = writer + .field_id_to_column_indices() + .iter() + .map(|(field_id, _)| *field_id as i32) + .collect::>() + .into(); + data_file.column_indices = writer + .field_id_to_column_indices() + .iter() + .map(|(_, column_index)| *column_index as i32) + .collect::>() + .into(); + data_file.file_size_bytes = CachedFileSize::new(summary.size_bytes); + + let overlay = DataOverlayFile { + data_file, + coverage, + committed_version: 0, + }; + Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id, + overlays: vec![overlay], + }], + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap() +} + +/// Sorted `id` values returned by a filtered scan. +async fn ids_matching(dataset: &Dataset, filter: &str) -> Vec { + ids_matching_opts(dataset, filter, false).await +} + +/// Like [`ids_matching`] but lets a test enable `fast_search()`, which skips unindexed +/// fragments. Overlay masking on indexed fragments must still apply regardless. +async fn ids_matching_opts(dataset: &Dataset, filter: &str, fast_search: bool) -> Vec { + let mut scanner = dataset.scan(); + scanner.filter(filter).unwrap().project(&["id"]).unwrap(); + if fast_search { + scanner.fast_search(); + } + let batch = scanner.try_into_batch().await.unwrap(); + let mut ids = ids_from_batches(std::slice::from_ref(&batch)); + ids.sort_unstable(); + ids +} + +/// Concatenate the `id` (Int32) column from each batch, in batch order. +fn ids_from_batches(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|b| { + b.column_by_name("id") + .unwrap() + .as_primitive::() + .values() + .to_vec() + }) + .collect() +} + +fn i32_array(values: impl IntoIterator>) -> ArrayRef { + Arc::new(Int32Array::from_iter(values)) +} + +fn fsl(rows: Vec>, dim: i32) -> ArrayRef { + let flat: Vec = rows.into_iter().flatten().collect(); + let item = Arc::new(ArrowField::new("item", DataType::Float32, true)); + Arc::new( + arrow_array::FixedSizeListArray::try_new( + item, + dim, + Arc::new(arrow_array::Float32Array::from(flat)), + None, + ) + .unwrap(), + ) +} + +/// A newer overlay on the indexed field drops stale index hits (the old value no longer +/// matches) and surfaces new matches (the new value is found even though the index never +/// saw it). Mirrors the spec's Bob 25 -> 26 worked example. +/// +/// Parametrized over `stable_row_ids` to cover the address-based stale-Take path under both +/// row-id schemes. +#[rstest] +#[tokio::test] +async fn test_overlay_stale_drop_and_new_match(#[values(false, true)] stable_row_ids: bool) { + let mut dataset = create_base_dataset_with(stable_row_ids).await; + build_age_index(&mut dataset).await; + + // Fragment 0, offset 1 is id=1, age=10. The overlay (committed after the index) + // changes its age to 999. + let dataset = commit_overlay( + dataset, + "age_overlay", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![i32_array([Some(999)])], + ) + .await; + + // Stale-drop: the index still holds age=10 for id=1, but its current value is 999, + // so it must not be returned. + assert_eq!(ids_matching(&dataset, "age = 10").await, Vec::::new()); + // New-match: the index never saw age=999, but re-evaluation finds it. + assert_eq!(ids_matching(&dataset, "age = 999").await, vec![1]); + // An untouched indexed value is unaffected. + assert_eq!(ids_matching(&dataset, "age = 20").await, vec![2]); +} + +/// Row-level BTree precision: when one row in a covered fragment is stale, only that row is +/// blocked from the index result and re-evaluated on the stale-Take path. Non-stale rows in +/// the same fragment (including one that matches the predicate) remain on the indexed path. +/// +/// Setup: fragment 0 has id=5 → age=50 (not stale). Overlay id=1 → age=50 (stale). +/// After the overlay two rows in fragment 0 have age=50. The row-level optimization must +/// return both: id=5 from the index and id=1 from the stale-Take path. +/// +/// Parametrized over `stable_row_ids`: with stable row ids enabled the stale-Take path must +/// identify rows by physical address, not `_rowid`, or it would take the wrong rows. +#[rstest] +#[tokio::test] +async fn test_btree_overlay_row_level_precision(#[values(false, true)] stable_row_ids: bool) { + let mut dataset = create_base_dataset_with(stable_row_ids).await; + build_age_index(&mut dataset).await; + + // Fragment 0: ids 0-5, ages 0,10,20,30,40,50. Overlay offset 1 (id=1): age 10→50. + // After this both id=1 and id=5 have age=50, in the same fragment. + let dataset = commit_overlay( + dataset, + "age_row_level", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![i32_array([Some(50)])], + ) + .await; + + // Stale drop: id=1's old age=10 entry must not appear. + assert_eq!(ids_matching(&dataset, "age = 10").await, Vec::::new()); + + // id=5 via index + id=1 via stale-Take path — both in fragment 0. + assert_eq!(ids_matching(&dataset, "age = 50").await, vec![1, 5]); + + // Non-stale rows in the same fragment still return correctly. + assert_eq!(ids_matching(&dataset, "age = 20").await, vec![2]); + assert_eq!(ids_matching(&dataset, "age = 30").await, vec![3]); +} + +/// `fast_search` skips *unindexed fragments*, but overlay masking on indexed fragments must +/// still apply: the drop-stale block and the stale-Take re-eval both run regardless of +/// `fast_search` on the scalar path. A regression that gated overlay masking behind +/// `!fast_search` would leak id=1's stale age=10 hit here. +#[tokio::test] +async fn test_btree_overlay_masked_under_fast_search() { + let mut dataset = create_base_dataset().await; + build_age_index(&mut dataset).await; + + // Fragment 0, offset 1 is id=1, age=10. Overlay (committed after the index) → age=999. + let dataset = commit_overlay( + dataset, + "age_fast_search", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![i32_array([Some(999)])], + ) + .await; + + // Stale hit dropped even under fast_search — the block is not gated by fast_search. + assert_eq!( + ids_matching_opts(&dataset, "age = 10", true).await, + Vec::::new() + ); + // The scalar re-eval path is likewise not gated, so the new value is still surfaced. + assert_eq!( + ids_matching_opts(&dataset, "age = 999", true).await, + vec![1] + ); + // An untouched indexed value on the same fragment is unaffected. + assert_eq!(ids_matching_opts(&dataset, "age = 20", true).await, vec![2]); +} + +/// An overlay touching only a non-indexed field excludes nothing from the index on `age`. +#[tokio::test] +async fn test_overlay_on_unrelated_field_excludes_nothing() { + let mut dataset = create_base_dataset().await; + build_age_index(&mut dataset).await; + + // Overlay field 0 (`id`), not the indexed `age`. The age index stays fully trusted. + let dataset = commit_overlay( + dataset, + "id_overlay", + 0, + &[0], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![i32_array([Some(777)])], + ) + .await; + + // The age index is still trusted: age=10 finds the offset-1 row, whose id now reads + // through the overlay as 777. The fragment was not routed to the flat path on account + // of an overlay that touches no indexed field. + assert_eq!(ids_matching(&dataset, "age = 10").await, vec![777]); + // An untouched row is unaffected. + assert_eq!(ids_matching(&dataset, "age = 20").await, vec![2]); + // The overlaid id is the new value on read, and the old one is gone. + assert_eq!(ids_matching(&dataset, "id = 777").await, vec![777]); + assert_eq!(ids_matching(&dataset, "id = 1").await, Vec::::new()); +} + +/// An overlay whose `committed_version <= index.dataset_version` is already incorporated by +/// the index (the index was built reading merged values) and is not excluded. +#[tokio::test] +async fn test_overlay_older_than_index_not_excluded() { + let dataset = create_base_dataset().await; + + // Commit the overlay first (age of id=1 becomes 999), then build the index on top. + let mut dataset = commit_overlay( + dataset, + "age_overlay_old", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![i32_array([Some(999)])], + ) + .await; + build_age_index(&mut dataset).await; + + // The index incorporates the overlay, so it returns the merged value directly. + assert_eq!(ids_matching(&dataset, "age = 999").await, vec![1]); + assert_eq!(ids_matching(&dataset, "age = 10").await, Vec::::new()); +} + +/// A covered offset whose overlay value is NULL overrides the cell to NULL, so the stale +/// index hit for its old value is dropped. +#[tokio::test] +async fn test_overlay_null_override() { + let mut dataset = create_base_dataset().await; + build_age_index(&mut dataset).await; + + // id=1 (age=10) is overridden to NULL. + let dataset = commit_overlay( + dataset, + "age_overlay_null", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![i32_array([None])], + ) + .await; + + assert_eq!(ids_matching(&dataset, "age = 10").await, Vec::::new()); + assert_eq!(ids_matching(&dataset, "age IS NULL").await, vec![1]); +} + +/// Overlays on a non-first fragment are masked correctly, and a query spanning both +/// fragments returns the right rows. +/// +/// Parametrized over `stable_row_ids`, and crucially overlays fragment 1 (ids 6..12), where a +/// physical address diverges from the stable row id — so this exercises the address-vs-row-id +/// distinction that a fragment-0 overlay cannot. +#[rstest] +#[tokio::test] +async fn test_overlay_multi_fragment(#[values(false, true)] stable_row_ids: bool) { + let mut dataset = create_base_dataset_with(stable_row_ids).await; + build_age_index(&mut dataset).await; + + // Fragment 1 holds ids 6..12 (ages 60..110). Offset 2 within fragment 1 is id=8, + // age=80; change it to 60 (a value that also legitimately exists at id=6). + let dataset = commit_overlay( + dataset, + "age_overlay_frag1", + 1, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([2])), + vec![i32_array([Some(60)])], + ) + .await; + + // id=8 no longer has age=80 (stale-drop on fragment 1). + assert_eq!(ids_matching(&dataset, "age = 80").await, Vec::::new()); + // Both id=6 (base) and id=8 (overlay) now have age=60 (new-match added to base hit). + assert_eq!(ids_matching(&dataset, "age = 60").await, vec![6, 8]); + // A value in the untouched fragment 0 is still served correctly. + assert_eq!(ids_matching(&dataset, "age = 30").await, vec![3]); +} + +const VEC_DIM: i32 = 8; + +fn vec_query() -> Vec { + vec![1.0_f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] +} + +/// 64-row two-fragment vector dataset with a single-partition IVF_FLAT index, then an overlay +/// on fragment 1 that moves id=35 (offset 3) onto `far` (away from the query) and id=40 +/// (offset 8) onto the query. Built before the overlay, the index still believes id=35 is the +/// query and has never seen id=40 near it. Every other base vector is orthogonal to the query. +/// +/// Overlaying fragment 1 (ids 32..64) is deliberate: a physical address diverges from the +/// stable row id there, so both the ANN prefilter block and the flat re-score take must operate +/// in the row-id domain when `stable_row_ids` is enabled. +async fn create_vector_overlay_dataset(stable_row_ids: bool) -> Dataset { + let query = vec_query(); + let far = vec![0.0_f32, 100.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]; + + let mut vectors: Vec> = Vec::with_capacity(64); + for i in 0..64 { + if i == 35 { + vectors.push(query.clone()); + } else { + let mut v = vec![0.0_f32; VEC_DIM as usize]; + v[1] = (i + 2) as f32; // orthogonal to the query, distinct, far + vectors.push(v); + } + } + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + VEC_DIM, + ), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..64)), + fsl(vectors, VEC_DIM), + ], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: 32, + enable_stable_row_ids: stable_row_ids, + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let mut dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + + // Single-partition IVF_FLAT: the ANN searches every indexed row with exact distances. + let params = VectorIndexParams::ivf_flat(1, MetricType::L2); + dataset + .create_index(&["vec"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + commit_overlay( + dataset, + "vec_overlay", + 1, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([3, 8])), + vec![fsl(vec![far, query], VEC_DIM)], + ) + .await +} + +/// Run a top-`k` ANN search for the standard query vector and return the returned `id`s, +/// optionally with `fast_search()` enabled. +async fn vector_query_ids(dataset: &Dataset, k: usize, fast_search: bool) -> Vec { + let mut scanner = dataset.scan(); + scanner + .nearest("vec", &arrow_array::Float32Array::from(vec_query()), k) + .unwrap() + .minimum_nprobes(1) + .project(&["id"]) + .unwrap(); + if fast_search { + scanner.fast_search(); + } + let results = scanner + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + ids_from_batches(&results) +} + +/// A vector index masks overlays: a row whose vector was moved (by a newer overlay) away +/// from the query is dropped from results, and a row moved *onto* the query is found by +/// re-scoring its current vector on the flat path — even though the index never saw it. +/// +/// Parametrized over `stable_row_ids` to cover the row-id domain for both block and re-score. +#[rstest] +#[tokio::test] +async fn test_vector_index_rescore_on_overlay(#[values(false, true)] stable_row_ids: bool) { + let dataset = create_vector_overlay_dataset(stable_row_ids).await; + let ids = vector_query_ids(&dataset, 3, false).await; + + // id=40 was moved onto the query and is found by re-scoring (new-match recall). + assert!( + ids.contains(&40), + "expected id=40 (re-scored to query) in {ids:?}" + ); + // id=35's stale index entry (the query) must not resurface: its current vector is far. + assert!( + !ids.contains(&35), + "stale vector for id=35 should be dropped, got {ids:?}" + ); +} + +/// The ANN prefilter block that drops stale overlay rows runs regardless of `fast_search`; +/// only the flat re-score is gated by it. So under `fast_search` id=35's stale hit must still +/// be dropped, while id=40 (moved onto the query) is intentionally not re-scored — the same +/// recall tradeoff `fast_search` already makes for unindexed data. A regression that moved the +/// `overlay_block` computation inside the `!fast_search` guard would leak id=35's stale vector. +#[tokio::test] +async fn test_vector_overlay_stale_dropped_under_fast_search() { + let dataset = create_vector_overlay_dataset(false).await; + let ids = vector_query_ids(&dataset, 3, true).await; + + // Correctness: the stale index hit is dropped even though the re-score is skipped. + assert!( + !ids.contains(&35), + "stale vector for id=35 must be dropped under fast_search, got {ids:?}" + ); + // Recall tradeoff: fast_search skips the flat re-score, so the moved-on match is not surfaced. + assert!( + !ids.contains(&40), + "fast_search skips re-score, so id=40 should be absent, got {ids:?}" + ); +} + +/// A compound boolean predicate (age AND id) exercises the ScalarIndexExpr tree-walk in +/// `overlay_stale_index_rows`. An overlay on `age` marks fragment 0 stale from the `age` +/// index's perspective, so the compound query must re-evaluate fragment 0 on the flat path. +#[tokio::test] +async fn test_overlay_stale_with_compound_index_expression() { + let mut dataset = create_base_dataset().await; + // Build BTree indexes on both columns so a compound filter can use both. + build_age_index(&mut dataset).await; + dataset + .create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + + // Fragment 0 covers id=0..5, age=0..50. Overlay changes id=1's age from 10 to 999. + let dataset = commit_overlay( + dataset, + "age_compound", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![i32_array([Some(999)])], + ) + .await; + + // Compound query: both the `age` and `id` index are involved. The overlay on `age` + // makes fragment 0 stale for the `age` index; it falls to the flat path, which uses + // the merged (overlay) value. Result: the stale age=10 hit is gone, age=999 appears. + assert_eq!(ids_matching(&dataset, "age = 10").await, Vec::::new()); + assert_eq!(ids_matching(&dataset, "age = 999").await, vec![1]); + // A pure `id` query on an unaffected fragment still works correctly. + assert_eq!(ids_matching(&dataset, "id = 2").await, vec![2]); +} + +/// Text dataset: two fragments, 6 rows each. Schema: id (Int32), text (Utf8). +/// Texts are unique tokens so each row can be identified by its term. +async fn create_text_dataset() -> Dataset { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("text", DataType::Utf8, true), + ])); + let texts: Vec<&str> = vec![ + "apple pie", + "apple banana", // row 1, fragment 0 — will be overlaid in tests + "cherry cake", + "banana split", + "orange juice", + "grape vine", + "mango sorbet", // fragment 1 starts here + "pear tart", + "lemon curd", + "peach cobbler", + "plum pudding", + "fig newton", + ]; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..12)), + Arc::new(StringArray::from(texts)), + ], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: 6, + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap() +} + +async fn build_text_fts_index(dataset: &mut Dataset) { + dataset + .create_index( + &["text"], + IndexType::Inverted, + None, + &InvertedIndexParams::default(), + true, + ) + .await + .unwrap(); +} + +/// FTS index with token positions stored, required for phrase queries. +async fn build_text_fts_index_with_positions(dataset: &mut Dataset) { + dataset + .create_index( + &["text"], + IndexType::Inverted, + None, + &InvertedIndexParams::default().with_position(true), + true, + ) + .await + .unwrap(); +} + +/// Collect sorted IDs of rows returned by an FTS query on `text`. +async fn fts_ids(dataset: &Dataset, query: FullTextSearchQuery) -> Vec { + let results = dataset + .scan() + .full_text_search(query) + .unwrap() + .project(&["id"]) + .unwrap() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut ids = ids_from_batches(&results); + ids.sort_unstable(); + ids +} + +async fn fts_ids_matching(dataset: &Dataset, term: &str) -> Vec { + fts_ids(dataset, FullTextSearchQuery::new(term.to_owned())).await +} + +async fn fts_phrase_ids_matching(dataset: &Dataset, phrase: &str) -> Vec { + use lance_index::scalar::inverted::query::{FtsQuery, PhraseQuery}; + + let query = FullTextSearchQuery::new_query(FtsQuery::Phrase( + PhraseQuery::new(phrase.to_owned()).with_column(Some("text".to_owned())), + )); + fts_ids(dataset, query).await +} + +/// An overlay committed after the FTS index is built replaces a row's text. Searching for +/// the old term must not return the stale row; searching for the new term must find it. +#[tokio::test] +async fn test_fts_overlay_stale_drop_and_new_match() { + let mut dataset = create_text_dataset().await; + build_text_fts_index(&mut dataset).await; + + // fragment 0, row offset 1 (id=1): "apple banana" → "cherry mango" + // field ID 1 is the `text` column. + let dataset = commit_overlay( + dataset, + "text_overlay", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![Arc::new(StringArray::from(vec![Some("cherry mango")]))], + ) + .await; + + // "apple" now matches only id=0 ("apple pie"); id=1's stale index entry must be dropped. + assert_eq!(fts_ids_matching(&dataset, "apple").await, vec![0]); + + // "banana" matched id=1 and id=3 before; after overlay id=1's stale entry must be gone. + assert_eq!(fts_ids_matching(&dataset, "banana").await, vec![3]); + + // "cherry" now matches id=1 (via flat path on stale fragment) and id=2 ("cherry cake"). + let cherry_ids = fts_ids_matching(&dataset, "cherry").await; + assert!( + cherry_ids.contains(&1), + "id=1 overlay→cherry mango should be found: {cherry_ids:?}" + ); + assert!( + cherry_ids.contains(&2), + "id=2 cherry cake should still be found: {cherry_ids:?}" + ); + + // "mango" now matches id=1 (overlay) and id=6 ("mango sorbet" in fragment 1). + let mango_ids = fts_ids_matching(&dataset, "mango").await; + assert!( + mango_ids.contains(&1), + "id=1 overlay→cherry mango should be found: {mango_ids:?}" + ); + assert!( + mango_ids.contains(&6), + "id=6 mango sorbet should still be found: {mango_ids:?}" + ); +} + +/// A phrase query must not return a stale hit for an overlaid FTS-indexed row. Phrase queries +/// have no flat re-evaluation path, so the fragment is excluded from the indexed phrase search +/// (like an unindexed fragment) rather than re-scored — the point of this test is that the +/// pre-overlay phrase hit is dropped, not that the new value is found. +#[tokio::test] +async fn test_fts_phrase_overlay_stale_drop() { + let mut dataset = create_text_dataset().await; + build_text_fts_index_with_positions(&mut dataset).await; + + // Before any overlay the phrase "apple banana" matches only id=1. + assert_eq!( + fts_phrase_ids_matching(&dataset, "apple banana").await, + vec![1] + ); + + // Overlay id=1's text (field 1) so the phrase no longer applies to its current value. + let dataset = commit_overlay( + dataset, + "phrase_overlay", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![Arc::new(StringArray::from(vec![Some("cherry mango")]))], + ) + .await; + + // The stale inverted-index positions for "apple banana" on id=1 must not be returned. + assert_eq!( + fts_phrase_ids_matching(&dataset, "apple banana").await, + Vec::::new() + ); +} + +/// An overlay on a non-FTS field must not exclude the fragment from phrase search. +#[tokio::test] +async fn test_fts_phrase_overlay_unrelated_field_not_excluded() { + let mut dataset = create_text_dataset().await; + build_text_fts_index_with_positions(&mut dataset).await; + + // Overlay field 0 (`id`), not the FTS-indexed `text` column: phrase coverage is untouched. + let dataset = commit_overlay( + dataset, + "id_overlay", + 0, + &[0], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![i32_array([Some(777)])], + ) + .await; + + assert_eq!( + fts_phrase_ids_matching(&dataset, "apple banana").await, + vec![777] + ); +} + +/// An overlay on a field the FTS index does NOT cover must not exclude anything. +#[tokio::test] +async fn test_fts_overlay_unrelated_field_not_excluded() { + let mut dataset = create_text_dataset().await; + build_text_fts_index(&mut dataset).await; + + // Overlay field 0 (id) — not covered by the FTS index on `text`. + let dataset = commit_overlay( + dataset, + "id_overlay_for_fts", + 0, + &[0], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![i32_array([Some(999)])], + ) + .await; + + // FTS coverage must be unchanged — both rows containing "apple" are still returned. + // The `id` overlay changes row offset 1's id from 1 to 999, so the projected id column + // reflects the overlay even though the FTS index correctly returned that row. + assert_eq!(fts_ids_matching(&dataset, "apple").await, vec![0, 999]); + assert_eq!(fts_ids_matching(&dataset, "banana").await, vec![3, 999]); +} + +/// Benchmark: measure query latency for BTree, FTS, and vector ANN with 0/4/16 overlay layers. +/// +/// Run with: cargo test -p lance --lib --release -- overlay_index_masking::bench --ignored --nocapture +#[tokio::test] +#[ignore = "benchmark"] +#[allow(clippy::print_stdout)] +async fn bench_index_query_overlay_overhead() { + use std::time::Instant; + + use arrow_array::Float32Array; + + const DIM: i32 = 32; + const ROWS: i32 = 1_000_000; + const ROWS_PER_FRAG: i32 = 100_000; // 10 fragments + const ITERS: u32 = 10; // large scans — 10 is enough for stable averages + + // Fixed disk path so timings are comparable across runs. Deleted and recreated fresh. + let uri = "/tmp/lance-bench-overlay-oss1325"; + if std::path::Path::new(uri).exists() { + std::fs::remove_dir_all(uri).unwrap(); + } + + // --- Build 1M-row dataset on local disk -------------------------------- + // Schema: id(0), age(1), vec(2) — 3 top-level fields. + // Lance field IDs (depth-first): id=0, age=1, vec=2, vec.item=3. + + println!("Building {ROWS}-row dataset at {uri} (this takes ~30 s)..."); + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("age", DataType::Int32, false), + ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + DIM, + ), + false, + ), + ])); + + let row_ids: Vec = (0..ROWS).collect(); + let ages: Vec = row_ids.iter().map(|&i| i * 10).collect(); + // Build the 128 MB flat float array directly (avoids 1M per-row Vec allocations). + let flat_vecs: Vec = (0..(ROWS as usize * DIM as usize)) + .map(|j| (j / DIM as usize) as f32 % 1000.0) + .collect(); + let vec_col = Arc::new( + arrow_array::FixedSizeListArray::try_new( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + DIM, + Arc::new(Float32Array::from(flat_vecs)), + None, + ) + .unwrap(), + ); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(row_ids)), + Arc::new(Int32Array::from(ages)), + vec_col, + ], + ) + .unwrap(); + + let write_params = WriteParams { + max_rows_per_file: ROWS_PER_FRAG as usize, + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let mut dataset = Dataset::write(reader, uri, Some(write_params)) + .await + .unwrap(); + + println!("Building BTree index on age..."); + dataset + .create_index( + &["age"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + + println!("Building IVF_FLAT(1 partition) index on vec..."); + dataset + .create_index( + &["vec"], + IndexType::Vector, + None, + &VectorIndexParams::ivf_flat(1, MetricType::L2), + true, + ) + .await + .unwrap(); + + println!("Indexes built.\n"); + + // --- Timing helper --------------------------------------------------- + + async fn timeit(iters: u32, mut f: F) -> f64 + where + F: FnMut() -> Fut, + Fut: std::future::Future, + { + f().await; // warmup + let t0 = Instant::now(); + for _ in 0..iters { + f().await; + } + t0.elapsed().as_secs_f64() * 1000.0 / iters as f64 + } + + // === Scenario A: BTree query overhead ================================ + // + // Overlay on `age` (field 1), covering only offset 0 of fragment 0. + // Fragment granularity: the entire fragment 0 (100k rows) falls to flat-scan. + // + // btree_cold: `age = 420` → id=42 → in fragment 0 (rows 0..99999). + // With overlays: 100k-row flat scan + per-overlay merge instead of index lookup. + // Without overlays: O(log n) BTree lookup. + // + // btree_warm: `age = 1000420` → id=100042 → in fragment 1 (rows 100000..199999). + // Always served by the BTree index regardless of overlay count on fragment 0. + // This isolates the index-lookup baseline. + println!("=== Scenario A: BTree (overlay on `age`, fragment 0 becomes stale) ==="); + println!( + "{:>10} {:>14} {:>14}", + "overlays", "cold_frag0_ms", "warm_frag1_ms" + ); + + let mut committed_a = 0u32; + for num_overlays in [0u32, 1, 4, 16] { + // Commit only the delta since the last iteration. + for layer in committed_a..num_overlays { + dataset = commit_overlay( + dataset, + &format!("age_ol{layer}"), + 0, // fragment 0 + &[1], // field 1 = age + OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + vec![i32_array([Some(999)])], + ) + .await; + } + committed_a = num_overlays; + + let ds = Arc::new(dataset.clone()); + + // Cold path: stale fragment falls to flat scan when overlays > 0. + let ds2 = ds.clone(); + let cold_ms = timeit(ITERS, || { + let ds = ds2.clone(); + async move { + ds.scan() + .filter("age = 420") + .unwrap() + .project(&["age"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + } + }) + .await; + + // Warm path: fragment 1 never stale, always index-served. + let ds2 = ds.clone(); + let warm_ms = timeit(ITERS, || { + let ds = ds2.clone(); + async move { + ds.scan() + .filter("age = 1000420") + .unwrap() + .project(&["age"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + } + }) + .await; + + println!("{num_overlays:>10} {cold_ms:>14.1} {warm_ms:>14.1}"); + } + + // === Scenario B: Vector ANN overhead ================================= + // + // Overlay on `vec` (field 2), covering only offset 0 of fragment 0. + // The field-aware check means the 16 age overlays from Scenario A do NOT affect + // the vector index (they touch field 1, not field 2). Only a vec overlay (field 2) + // marks fragment 0 stale for the vector index. + // + // With a vec overlay: 100k rows of fragment 0 are excluded from ANN prefilter + // bitmaps and re-scored brute-force (O(100k × DIM) distance computations). + println!("\n=== Scenario B: Vector ANN (overlay on `vec`, 100k rows brute-forced) ==="); + println!("{:>12} {:>10}", "vec_overlays", "ann_ms"); + + let query_vec = Float32Array::from(vec![0.5f32; DIM as usize]); + + for num_vec_overlays in [0u32, 1] { + if num_vec_overlays == 1 { + dataset = commit_overlay( + dataset, + "vec_ol0", + 0, // fragment 0 + &[2], // field 2 = vec (FixedSizeList top-level field) + OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + vec![fsl(vec![vec![0.0f32; DIM as usize]], DIM)], + ) + .await; + } + + let ds = Arc::new(dataset.clone()); + let ds2 = ds.clone(); + let qv = query_vec.clone(); + let ann_ms = timeit(ITERS, || { + let ds = ds2.clone(); + let q = qv.clone(); + async move { + ds.scan() + .nearest("vec", &q, 10) + .unwrap() + .minimum_nprobes(1) + .project(&["id"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + } + }) + .await; + + println!("{num_vec_overlays:>12} {ann_ms:>10.1}"); + } +} diff --git a/rust/lance/src/dataset/tests/mod.rs b/rust/lance/src/dataset/tests/mod.rs index ecc64587b0c..76713abb8d1 100644 --- a/rust/lance/src/dataset/tests/mod.rs +++ b/rust/lance/src/dataset/tests/mod.rs @@ -11,6 +11,7 @@ mod dataset_index; mod dataset_io; mod dataset_merge_update; mod dataset_migrations; +mod dataset_overlay_index_masking; mod dataset_scanner; mod dataset_schema_evolution; mod dataset_transactions; diff --git a/rust/lance/src/index/prefilter.rs b/rust/lance/src/index/prefilter.rs index 071f1b8893d..1e02b6f807a 100644 --- a/rust/lance/src/index/prefilter.rs +++ b/rust/lance/src/index/prefilter.rs @@ -52,6 +52,10 @@ pub struct DatasetPreFilter { // Fragment IDs whose data is still in the index but has been removed from the dataset. // Used by FTS merge-on-read to prune stale fragments at search time. pub(super) deleted_fragments: Option, + // Row addresses whose index entries are stale due to a newer data overlay committed after + // the index was built. Computed synchronously at plan time and ANDead into the final mask + // so the index never returns those rows. + pub(super) overlay_block: Option, // When the tasks are finished this is the combined filter pub(super) final_mask: Mutex>>, } @@ -83,6 +87,7 @@ impl DatasetPreFilter { deleted_ids, filtered_ids, deleted_fragments: None, + overlay_block: None, final_mask: Mutex::new(OnceCell::new()), } } @@ -226,6 +231,13 @@ impl DatasetPreFilter { self.deleted_fragments = Some(fragments); } + /// Block specific row addresses from index results because their index entries are stale + /// due to a data overlay committed after the index was built. + pub fn with_overlay_block(mut self, block: RowAddrMask) -> Self { + self.overlay_block = Some(block); + self + } + /// Creates a task to load a mask that filters out deleted rows and, /// when `restrict_to_fragments` is true, also restricts results to only /// the given `fragments`. @@ -383,6 +395,9 @@ impl PreFilter for DatasetPreFilter { } combined = combined & RowAddrMask::from_block(block_list); } + if let Some(overlay_block) = &self.overlay_block { + combined = combined & overlay_block.clone(); + } Arc::new(combined) }); @@ -393,6 +408,7 @@ impl PreFilter for DatasetPreFilter { self.deleted_ids.is_none() && self.filtered_ids.is_none() && self.deleted_fragments.is_none() + && self.overlay_block.is_none() } /// Get the row id mask for this prefilter diff --git a/rust/lance/src/index/vector/fixture_test.rs b/rust/lance/src/index/vector/fixture_test.rs index a03cc313466..2af020e1df2 100644 --- a/rust/lance/src/index/vector/fixture_test.rs +++ b/rust/lance/src/index/vector/fixture_test.rs @@ -278,6 +278,7 @@ mod test { deleted_ids: None, filtered_ids: None, deleted_fragments: None, + overlay_block: None, final_mask: Mutex::new(OnceCell::new()), }), &NoOpMetricsCollector, diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index 3e264f192ad..783d2aff9d3 100644 --- a/rust/lance/src/io/exec/filtered_read.rs +++ b/rust/lance/src/io/exec/filtered_read.rs @@ -100,6 +100,18 @@ impl EvaluatedIndex { applicable_fragments, }) } + + /// Block `rows` (stale overlay row addresses) from the index result so the index never + /// emits them. Their fragments stay in the covered set, so non-stale rows keep the index; + /// the blocked rows are re-evaluated against their current (overlay-merged) values on a + /// separate targeted take path built by the scanner. + fn without_rows(mut self, block_list: &RowAddrTreeMap) -> Self { + self.index_result.upper = + std::mem::take(&mut self.index_result.upper).also_block(block_list.clone()); + self.index_result.lower = + std::mem::take(&mut self.index_result.lower).also_block(block_list.clone()); + self + } } /// A fragment along with ranges of row offsets to read @@ -1505,6 +1517,12 @@ pub struct FilteredReadOptions { pub io_buffer_size_bytes: Option, /// If true, skip fragments that are not covered by the scalar index result. pub only_indexed_fragments: bool, + /// Row addresses whose index entries may be stale because an overlay committed after the + /// index was built touches an indexed field. They are blocked from the index result so the + /// index never emits them; the scanner re-evaluates just these rows against their current + /// (overlay-merged) values on a targeted take path. Their fragments stay in the covered set, + /// so non-stale rows keep the index. `None` on the common no-overlay fast path. + pub overlay_block: Option, } impl FilteredReadOptions { @@ -1534,12 +1552,20 @@ impl FilteredReadOptions { full_filter: None, io_buffer_size_bytes: None, only_indexed_fragments: false, + overlay_block: None, threading_mode: FilteredReadThreadingMode::OnePartitionMultipleThreads( get_num_compute_intensive_cpus(), ), } } + /// Block the given stale overlay row addresses (see the `overlay_block` field) from the + /// scalar index result so the index never emits them. + pub fn with_overlay_block(mut self, block: RowAddrMask) -> Self { + self.overlay_block = Some(block); + self + } + /// Include deleted rows in the scan /// /// This is currently only supported if there is no scan_range specified @@ -2117,9 +2143,15 @@ impl FilteredReadExec { let index_search_result = index_search.next().await.ok_or_else(|| { Error::internal("Index search did not yield any results".to_string()) })??; - evaluated_index = Some(Arc::new(EvaluatedIndex::try_from_arrow( - &index_search_result, - )?)); + let mut idx = EvaluatedIndex::try_from_arrow(&index_search_result)?; + // `overlay_block` is always constructed as a block list (see + // `Scanner::stale_rows_block_mask`), so `block_list()` is always `Some`. + if let Some(block_list) = + options.overlay_block.as_ref().and_then(|b| b.block_list()) + { + idx = idx.without_rows(block_list); + } + evaluated_index = Some(Arc::new(idx)); } // Load fragments to compute the plan diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index fde0289621d..e5157f78818 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -60,6 +60,8 @@ use lance_table::format::IndexMetadata; use tokio::sync::Notify; use uuid::Uuid; +use lance_select::RowAddrMask; + use crate::dataset::Dataset; use crate::index::DatasetIndexInternalExt; use crate::index::prefilter::{DatasetPreFilter, FilterLoader}; @@ -1099,11 +1101,14 @@ pub static KNN_PARTITION_SCHEMA: LazyLock = LazyLock::new(|| { ])) }); +/// Create a new ANN execution node. `overlay_block`, when `Some`, excludes rows whose index +/// entries may be stale due to a newer data overlay (see [`ANNIvfSubIndexExec::with_overlay_block`]). pub fn new_knn_exec( dataset: Arc, indices: &[IndexMetadata], query: &Query, prefilter_source: PreFilterSource, + overlay_block: Option, ) -> Result> { let ivf_node = ANNIvfPartitionExec::try_new( dataset.clone(), @@ -1111,13 +1116,16 @@ pub fn new_knn_exec( query.clone(), )?; - let sub_index = ANNIvfSubIndexExec::try_new( + let mut sub_index = ANNIvfSubIndexExec::try_new( Arc::new(ivf_node), dataset, indices.to_vec(), query.clone(), prefilter_source, )?; + if let Some(overlay_block) = overlay_block { + sub_index = sub_index.with_overlay_block(overlay_block); + } Ok(Arc::new(sub_index)) } @@ -1377,6 +1385,10 @@ pub struct ANNIvfSubIndexExec { /// Prefiltering input prefilter_source: PreFilterSource, + /// Row addresses whose index entries are stale due to a newer data overlay. Blocked from + /// index results at execution time via [`DatasetPreFilter::with_overlay_block`]. + overlay_block: Option, + /// Datafusion Plan Properties properties: Arc, @@ -1409,11 +1421,18 @@ impl ANNIvfSubIndexExec { indices, query, prefilter_source, + overlay_block: None, properties, metrics: ExecutionPlanMetricsSet::new(), }) } + /// Block stale row addresses (see the `overlay_block` field) from index results. + pub fn with_overlay_block(mut self, overlay_block: RowAddrMask) -> Self { + self.overlay_block = Some(overlay_block); + self + } + /// Returns a reference to the vector query. pub fn query(&self) -> &Query { &self.query @@ -1880,6 +1899,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { indices: self.indices.clone(), query: self.query.clone(), prefilter_source, + overlay_block: self.overlay_block.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), } @@ -1960,11 +1980,13 @@ impl ExecutionPlan for ANNIvfSubIndexExec { PreFilterSource::None => None, }; - let pre_filter = Arc::new(DatasetPreFilter::new( - ds.clone(), - &indices, - prefilter_loader, - )); + let pre_filter = { + let mut pf = DatasetPreFilter::new(ds.clone(), &indices, prefilter_loader); + if let Some(block) = self.overlay_block.clone() { + pf = pf.with_overlay_block(block); + } + Arc::new(pf) + }; let state = Arc::new(ANNIvfEarlySearchResults::new(indices.len(), query.k)); diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index e67ba3b7b34..b943ad12d3a 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -6,7 +6,7 @@ use std::sync::{Arc, LazyLock}; use super::utils::{IndexMetrics, InstrumentedRecordBatchStreamAdapter}; use crate::{ Dataset, - dataset::rowids::{load_row_id_sequence, load_row_id_sequences}, + dataset::rowids::{load_row_id_sequences, translate_addr_treemap_to_row_ids}, index::{ prefilter::DatasetPreFilter, scalar_logical::{open_named_scalar_index, scalar_index_fragment_bitmap}, @@ -43,7 +43,7 @@ use lance_index::{ }; use lance_select::{ IndexExprResult, NullableIndexExprResult, NullableRowAddrMask, NullableRowAddrSet, RowAddrMask, - RowAddrSelection, RowAddrTreeMap, RowSetOps, result::IndexExprResultWireFormat, + RowAddrTreeMap, RowSetOps, result::IndexExprResultWireFormat, }; use lance_table::format::Fragment; use roaring::RoaringBitmap; @@ -105,68 +105,6 @@ async fn translate_addr_set_to_row_ids( Ok(NullableRowAddrSet::new(selected, nulls)) } -/// Map a set of physical row addresses to their stable row ids -/// -/// For each fragment present in `addrs`, the live rows in physical order carry -/// the stable ids yielded by the fragment's [`RowIdSequence`] in the same -/// order. Zipping the two (skipping deleted physical offsets) gives the -/// `physical offset -> stable id` mapping. Addresses that point at deleted rows -/// have no live counterpart and are dropped, which is correct: those rows are -/// not part of the answer. -async fn translate_addr_treemap_to_row_ids( - dataset: &Dataset, - addrs: &RowAddrTreeMap, -) -> Result { - let mut row_ids = RowAddrTreeMap::new(); - for (fragment_id, selection) in addrs.iter() { - let file_fragment = dataset.get_fragment(*fragment_id as usize).ok_or_else(|| { - Error::internal(format!( - "fragment {fragment_id} referenced by an address-domain index result \ - was not found in the dataset" - )) - })?; - let sequence = load_row_id_sequence(dataset, file_fragment.metadata()).await?; - - match selection { - RowAddrSelection::Full => { - // The whole fragment is selected: every live row's id qualifies. - row_ids |= RowAddrTreeMap::from(sequence.as_ref()); - } - RowAddrSelection::Partial(offsets) => { - let Some(max_offset) = offsets.max() else { - continue; - }; - let (deletion_vector, num_physical_rows) = futures::try_join!( - file_fragment.get_deletion_vector(), - file_fragment.physical_rows() - )?; - let num_physical_rows = num_physical_rows as u32; - let mut ids = sequence.iter(); - for physical_offset in 0..num_physical_rows { - if physical_offset > max_offset { - break; - } - let deleted = deletion_vector - .as_ref() - .is_some_and(|dv| dv.contains(physical_offset)); - if deleted { - continue; - } - match ids.next() { - Some(id) => { - if offsets.contains(physical_offset) { - row_ids.insert(id); - } - } - None => break, - } - } - } - } - } - Ok(row_ids) -} - /// An execution node that performs a scalar index search /// /// This does not actually scan any data. We only look through the index to determine @@ -660,6 +598,10 @@ pub struct MaterializeIndexExec { dataset: Arc, expr: ScalarIndexExpr, fragments: Arc>, + /// Row addresses blocked from the index result due to data overlay files committed after the + /// index was built. ANDead into the candidate mask before row ID materialisation so that stale + /// index entries never reach downstream operators. + overlay_block: Option, properties: Arc, metrics: ExecutionPlanMetricsSet, } @@ -732,16 +674,24 @@ impl MaterializeIndexExec { dataset, expr, fragments, + overlay_block: None, properties, metrics: ExecutionPlanMetricsSet::new(), } } + /// Block specific row addresses (see the `overlay_block` field) from the index result. + pub fn with_overlay_block(mut self, block: RowAddrMask) -> Self { + self.overlay_block = Some(block); + self + } + #[instrument(name = "materialize_scalar_index", skip_all, level = "debug")] async fn do_execute( expr: ScalarIndexExpr, dataset: Arc, fragments: Arc>, + overlay_block: Option, metrics: Arc, ) -> Result { let expr_result = expr.evaluate(dataset.as_ref(), metrics.as_ref()); @@ -769,12 +719,15 @@ impl MaterializeIndexExec { } Ok(result.upper) }; - let mask = if let Some(prefilter) = prefilter { + let mut mask = if let Some(prefilter) = prefilter { let (expr_result, prefilter) = futures::try_join!(expr_result, prefilter)?; take_upper(expr_result)? & (*prefilter).clone() } else { take_upper(expr_result.await?)? }; + if let Some(block) = overlay_block { + mask = mask & block; + } let ids = row_ids_for_mask(mask, &dataset, &fragments).await?; let ids = UInt64Array::from(ids); Ok(RecordBatch::try_new( @@ -906,6 +859,7 @@ impl ExecutionPlan for MaterializeIndexExec { self.expr.clone(), self.dataset.clone(), self.fragments.clone(), + self.overlay_block.clone(), metrics, ); let stream = futures::stream::iter(vec![batch_fut])