feat(merge_insert): probe scalar indices on the DataFusion plan path

[wjones127]

Stacked on #8052 — its benchmark commit shows up in this diff until that one merges. Review only feat(merge_insert): probe scalar indices on the DataFusion plan path.

merge_insert has two execution paths. The DataFusion path reaches matched target rows by scanning the whole table and hash-joining it against the source. The legacy path can instead probe a scalar index on the merge keys, reading only the pages that could contain a source key — but it is the only path that can, so any merge that needs the DataFusion path scans the full target even when an index would have found the same rows for a fraction of the IO.

This PR teaches the DataFusion path to probe. The target scan becomes:

Union
├─ FilteredReadExec              read the projected columns of probed rows
│  └─ CoalescePartitionsExec
│     └─ MapIndexExec            AND of one IsIn probe per indexed key
│        └─ <source scan>        key columns only
└─ FilteredReadExec              fragments no index covers

It is a new TableProvider for the target rather than a change to the plan builder, and it reports the same schema as the full-scan provider it replaces, so the join, the action expression, and the write sink are untouched.

Correctness rests on the probe being allowed to over-match but never under-match, since the join is what applies the real key predicate. Two consequences: only the indexed merge keys need probing, so a composite key with one indexed column still prunes by that column; and rows in fragments no chosen index covers would be invisible to the probe, so those fragments are scanned and unioned in.

The probe replaces the scan when every one of these holds:

• the source can be scanned more than once, because the probe reads its key columns separately from the join
• at least one merge key has a scalar index supporting exact equality
• when_not_matched_by_source is Keep — the other modes need target rows the source never mentions
• the dataset uses the v2 storage format
• use_index was not turned off

Which merge configurations that adds up to is deliberately narrow: this PR does not change which path a merge takes, so it only affects merges that already used the DataFusion path and have an indexed key. The main one is a composite key where only some columns are indexed. Routing indexed keys away from the legacy path is a separate change, and it needs the benchmark comparison to justify it.

Example

A merge on ["a", "b"] where only a is indexed, previously a full scan of the target:

HashJoinExec: mode=CollectLeft, join_type=Left, on=[(a@0, a@0), (b@1, b@1)], projection=[...]
  ProjectionExec: expr=[a@0 as a, b@1 as b, value@2 as value, true as __merge_source_sentinel]
    DataSourceExec: partitions=1, partition_sizes=[1]
  RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1
    ProjectionExec: expr=[a@1 as a, b@2 as b, _rowid@0 as _rowid, _rowaddr@3 as _rowaddr]
      LanceRead: uri=data, projection=[a, b, _rowaddr], source=stream(_rowid)
        IndexedLookup
          DataSourceExec: partitions=1, partition_sizes=[1]

Note projection=[a, b, _rowaddr]: a full-schema upsert builds its output rows entirely from the source, so the target read fetches only the join keys and the row address needed to supersede the old rows. value is never read. There is a test pinning this, because widening it would quietly turn a cheap keyed read into a full-row one.

Behavior changes

explain_plan now stands the source up as an in-memory table instead of a single-use stream. Without that it could never show the probe, since a single-use source rules it out — and it also means the plan shown is the one the common entry points build. Because the source now reports exact statistics, DataFusion collects it as the join's build side, so the source appears above the target in the plan and the join type is the swapped form of what it was (Left where it read Right). Same join, sides exchanged.

analyze_plan now builds its source provider the same way execution does, so the plan it analyzes is the plan that would have run. For a stream source with retries enabled that means the source is buffered (in memory, then spilled) rather than consumed once.

Not included

• Routing indexed merge keys away from the legacy path. That change is small but needs benchmark evidence, and for partial-schema sources it has to wait for the DataFusion path to gain a column-patch write sink — otherwise it would trade a full target scan for a full row rewrite.
• Emitting (key, row_addr) pairs from the probe, which would let a plain upsert skip the target read entirely. The index API returns a row-address mask today, so the key that matched is not recoverable without reading it back.
• Any source-size condition on the routing. The break-even depends on key clustering and index cache residency, not on a row-count ratio, so it belongs with a real cost model rather than a guessed constant.

[codecov]

Codecov Report

❌ Patch coverage is 89.88235% with 43 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
rust/lance/src/dataset/write/merge_insert/probe.rs	81.98%	18 Missing and 11 partials ⚠️
rust/lance/src/dataset/write/merge_insert.rs	94.57%	2 Missing and 12 partials ⚠️

