fix(index): keep overlays masked when updating a non-indexed column

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::<i32>::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::<i32>::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 {

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<u32, &Fragment> = 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<u32, &Fragment>,
+    ) -> 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