MergeInsertBuilder produces spurious "Ambiguous merge inserts" against rows previously rewritten by merge_insert

[ragnorc]

Summary

Two sequential merge_insert operations against the same target row (or set of rows) fail on the second call with:

Invalid user input: Ambiguous merge inserts are prohibited:
multiple source rows match the same target row on (id = "<key>")

even when the source batch has exactly one row per <key>. Reproduces on Lance 4.0.0 and 4.0.1 (byte-identical in this code path).

Minimal repro

use arrow_array::{RecordBatch, StringArray, Int32Array};
use arrow_schema::{DataType, Field, Schema};
use lance::dataset::{
    Dataset, WriteParams, WriteMode,
    MergeInsertBuilder, WhenMatched, WhenNotMatched,
};
use lance_index::{DatasetIndexExt, IndexType};
use lance_index::scalar::ScalarIndexParams;
use std::sync::Arc;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    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])),
    ])?;

    let uri = "memory://test"; // or a tmpdir path
    let mut ds = Dataset::write(
        Box::new(arrow_array::RecordBatchIterator::new(
            vec![Ok(initial)], schema.clone()
        )),
        uri,
        Some(WriteParams {
            mode: WriteMode::Overwrite,
            enable_stable_row_ids: true,
            ..Default::default()
        }),
    ).await?;

    // Build BTREE on \`id\` (mirrors what downstreams do on every commit).
    ds.create_index(
        &["id"],
        IndexType::Scalar,
        None,
        &ScalarIndexParams::default(),
        false,
    ).await?;

    // First merge_insert: update A's value from 1 to 11.
    let update_a = RecordBatch::try_new(schema.clone(), vec![
        Arc::new(StringArray::from(vec!["A"])),
        Arc::new(Int32Array::from(vec![11])),
    ])?;

    let mut builder = MergeInsertBuilder::try_new(
        Arc::new(ds.clone()),
        vec!["id".to_string()],
    )?;
    builder.when_matched(WhenMatched::UpdateAll);
    builder.when_not_matched(WhenNotMatched::DoNothing);
    let job = builder.try_build()?;
    let reader = arrow_array::RecordBatchIterator::new(
        vec![Ok(update_a.clone())], schema.clone(),
    );
    let (ds, _) = job.execute(
        lance_datafusion::utils::reader_to_stream(Box::new(reader))
    ).await?;
    // OK so far.

    // Second merge_insert: update A's value from 11 to 22.
    let update_a_again = RecordBatch::try_new(schema.clone(), vec![
        Arc::new(StringArray::from(vec!["A"])),
        Arc::new(Int32Array::from(vec![22])),
    ])?;

    let mut builder = MergeInsertBuilder::try_new(
        Arc::new(ds.clone()),
        vec!["id".to_string()],
    )?;
    builder.when_matched(WhenMatched::UpdateAll);
    builder.when_not_matched(WhenNotMatched::DoNothing);
    let job = builder.try_build()?;
    let reader = arrow_array::RecordBatchIterator::new(
        vec![Ok(update_a_again)], schema,
    );
    let (_, _) = job.execute(
        lance_datafusion::utils::reader_to_stream(Box::new(reader))
    ).await?;
    // Errors here:
    // "Ambiguous merge inserts are prohibited: multiple source rows
    //  match the same target row on (id = \"A\")"

    Ok(())
}

Workarounds that suppress the symptom

Either of these on the second MergeInsertBuilder is sufficient (cross-validated downstream — each alone suppresses the bug; we ship both as defense in depth):

builder.use_index(false);
// OR
builder.source_dedupe_behavior(SourceDedupeBehavior::FirstSeen);

Internal analysis (best guess; please correct)

Two Lance internals appear to be involved; we have not fully disambiguated which is causal:

1. Indexed-scan join path. create_indexed_scan_joined_stream (src/dataset/write/merge_insert.rs:616+) appears to UNION a BTREE-on-id lookup with a full scan of unindexed fragments. After the first merge_insert rewrites the matched row, the BTREE's entry still points at the now-tombstoned slot in the original fragment, while the rewritten row lives in a new (unindexed) fragment. Both sides of the UNION surface the same stable _rowid (since enable_stable_row_ids = true).
2. Source-dedup HashSet. processed_row_ids: Mutex<HashSet<u64>> (src/dataset/write/merge_insert.rs:2099) reports the spurious duplicate. Setting SourceDedupeBehavior::FirstSeen silences it; setting use_index(false) (which presumably bypasses the UNION entirely) also silences it.

These could be one bug (UNION causes the dedup HashSet to see duplicates) or two interacting bugs. We have not traced the actual data flow.

Possible fixes

• DV-filter the indexed-scan side of the UNION so tombstoned _rowids don't reach the join.
• Default SourceDedupeBehavior::FirstSeen for datasets with stable row IDs enabled — the dedup is then a no-op invariant of the row-id machinery.
• Extend Fragment Reuse Index (FRI) coverage to Operation::Update transactions (it currently tracks compaction Rewrite only — https://lance.org/format/table/index/system/frag_reuse/).

[westonpace]

This should be fixed in 6.0.0. Can you try and reproduce there?

[ragnorc]

Thanks @westonpace - just looked at this again. The bug still reproduces in 6.0.0 (and on main at 7.1.0-beta.4) for the exact configuration in the original repro. Traced the mechanism end to end, validated the fix, and have opened #6965.

The failure mode is more subtle than the original analysis (mine included) and I think it generalizes a bit beyond merge_insert.

Why #6563 doesn't cover this case

#6563 added create_restricted_deletion_mask and routed MapIndexExec through it. That allow-list correctly excludes addresses from fragments outside the index's fragment_bitmap - but only on the non-stable-row-id branch of create_deletion_mask_impl. At rust/lance/src/index/prefilter.rs:270:

} else if dataset.manifest.uses_stable_row_ids() {
    Some(Self::do_create_deletion_mask_row_id(dataset.clone()).boxed())
}

The stable-row-id branch shortcuts straight to do_create_deletion_mask_row_id, which ignores the restrict_to_fragments flag and iterates every fragment in the dataset. So for stable-row-id datasets the fragment-bitmap filter is a no-op - exactly the configuration in this issue.

What the actual mechanism is

My first instinct was the framing from the original analysis - "BTREE points at the now-tombstoned slot in the original fragment, deletion-vector filter doesn't catch it". That's wrong on stable-row-id datasets, and I want to walk through why because it's the part that took the longest to pin down.

To disambiguate "the BTREE u64 is a physical address" vs "the BTREE u64 is a stable row id" I built a variant of the repro where the target row's physical address would be far above any plausible stable row id. Three pre-allocated fragments via append (so the target row "A" lives in frag 2 at offset 0, making its physical address (2<<32)|0 = 8,589,934,592), with only 5 rows total in the dataset (so the largest stable row id is 4). Then instrumented MapIndexExec to log what the BTREE returns and what the mask does:

[probe-mapidx] BTREE pre-filter: AllowList({0: [4]}) | deletion_mask present: true
[probe-mapidx] mask after AND:    AllowList({0: [4]})   ← stale entry passes
[probe] dedup batch: row_ids=[4, 4]
[probe] DUPLICATE row_id=4 already_seen={4}

AllowList({0: [4]}) is RowAddrTreeMap's debug repr - upper-32-bits 0, lower-32-bits 4, so the stored u64 is 4. That's A's stable row id, not A's physical address (which would render as {2: [0]}). So on stable-row-id datasets the BTREE's posting list is keyed by stable row id, which lines up with btree.rs training: TrainingCriteria::new(TrainingOrdering::Values).with_row_id() at rust/lance-index/src/scalar/btree.rs:1916 (and _rowid is the stable row id when enable_stable_row_ids=true). The deletion-mask allow-list is also stable-row-id-keyed via RowAddrTreeMap::from(&RowIdSequence) at rust/lance-table/src/rowids.rs:529, so the AND in MapIndexExec is at least operating in the same namespace - it just isn't filtering enough of it.

