merge_insert reports zero updated rows for a matching partial upsert

[mmatczuk]

Description

A merge_insert with WhenMatched::UpdateAll reports num_updated_rows: 0 when its source contains one key that is present in the dataset. The expected count is 1.

The failure occurs on storage version 2.1 with stable row IDs disabled, max_rows_per_group: 1, and max_rows_per_file: 8. The dataset's only value column has Arrow Null type. Before the failing merge, the dataset undergoes UpsertIf, UpsertFull, ZoneMap index creation, another UpsertFull, and an append. The failing partial upsert supplies key 166 and the value column, using WhenNotMatched::DoNothing and WhenNotMatchedBySource::Keep. A full-scan oracle verifies the dataset after each preceding operation, including the presence of key 166 before this merge.

The merge implementation assigns UpdateAll when both the source and target contribute a row to the join, and the write path increments num_updated_rows for every UpdateAll action. The root cause of the missing count has not yet been isolated.

It is not yet clear whether this is only an incorrect statistic or a dropped update. Continuing after the statistics mismatch produces no scan divergence, but a Null-typed column has the same value before and after every update, so the data oracle cannot distinguish a successful write from a lost one.

A deterministic 158-byte harness input reproduces the failure, but no standalone public-API reproducer has been found. Simpler cases all report num_updated_rows: 1, including a partial source schema, a Null-typed value column, a scalar index with and without a preceding update, stable row IDs on and off, and small row-group and file limits. The longer mutation history appears to be required.

Expected behavior

num_updated_rows should be 1 when the source key matches one live target row and WhenMatched::UpdateAll is configured. The row should also be updated.

Lance version

v10.1.0-beta.2 (94cca93d13b1c731afd6660d46ad2f8d4b4345c4)

Language binding

Rust

Environment

Fedora Linux 44, x86_64, in-memory object store, rustc 1.97.1

Logs / traceback

thread 'main' panicked at src/dataset_ops.rs:1667:21:
assertion `left == right` failed: step 4, candidate 5: op UpsertPartial { keys: [166], mask: 0, rotate: 255, data: DataParams { len: Big, nulls: None, nans: true } }: updated_rows
  left: 0
 right: 1

[lance-gatefixer]

I could not establish this failure on current main at 86d3bd1d50ea12f250ea20c1e95dcdbab454320c. The repository does not contain the referenced dataset_ops harness or its 158-byte input, and the closest in-tree partial-schema/all-null merge regression passes for storage versions 2.0, 2.1, and 2.2 (cargo test -p lance test_repro_3515_partial_schema_fully_indexed -- --test-threads=1). The current planner also does not select a ZoneMap for merge-insert exact-equality lookup. Please attach the exact 158-byte input together with the harness revision and replay command (or an equivalent standalone reproducer) so the full mutation history can be replayed and the missing target row isolated.

[lance-gatefixer]

I still could not establish this failure on current main at 86d3bd1d50ea12f250ea20c1e95dcdbab454320c. The closest in-tree partial-schema/all-null regression passes for storage versions 2.0, 2.1, and 2.2, and a reconstructed storage-2.1 replay plus 100 deterministic variants of the reported mutation shape all reported one updated row. The reporter's public fork contains neither the referenced dataset_ops harness nor its 158-byte input. Please attach that exact input together with the harness commit and replay command (or an equivalent standalone reproducer) so the specific failing history can be isolated.

[mmatczuk]

This reproduces independently on v11.0.0-beta.4 (d20ca47) without the earlier 158-byte session. It is a dropped update, not only a bad statistic: after the merge, the target column remains NULL, and num_updated_rows is correctly zero because no replacement batch was emitted.

Steps to reproduce

Both preconditions are required:

1. A scalar index on the join key (id). This makes MergeInsertJob::can_use_create_plan return
   false, which routes the merge onto the legacy Merger + create_indexed_scan_joined_stream
   path instead of the v2 fast path.
2. A source batch whose column order differs from the dataset's column order. Values are
   identical either way — only the field order differs.

With either precondition removed, the same merge counts and applies correctly.

Root cause

