diff --git a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs index cfabdc13266..e6efb3f6ad8 100644 --- a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs +++ b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs @@ -589,6 +589,72 @@ async fn test_overlay_stale_with_compound_index_expression() { assert_eq!(ids_matching(&dataset, "id = 2").await, vec![2]); } +/// A `RewriteRows` update (under stable row ids) that touches only a *non-indexed* column moves +/// the matched rows to a new fragment and, because the scalar index's field was not modified, +/// extends that index's fragment coverage onto the new fragment +/// (`register_pure_rewrite_rows_update_frags_in_indices`) so its existing entries are reused. +/// +/// That reuse is unsound when a moved row carried a data overlay on the *indexed* field: the +/// update materializes the overlay's current value into the new fragment, but the reused index +/// entry still holds the stale pre-overlay value, and the new fragment (now marked covered) no +/// longer falls to the flat path that previously served the correct value via overlay masking. +/// +/// Here `age` is indexed and overlaid (id=1: age 10 -> 999); the update sets the non-indexed +/// `id` column on that row. After it, `age = 10` must stay dropped and `age = 999` must still +/// find the row — otherwise the stale index entry has resurfaced. +#[tokio::test] +async fn test_update_nonindexed_column_preserves_overlay_masking() { + use crate::dataset::UpdateBuilder; + + let mut dataset = create_base_dataset_with(true).await; + build_age_index(&mut dataset).await; + + // Overlay fragment 0, offset 1 (id=1): age 10 -> 999, committed after the index. + let dataset = commit_overlay( + dataset, + "age_update", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![i32_array([Some(999)])], + ) + .await; + + // Masking works before the update. + assert_eq!(ids_matching(&dataset, "age = 10").await, Vec::::new()); + assert_eq!(ids_matching(&dataset, "age = 999").await, vec![1]); + + // Update only the non-indexed `id` column of the overlaid row. This is a rewrite-rows move: + // the row (with age materialized to 999) is written to a new fragment and deleted from + // fragment 0, keeping its stable row id. + let dataset = UpdateBuilder::new(Arc::new(dataset)) + .update_where("id = 1") + .unwrap() + .set("id", "100") + .unwrap() + .build() + .unwrap() + .execute() + .await + .unwrap() + .new_dataset; + + // Still masked: the stale age=10 entry must stay dropped and the overlaid age=999 value must + // still be found (now on the moved row, whose id is 100). + assert_eq!( + ids_matching(&dataset, "age = 10").await, + Vec::::new(), + "stale index entry age=10 resurfaced after updating a non-indexed column" + ); + assert_eq!( + ids_matching(&dataset, "age = 999").await, + vec![100], + "overlaid value age=999 lost after updating a non-indexed column" + ); + // A row untouched by the overlay is unaffected. + assert_eq!(ids_matching(&dataset, "age = 20").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 { diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs index 49a5c154d32..37a5367a9b5 100644 --- a/rust/lance/src/dataset/transaction.rs +++ b/rust/lance/src/dataset/transaction.rs @@ -14,6 +14,7 @@ use super::ManifestWriteConfig; use super::write::merge_insert::inserted_rows::KeyExistenceFilter; +use crate::dataset::overlay::collect_overlay_stale_frags; use crate::dataset::transaction::UpdateMode::{RewriteColumns, RewriteRows}; use crate::index::mem_wal::update_mem_wal_index_merged_generations; use crate::utils::temporal::timestamp_to_nanos; @@ -2043,12 +2044,21 @@ impl Transaction { .copied() .collect(); + // The original fragments that carried an overlay: their moved rows may have a + // stale index entry (see `register_pure_rewrite_rows_update_frags_in_indices`). + let original_overlaid_frags: HashMap = existing_fragments + .iter() + .filter(|f| original_fragment_ids.contains(&f.id) && !f.overlays.is_empty()) + .map(|f| (f.id as u32, f)) + .collect(); + Self::register_pure_rewrite_rows_update_frags_in_indices( &mut final_indices, &pure_updated_frag_ids, &original_fragment_ids, fields_for_preserving_frag_bitmap, - ); + &original_overlaid_frags, + )?; } if let Some(next_row_id) = &mut next_row_id { @@ -2668,9 +2678,10 @@ impl Transaction { pure_update_frag_ids: &[u64], original_fragment_ids: &[u64], fields_for_preserving_frag_bitmap: &[u32], - ) { + original_overlaid_frags: &HashMap, + ) -> Result<()> { if pure_update_frag_ids.is_empty() { - return; + return Ok(()); } let value_updated_field_set = fields_for_preserving_frag_bitmap @@ -2681,25 +2692,43 @@ impl Transaction { let index_covers_modified_field = index.fields.iter().any(|field_id| { value_updated_field_set.contains(&u32::try_from(*field_id).unwrap()) }); + if index_covers_modified_field { + continue; + } + let Some(fragment_bitmap) = index.fragment_bitmap.as_ref() else { + continue; + }; - if !index_covers_modified_field - && let Some(fragment_bitmap) = &mut index.fragment_bitmap - { - // check if all the original fragments contains the updating rows are covered - // by the index(index fragment bitmap contains these frag ids). - // if not, that means not all the updating rows are indexed, so we could not - // index them. - let index_covers_all_original_fragments = original_fragment_ids - .iter() - .all(|&fragment_id| fragment_bitmap.contains(fragment_id as u32)); + // Check that all the original fragments containing the updated rows are covered by + // the index. If not, some updated rows were not indexed, so we cannot index them. + let index_covers_all_original_fragments = original_fragment_ids + .iter() + .all(|&fragment_id| fragment_bitmap.contains(fragment_id as u32)); + if !index_covers_all_original_fragments { + continue; + } - if index_covers_all_original_fragments { - for fragment_id in pure_update_frag_ids.iter().map(|f| *f as u32) { - fragment_bitmap.insert(fragment_id); - } + // Reusing the index entries of the moved rows is only sound when those rows' current + // value equals what the index holds. A moved row that carried an overlay on this + // index's fields (committed after the index was built) breaks that: the update + // materialized the overlay's value into the new fragment, but the index still holds + // the stale pre-overlay value. Extending coverage onto the new fragment would resurface + // that stale entry and strip the flat-path re-evaluation overlay masking relied on. So + // leave the new fragments unindexed for such an index — they fall to the flat path + // against current values. + let mut overlay_stale = RoaringBitmap::new(); + collect_overlay_stale_frags(index, original_overlaid_frags, &mut overlay_stale)?; + if !overlay_stale.is_empty() { + continue; + } + + if let Some(fragment_bitmap) = index.fragment_bitmap.as_mut() { + for fragment_id in pure_update_frag_ids.iter().map(|f| *f as u32) { + fragment_bitmap.insert(fragment_id); } } } + Ok(()) } /// If an operation modifies one or more fields in a fragment then we need to remove