The actual sequence in the probed test:

1. First merge_insert UpdateAll rewrites A. A's stable row id (4) is preserved, but its physical home moves to a new fragment (frag 3), which is outside the index's fragment_bitmap.
2. Second merge_insert UpdateAll on the same key. create_indexed_scan_joined_stream builds the standard "indexed scan ∪ unindexed_fragments scan" UNION.
3. Indexed side: BTREE returns stable row id 4. do_create_deletion_mask_row_id builds the allow-list by iterating every fragment - including frag 3. So the allow-list is {0,1,2,3,4} (frag 0 contributes {0,1}, frag 1 contributes {2,3}, frag 2 contributes {} because A's slot was tombstoned via the deletion vector, frag 3 contributes {4} because A's stable row id was preserved on rewrite). mask & {4} = {4} - passes the filter. TakeExec, via RowIdIndex, resolves stable row id 4 to its current address (frag 3) and reads row A.
4. Unindexed side: full scan of unindexed_fragments = [frag 3] also reads row A.
5. Both branches emit _rowid = 4. The processed_row_ids HashSet in merge_insert's source-dedup loop sees 4 twice and aborts.

So the design assumption being violated is the merge_insert UNION's: "indexed-side rows and unindexed-side rows are disjoint." On stable-row-id datasets that's only true if the allow-list is restricted to fragments inside the bitmap. Without the restriction, a row whose stable row id is in the BTREE but whose current physical home is in an unindexed fragment shows up on both sides.

I also empirically validated that this is stable-row-id-specific via a matrix that toggled each factor independently. Flipping enable_stable_row_ids=false (everything else identical to the original repro) removes the bug; flipping the BTREE off also removes it; both documented workarounds (use_index(false), SourceDedupeBehavior::FirstSeen) suppress it as expected.

Considered alternatives

Before settling on the fix below I thought about a few less-good options:

• Default SourceDedupeBehavior::FirstSeen - symptomatic. Masks the duplicate in merge_insert specifically and silently dedups legitimate ambiguous source rows that Fail exists to catch. Doesn't help FTS / vector queries that hit the same allow-list path.
• DV-filter the indexed-scan side of the UNION (the original analysis' proposal) - fixes merge_insert in isolation, leaves the latent bug in every other consumer of do_create_deletion_mask_row_id.
• Drop a fragment from the index's fragment_bitmap on any row-level mutation - aggressive, kills the indexed read on every modified fragment. Not in keeping with how Lance treats indexes as eventually-consistent layers.
• Index stable_row_id with a separate always-current address resolution layer - architecturally the cleanest answer (eliminates the entire "stale index entry" class) but a much larger surgery. Could extend the Fragment Reuse Index to cover Operation::Update as a smaller variant of this.

The fix in #6965

Threads restrict_to_fragments into do_create_deletion_mask_row_id so that on stable-row-id datasets the allow-list contains stable row ids whose current home is inside the restriction - semantically equivalent to the address-keyed allow-list in the non-stable-row-id branch. Three minimal edits:

• prefilter.rs: do_create_deletion_mask_row_id takes restrict_to: Option<RoaringBitmap> and filters dataset.get_fragments() by it.
• prefilter.rs:270: pass Some(fragments) when restrict_to_fragments=true.
• caches.rs: add restrict_hash: Option<u64> to RowAddrMaskKey so two consumers requesting different fragment subsets don't poison each other's cached mask.

Validation

Stashed the fix, re-ran the new regression tests on main - both fail with the documented errors. Restored, both pass. Full prefilter (18/18) and merge_insert (145/145) suites green; cargo fmt + cargo clippy --all --tests --benches -- -D warnings clean.

While writing the FTS-side regression I confirmed the same bug class also affects full_text_search after a merge_insert on stable-row-id datasets (30 unique IDs out of 40 expected, pre-fix). So this PR also closes that previously-unfiled variant - included as test_issue_6877_fts_no_duplicates_stable_row_ids.

PR: #6965.