rust/lance/src/dataset/write/merge_insert.rs, Merger::extract_selections (~line 2992) locates
the source key columns by name, then assumes the target half of the joined batch repeats them
at the same relative offset:

let target_key_cols = source_key_cols.iter().map(|c| c + right_offset).collect::<Vec<_>>();
...
let in_right = Self::not_all_null(combined_batch, &target_key_cols)?;

That holds on the full-table-scan path, whose target projection is built in source column
order (create_full_table_joined_stream, ~line 1259). It does not hold on
create_indexed_scan_joined_stream (~line 1043), which projects the target in dataset-schema
order via self.dataset.empty_projection().union_arrow_schema(schema).

So for dataset (id, a, b) and a source ordered (a, id), the joined batch is:

 0: a          1: id                        <- source half
 2: target_id  3: target_a                  <- target half (dataset order), right_offset = 2
 4: _rowid     5: _rowaddr

index_of("id") is 1, so the target key column is taken to be 1 + 2 = 3, i.e. target_a —
not target_id. If target_a happens to be NULL for the matched row, in_right is false,
in_both ends up empty, and the num_updated_rows counter at line 3143 is never incremented —
and no replacement batch is emitted for that row at all.

Suggested fix: resolve the target key columns by name (target_{key}, the same naming
prefix_columns_phys(target, "target_") already produces) instead of source_idx + right_offset,
or make the indexed-scan path project the target in source order the way the full-scan path does.
Using the join's own presence sentinel, the way the v2 assign_action.rs path does, would be more
robust still.

The complete reproducer uses only the public Rust API.

Create Cargo.toml:

[package]
name = "merge-insert-partial-undercount-repro"
version = "0.0.0"
edition = "2021"

[dependencies]
arrow = "58"
futures = "0.3"
tempfile = "3"
tokio = { version = "1.23", features = ["rt", "macros"] }
lance = { git = "https://github.com/lance-format/lance", tag = "v11.0.0-beta.4", default-features = false }
lance-index = { git = "https://github.com/lance-format/lance", tag = "v11.0.0-beta.4", default-features = false }

Create src/main.rs:

use std::sync::Arc;

use arrow::array::{
    Array, AsArray as _, Int32Array, RecordBatch, RecordBatchIterator, RecordBatchReader,
    UInt64Array,
};
use arrow::datatypes::{DataType, Field, Schema};
use futures::TryStreamExt as _;
use lance::Dataset;
use lance::dataset::{
    MergeInsertBuilder, WhenMatched, WhenNotMatched, WhenNotMatchedBySource, WriteParams,
};
use lance::index::DatasetIndexExt as _;
use lance_index::IndexType;
use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams};

fn reader(batch: RecordBatch) -> Box<dyn RecordBatchReader + Send> {
    let schema = batch.schema();
    Box::new(RecordBatchIterator::new([Ok(batch)], schema))
}

