From f1c4674fbb0c84f94116e9387148a67eef3e9aa2 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Wed, 22 Jul 2026 17:47:54 -0700 Subject: [PATCH 1/2] fix(index): keep overlays masked when updating a non-indexed column A rewrite-rows update (under stable row ids) that sets only a non-indexed column moves the matched rows to a new fragment. Because the update did not touch the scalar index's field, `register_pure_rewrite_rows_update_frags_in_indices` extends that index's coverage onto the new fragment and reuses its existing entries, on the assumption that the moved rows' indexed value is unchanged. That assumption breaks 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 marking the new fragment covered removes the flat-path re-evaluation that overlay masking relied on. The overlaid row then vanishes from index-path queries (a range scan returns 99 rows where 100 are expected). Skip extending coverage for an index whose fields are touched by an overlay committed after the index was built on any of the moved rows' original fragments, so those new fragments fall to the flat path against current values. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../tests/dataset_overlay_index_masking.rs | 66 +++++++++++++++++++ rust/lance/src/dataset/transaction.rs | 62 ++++++++++++----- 2 files changed, 111 insertions(+), 17 deletions(-) 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..5b8c937fb5e 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,42 @@ 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 || index.fragment_bitmap.is_none() { + 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 + // 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 = { + let fragment_bitmap = index.fragment_bitmap.as_ref().unwrap(); + original_fragment_ids .iter() - .all(|&fragment_id| fragment_bitmap.contains(fragment_id as u32)); + .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; + } + + let fragment_bitmap = index.fragment_bitmap.as_mut().unwrap(); + 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 From 58d5b2a66dca115fc2eae343f17d0f76ddd5a207 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 23 Jul 2026 10:00:53 -0700 Subject: [PATCH 2/2] refactor: bind fragment_bitmap with let-else instead of is_none + unwrap Addresses CodeRabbit review: replace the check-then-unwrap pattern (`is_none()` followed by `as_ref().unwrap()`/`as_mut().unwrap()`) with a `let ... else` binding and an `if let`, per the crate's coding guidelines. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/transaction.rs | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs index 5b8c937fb5e..37a5367a9b5 100644 --- a/rust/lance/src/dataset/transaction.rs +++ b/rust/lance/src/dataset/transaction.rs @@ -2692,18 +2692,18 @@ 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 || index.fragment_bitmap.is_none() { + if index_covers_modified_field { continue; } + let Some(fragment_bitmap) = index.fragment_bitmap.as_ref() else { + continue; + }; // 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 = { - let fragment_bitmap = index.fragment_bitmap.as_ref().unwrap(); - original_fragment_ids - .iter() - .all(|&fragment_id| fragment_bitmap.contains(fragment_id as u32)) - }; + 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; } @@ -2722,9 +2722,10 @@ impl Transaction { continue; } - let fragment_bitmap = index.fragment_bitmap.as_mut().unwrap(); - for fragment_id in pure_update_frag_ids.iter().map(|f| *f as u32) { - fragment_bitmap.insert(fragment_id); + 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(())