📢 Thoughts on this report? Let us know!

[wjones127]

Benchmark results, before and after

Ran the ci_benchmarks merge_insert suite on both sides: base is b198e39ec (the benchmark PR, #8052) and probe is 637744987. 61 cases per side, ~14 minutes each, on an M-series laptop against the full-scale 10M-row merge_insert_narrow target.

Only the narrow-target cases ran. The wide, frags, deleted and unindexed_tail targets were skipped because no case in them can reach the probe, so all they would measure is run-to-run variance. That includes test_upsert_unindexed_tail, which is the one benchmark whose shape matches the probe's unindexed-fragment union — its v1_indexed variant routes to the legacy path and its v2_hash variant turns the index off, so the union branch stays covered only by unit tests.

The run found two correctness bugs

Both in test_composite_key_partially_indexed, and both caught by that benchmark's own row-count assertion rather than by timing.

The probe missed most of the source. A materialized source is spread round-robin across partitions, and MapIndexExec probes only the partition it is asked for. The CoalescePartitionsExec in front of it declared no distribution requirement, so the optimizer removed it as unnecessary and the keys in every partition but the first were never probed. A 10K-row source updated 8,192 rows and inserted 1,808 duplicates.

Over-matching probes named the same target row twice. Probes are evaluated per source batch, so two batches can reach the same target row. That row was then read twice, which looks like two candidate matches for one source row, and the merge failed with Ambiguous merge inserts are prohibited.

Fixed in 637744987: MapIndexExec now requires a single input partition, and a DistinctRowAddrs node de-duplicates candidate row addresses before the take. The regression test gives every source batch a distinct key, which is what makes it fail on both counts without the fix — the previous test reused one key across batches and passed through both bugs.

The second bug is pre-existing in the legacy indexed path, which has the same per-batch probe with no cross-batch de-duplication. Filed as #8057; fixing it there changes legacy behavior and is out of scope here.

Affected case

test_composite_key_partially_indexed is the only configuration in the suite that reaches the probe. The suite's use_index knob is True (which routes a fully-indexed key to the legacy path) or False (which disables the probe as well), so a partially-indexed composite key is the only way into it.

benchmark	base	probe	delta
test_composite_key_partially_indexed	1.15 s	93.8 ms	-91.8%

Same merge measured for IO, 10K-row source, warm cache:

metric	base	probe
time	1.15 s	94.3 ms
read iops	21	5
read bytes	53.5 MiB	31.6 KiB
write bytes	441.4 KiB	452.3 KiB

Everything else is unchanged

46 timing cases, none of which reach the probe: median absolute delta 1.7%, 44 within ±10%, worst +18.1% on a 5-second case. All 14 IO cases report byte-identical read and write volumes on both sides, which is the reliable signal here — the timing spread is laptop variance, not a code difference.

The probe is not always faster

This is the result worth acting on. Same dataset, same 10K-row source, same merge shape — the only difference is which column carries the index:

probed column	distinct values	probe	scan
composite_a	10M (unique)	96 ms	1.16 s
composite_b	1024	12.3 s	1.21 s

composite_b is row_index % 1024, so a 10K-row source names all 1024 of its values and b IsIn [...] reaches every row in the target. Read volume is identical to the scan at 38.8 MiB, so the 10x is entirely access pattern: the probe takes 10M row addresses through keyed reads where the scan reads the same bytes sequentially.

Both of those are partially-indexed composite keys, so both take the probe as this PR stands. test_composite_key_partially_indexed uses the unique column, which is why the suite only shows the favourable side.

Context: the same merge across paths

10K-row source against the 10M-row target, warm cache.

configuration	path	time	read
composite key, a indexed only	v2 + probe	94 ms	31.6 KiB
composite key, both indexed	v1 probe	5.99 s	439.9 KiB
composite key, both indexed	v2 scan	1.17 s	38.8 MiB
single indexed key	v1 probe	89 ms	439.9 KiB
single indexed key	v2 scan	870 ms	26.8 MiB

Two things fall out of this. For a selective key the probe is roughly 10x faster than the scan, which is the case for routing indexed merges onto this path. But v1's 5.99 s on the fully-indexed composite is the low-cardinality effect again — it ANDs composite_b into the probe, and paying for that index costs more than the pruning saves even though the AND makes the result set no smaller. So the routing flip would regress that shape from 1.17 s to around 6 s unless the probe first learns to choose which indexed columns to probe, rather than using all of them.