/// Fresh dataset holding one row `(id=1, a=NULL, b=7)`, optionally indexed on `id`, then one
/// partial-column merge_insert that sets `a = 42` for `id = 1`. `id_first` only changes the
/// *column order* of the source batch. Returns `(num_updated_rows, value of a after the merge)`.
async fn trial(dir: &std::path::Path, index_on_id: bool, id_first: bool) -> (u64, Option<i32>) {
    let target_schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::UInt64, false),
        Field::new("a", DataType::Int32, true),
        Field::new("b", DataType::Int32, true),
    ]));
    let target = RecordBatch::try_new(
        target_schema,
        vec![
            Arc::new(UInt64Array::from(vec![1u64])),
            Arc::new(Int32Array::from(vec![None::<i32>])),
            Arc::new(Int32Array::from(vec![Some(7)])),
        ],
    )
    .unwrap();

    let uri = dir
        .join(format!("ds-{index_on_id}-{id_first}"))
        .to_str()
        .unwrap()
        .to_owned();
    let mut dataset = Dataset::write(reader(target), &uri, Some(WriteParams::default()))
        .await
        .unwrap();

    if index_on_id {
        dataset
            .create_index(
                &["id"],
                IndexType::Bitmap,
                Some("idx_id".to_owned()),
                &ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap),
                true,
            )
            .await
            .unwrap();
    }

    // Partial-column source: `id` plus `a` only — a strict subschema of the dataset.
    let id = Field::new("id", DataType::UInt64, false);
    let a = Field::new("a", DataType::Int32, true);
    let ids = Arc::new(UInt64Array::from(vec![1u64])) as Arc<dyn Array>;
    let vals = Arc::new(Int32Array::from(vec![Some(42)])) as Arc<dyn Array>;
    let source = if id_first {
        RecordBatch::try_new(Arc::new(Schema::new(vec![id, a])), vec![ids, vals])
    } else {
        RecordBatch::try_new(Arc::new(Schema::new(vec![a, id])), vec![vals, ids])
    }
    .unwrap();

    let mut builder =
        MergeInsertBuilder::try_new(Arc::new(dataset.clone()), vec!["id".into()]).unwrap();
    builder
        .when_matched(WhenMatched::UpdateAll)
        .when_not_matched(WhenNotMatched::DoNothing)
        .when_not_matched_by_source(WhenNotMatchedBySource::Keep);
    let job = builder.try_build().unwrap();
    let (dataset, stats) = job.execute_reader(reader(source)).await.unwrap();

    let batches: Vec<RecordBatch> = dataset
        .scan()
        .project(&["a"])
        .unwrap()
        .try_into_stream()
        .await
        .unwrap()
        .try_collect()
        .await
        .unwrap();
    let col = batches[0]
        .column(0)
        .as_primitive::<arrow::datatypes::Int32Type>();
    let a_after = (!col.is_null(0)).then(|| col.value(0));

    (stats.num_updated_rows, a_after)
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
    let dir = tempfile::TempDir::new().unwrap();

    // Controls: correct with either precondition missing.
    for (index_on_id, id_first) in [(false, false), (false, true), (true, true)] {
        let (updated, a) = trial(dir.path(), index_on_id, id_first).await;
        println!("control: index_on_id={index_on_id} id_first={id_first} -> num_updated_rows={updated} a={a:?}");
        assert_eq!((updated, a), (1, Some(42)));
    }

    // The bug: index on the join key + source column order != dataset column order.
    let (updated, a) = trial(dir.path(), true, false).await;
    println!("  BUG: index_on_id=true  id_first=false -> num_updated_rows={updated} a={a:?}");
    assert_eq!((updated, a), (1, Some(42)), "update was undercounted and dropped");
}

Run:

cargo run

Expected behavior

The matched row must be updated, regardless of source-field order. The final assertion must observe num_updated_rows = 1 and a = Some(42).

Lance version

v11.0.0-beta.4 (d20ca47)

Language binding

Rust

Environment

Linux x86_64, local filesystem

Logs / traceback

control: index_on_id=false id_first=false -> num_updated_rows=1 a=Some(42)
control: index_on_id=false id_first=true  -> num_updated_rows=1 a=Some(42)
control: index_on_id=true  id_first=true  -> num_updated_rows=1 a=Some(42)
  BUG: index_on_id=true  id_first=false  -> num_updated_rows=0 a=None
thread 'main' panicked at src/main.rs:...: update was undercounted and dropped
  left: (0, None)
 right: (1, Some(42))

A related (but distinct) symptom of the same root cause

The same misaligned target_key_cols offset also affects full-schema upserts
(WhenMatched::UpdateAll + WhenNotMatched::InsertAll), reachable with the identical pair of
preconditions. There, the outcome is different and less severe: the not-matched branch still
carries _rowaddr in its projection, and a full-schema source that's merely reordered gets
classified as a subschema update (is_full_schema at ~line 2372 doesn't set
ignore_field_order, unlike can_use_create_plan), so the row is patched in place via
update_fragments anyway. Net effect: data ends up correct, but the merge reports
num_updated_rows = 0, num_inserted_rows = 1 for a key that already existed — a stats-only
version of this same bug, not a second data-correctness issue. It's a different symptom from
#8282 (a debug_assert! row-address-ordering panic in the same file) — that one is a separate
defect, not the same root cause as far as I can tell.

Happy to open that as its own report if preferred, but since it's the same line of code and the
same two preconditions, a single fix to extract_selections should resolve both at once, so I'm
noting it here instead of splitting the report.