From 87b8894d7faefee0779d11494ef5434f6a5a7f06 Mon Sep 17 00:00:00 2001 From: Ragnor Comerford Date: Wed, 27 May 2026 17:19:23 +0200 Subject: [PATCH] fix: apply fragment-bitmap allow-list to stable-row-id deletion mask MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `do_create_deletion_mask_row_id` previously iterated every fragment in the dataset, so the resulting allow-list included stable row ids whose *current* physical home was outside the index's `fragment_bitmap`. On stable-row-id datasets this bypasses the allow-list added in #6563: `MapIndexExec` requests `create_restricted_deletion_mask`, but the stable-row-id branch at `prefilter.rs:270` ignores the restriction. Effect on merge_insert: after a `merge_insert UpdateAll` rewrites a row into a new fragment, a second merge_insert against the same key sees that row twice — once via the BTREE (which holds the row's stable row id and resolves to its new fragment through TakeExec) and once via the unindexed-fragments scan that also covers the new fragment. Both branches emit the same `_rowid`, so the source-dedup HashSet trips and the merge fails with "Ambiguous merge inserts are prohibited". Effect on FTS: a `full_text_search` after a merge_insert returns each moved row twice, because the inverted index's hits go through the same prefilter and the post-filter scan picks up the new fragment. Fix: thread `restrict_to_fragments` through `do_create_deletion_mask_row_id` so it only iterates fragments in the bitmap. The resulting allow-list now means "stable row ids whose current home is inside the restriction" — semantically equivalent to the non-stable-row-id branch. The cache key gains an optional bitmap hash so two consumers requesting different fragment subsets do not poison each other's cached mask. Closes #6877. --- rust/lance/src/dataset/write/merge_insert.rs | 247 +++++++++++++++++++ rust/lance/src/index/prefilter.rs | 101 +++++++- rust/lance/src/session/caches.rs | 9 +- 3 files changed, 350 insertions(+), 7 deletions(-) diff --git a/rust/lance/src/dataset/write/merge_insert.rs b/rust/lance/src/dataset/write/merge_insert.rs index e9432fbad27..422b271491e 100644 --- a/rust/lance/src/dataset/write/merge_insert.rs +++ b/rust/lance/src/dataset/write/merge_insert.rs @@ -8589,6 +8589,131 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n } } + // Regression test for GitHub issue #6877. + // + // Two sequential full-schema merge_insert UpdateAll calls against the same + // target row, on a dataset with stable_row_ids enabled and a BTREE scalar + // index on the join column, used to fail on the second call with + // "Ambiguous merge inserts are prohibited" — even though each call's + // source had exactly one row per key. + // + // Mechanism: with stable row ids the BTREE stores stable_row_ids (not + // physical addresses). After the first merge_insert, A's stable_row_id is + // preserved but its physical home moves to an unindexed fragment. The + // BTREE-side TakeExec resolves the stable_row_id to A's new location and + // emits a row; the unindexed-fragments scan also covers the new fragment + // and emits the same logical row. Both surface the same `_rowid`, so the + // merge_insert source-dedup HashSet sees a duplicate and aborts. + // + // Fix: thread `restrict_to_fragments` into `do_create_deletion_mask_row_id` + // so the allow-list only contains stable_row_ids whose current physical + // home is inside the index's fragment_bitmap. + #[tokio::test] + async fn test_issue_6877_repeated_merge_insert_stable_row_ids() { + use arrow_array::Int32Array; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("value", DataType::Int32, false), + ])); + + let initial = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["A", "B", "C"])), + Arc::new(Int32Array::from(vec![1, 2, 3])), + ], + ) + .unwrap(); + + let mut ds = Dataset::write( + Box::new(RecordBatchIterator::new([Ok(initial)], schema.clone())), + "memory://test_6877", + Some(WriteParams { + mode: WriteMode::Overwrite, + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + + ds.create_index( + &["id"], + IndexType::Scalar, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + + // First merge_insert: A 1 -> 11. + let update_a = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["A"])), + Arc::new(Int32Array::from(vec![11])), + ], + ) + .unwrap(); + let (ds, _) = MergeInsertBuilder::try_new(Arc::new(ds), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new( + [Ok(update_a)], + schema.clone(), + ))) + .await + .unwrap(); + + // Second merge_insert: A 11 -> 22. Used to fail before the fix. + let update_a_again = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["A"])), + Arc::new(Int32Array::from(vec![22])), + ], + ) + .unwrap(); + let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new( + [Ok(update_a_again)], + schema.clone(), + ))) + .await + .unwrap(); + + // Sanity check: A's value is now 22. + let batches = ds + .scan() + .filter("id = 'A'") + .unwrap() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let combined = concat_batches(&schema, &batches).unwrap(); + assert_eq!(combined.num_rows(), 1); + let values = combined + .column_by_name("value") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), 22); + } + // Regression test: partial-schema merge_insert followed by update (deleting all rows // in a fragment) followed by partial merge_insert should not produce // "fragment id N does not exist" errors. @@ -8945,6 +9070,128 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n ); } + // Companion regression test for issue #6877 on the FTS path. + // + // The FTS prefilter shares `do_create_deletion_mask_row_id` with the + // scalar-index path, so the same stable-row-id bypass that produced + // duplicate rows in merge_insert can produce duplicate hits in FTS search + // after a merge_insert moves rows to unindexed fragments. This test pins + // the contract for the FTS consumer. + #[tokio::test] + async fn test_issue_6877_fts_no_duplicates_stable_row_ids() { + let rows_per_frag = 10usize; + let num_frags = 3usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("text", DataType::Utf8, false), + ])); + + let make_batch = |frag_idx: usize| { + let start = frag_idx * rows_per_frag; + let ids: Vec = (start..start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let texts: Vec = (start..start + rows_per_frag) + .map(|j| format!("common unique{j:04}")) + .collect(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(texts)), + ], + ) + .unwrap() + }; + + let batch0 = make_batch(0); + let reader = Box::new(RecordBatchIterator::new([Ok(batch0)], schema.clone())); + let mut ds = Dataset::write( + reader, + "memory://fts_stable_row_id_test", + Some(WriteParams { + mode: WriteMode::Overwrite, + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + for frag_idx in 1..num_frags { + let batch = make_batch(frag_idx); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + ds.append(reader, None).await.unwrap(); + } + + let params = InvertedIndexParams::default(); + ds.create_index(&["text"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + // Full-schema merge_insert rewriting fragment 1's rows. After this, + // the original locations are tombstoned and the new locations live in + // a new (unindexed) fragment; the stable_row_ids are preserved. + let frag1_start = rows_per_frag; + let ids: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let texts: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("common updated{j:04}")) + .collect(); + let update_batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(texts)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], schema.clone())); + let (ds, _) = MergeInsertBuilder::try_new(Arc::new(ds), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .unwrap(); + + // FTS search for "common" — every row should match exactly once. + let query = FullTextSearchQuery::new("common".to_string()); + let results = ds + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + + let ids = results + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let unique_ids: std::collections::HashSet<&str> = + (0..ids.len()).map(|i| ids.value(i)).collect(); + assert_eq!( + unique_ids.len(), + ids.len(), + "Found duplicate ids in FTS results: {} unique out of {} total", + unique_ids.len(), + ids.len() + ); + assert_eq!( + unique_ids.len(), + rows_per_frag * num_frags, + "Expected {} rows but got {}", + rows_per_frag * num_frags, + unique_ids.len() + ); + } + // Regression test: after a partial-schema merge_insert invalidates a fragment, // compaction should succeed and subsequent searches should return correct results. // diff --git a/rust/lance/src/index/prefilter.rs b/rust/lance/src/index/prefilter.rs index ebb4c600e65..071f1b8893d 100644 --- a/rust/lance/src/index/prefilter.rs +++ b/rust/lance/src/index/prefilter.rs @@ -127,13 +127,33 @@ impl DatasetPreFilter { } #[instrument(level = "debug", skip_all)] - async fn do_create_deletion_mask_row_id(dataset: Arc) -> Result> { - // This can only be computed as an allow list, since we have no idea - // what the row ids were in the missing fragments. + async fn do_create_deletion_mask_row_id( + dataset: Arc, + restrict_to: Option, + ) -> Result> { + // The mask is an allow-list of stable row ids. When `restrict_to` is + // set the iteration is limited to the listed fragments, so the + // resulting list excludes stable row ids whose *current* physical home + // is outside the restriction. This is the missing piece for the + // stable-row-id branch of #6563: without it, the merge_insert UNION + // (indexed scan ∪ unindexed-fragments scan) sees the same logical row + // twice — once via the BTREE (which holds the row's stable_row_id) and + // once via the unindexed scan (which holds the fragment the row now + // lives in). See issue #6877. async fn load_row_ids_and_deletions( dataset: &Dataset, + restrict_to: Option<&RoaringBitmap>, ) -> Result, Option>)>> { - stream::iter(dataset.get_fragments()) + let frags: Vec<_> = dataset + .get_fragments() + .into_iter() + .filter(|f| { + restrict_to + .map(|allow| allow.contains(f.id() as u32)) + .unwrap_or(true) + }) + .collect(); + stream::iter(frags) .map(|frag| async move { let row_ids = load_row_id_sequence(dataset, frag.metadata()); let deletion_vector = frag.get_deletion_vector(); @@ -145,16 +165,31 @@ impl DatasetPreFilter { .await } + let restrict_hash = restrict_to.as_ref().map(|b| { + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + let mut h = DefaultHasher::new(); + // RoaringBitmap doesn't implement Hash; serialize the sorted u32s. + for v in b.iter() { + v.hash(&mut h); + } + h.finish() + }); + let dataset_clone = dataset.clone(); + let restrict_for_load = restrict_to.clone(); let key = crate::session::caches::RowAddrMaskKey { version: dataset.manifest().version, + restrict_hash, }; dataset .metadata_cache .as_ref() .get_or_insert_with_key(key, move || { async move { - let row_ids_and_deletions = load_row_ids_and_deletions(&dataset_clone).await?; + let row_ids_and_deletions = + load_row_ids_and_deletions(&dataset_clone, restrict_for_load.as_ref()) + .await?; // The process of computing the final mask is CPU-bound, so we spawn it // on a blocking thread. @@ -268,7 +303,12 @@ impl DatasetPreFilter { if missing_frags.is_empty() && frags_with_deletion_files.is_empty() && !needs_allow_list { None } else if dataset.manifest.uses_stable_row_ids() { - Some(Self::do_create_deletion_mask_row_id(dataset.clone()).boxed()) + let restrict_to = if restrict_to_fragments { + Some(fragments) + } else { + None + }; + Some(Self::do_create_deletion_mask_row_id(dataset.clone(), restrict_to).boxed()) } else if missing_frags.is_empty() && frags_with_deletion_files.is_empty() { // No deletions to load, but the dataset has fragments outside the // index bitmap. Return a synchronous allow-list mask. @@ -529,4 +569,53 @@ mod test { let mask = mask.unwrap().await.unwrap(); assert_eq!(mask.allow_list().and_then(|x| x.len()), Some(3)); // There were three rows left over; } + + // Regression test for issue #6877. + // + // `create_restricted_deletion_mask` on a stable-row-id dataset must honor + // the bitmap restriction by excluding stable row ids whose *current* + // physical home is outside the bitmap. Without this, the merge_insert + // UNION (indexed-scan ∪ unindexed-fragments scan) emits the same logical + // row twice — once via the BTREE (which holds the row's stable_row_id) + // and once via the unindexed scan. + #[tokio::test] + async fn test_restricted_deletion_mask_stable_row_id_honors_bitmap() { + // Dataset with three fragments, 3 rows each, stable_row_ids = {0..9}. + // Row x=8 is deleted, so the live stable_row_ids are {0..7}. + let datasets = test_datasets(true).await; + let ds = datasets.deletions_no_missing_frags.clone(); + + // Full bitmap: allow-list covers all currently-live stable row ids. + let mask = DatasetPreFilter::create_restricted_deletion_mask( + ds.clone(), + RoaringBitmap::from_iter(0..3), + ) + .expect("full-bitmap mask present on stable-row-id dataset with deletions") + .await + .unwrap(); + let expected_all = RowAddrTreeMap::from_iter(0..8); + assert_eq!(mask.allow_list(), Some(&expected_all)); + + // Restricted to fragments {0, 1}: allow-list must exclude stable row + // ids whose current home is in fragment 2 (rows 6, 7 — 8 was deleted). + let mask = DatasetPreFilter::create_restricted_deletion_mask( + ds.clone(), + RoaringBitmap::from_iter(0..2), + ) + .expect("restricted mask present") + .await + .unwrap(); + let expected_restricted = RowAddrTreeMap::from_iter(0..6); + assert_eq!(mask.allow_list(), Some(&expected_restricted)); + + // Restricted to empty bitmap: every BTREE-returned address is filtered + // out. Empty allow-list is the correct semantic ("no row's current home + // is in the restriction"). + let mask = + DatasetPreFilter::create_restricted_deletion_mask(ds.clone(), RoaringBitmap::new()) + .expect("empty-restriction mask present") + .await + .unwrap(); + assert_eq!(mask.allow_list().and_then(|x| x.len()), Some(0)); + } } diff --git a/rust/lance/src/session/caches.rs b/rust/lance/src/session/caches.rs index eab758418f7..82dc755f6c0 100644 --- a/rust/lance/src/session/caches.rs +++ b/rust/lance/src/session/caches.rs @@ -122,12 +122,19 @@ impl CacheKey for DeletionFileKey<'_> { #[derive(Debug)] pub struct RowAddrMaskKey { pub version: u64, + /// `Some(hash)` when the mask is restricted to a fragment subset; `None` + /// when it covers all fragments in the dataset. Two consumers that ask + /// for different subsets must not poison each other's cache entry. + pub restrict_hash: Option, } impl CacheKey for RowAddrMaskKey { type ValueType = RowAddrMask; fn key(&self) -> Cow<'_, str> { - Cow::Owned(format!("row_addr_mask/{}", self.version)) + match self.restrict_hash { + None => Cow::Owned(format!("row_addr_mask/{}", self.version)), + Some(h) => Cow::Owned(format!("row_addr_mask/{}/{:x}", self.version, h)), + } } fn type_name() -> &'static str { "RowAddrMask"