feat: add native support for MergeRowsExec (row-level MERGE INTO)

[unikdahal]

Which issue does this PR close?

Part of #5122.

Rationale for this change

Row-level MERGE plans currently fall back to Spark entirely because MergeRowsExec has no
native equivalent, even though the expensive parts of the plan (the source/target join, filters,
projections) are operators Comet already accelerates.

MergeRowsExec itself is close to a projection: it routes each joined row through matched /
not-matched / not-matched-by-source instruction lists that are ordinary Catalyst expressions,
emitting kept, inserted, or discarded rows.

Adding a native implementation allows the join and merge dispatch logic to remain in a native
Comet stage instead of falling back to Spark for the whole statement.

This is independent of native writes: MergeRowsExec can remain native while the final Iceberg
write is still performed by the JVM writer, and can also be combined with Comet's native Iceberg
writer when enabled.

What changes are included in this PR?

• New native MergeRowsExec operator
  (native/core/src/execution/operators/merge_rows.rs) reproducing Spark's dispatch semantics:
  • matched / not-matched / not-matched-by-source instruction groups
  • Keep / Discard / Split semantics (0/1/2 output projections)
  • first-match-wins clause evaluation
  • Spark's NULL-collapses-to-false predicate semantics
• Cardinality checking (MERGE_CARDINALITY_VIOLATION), mirroring Spark's
  BitmapCardinalityValidator, surfaced as a real SparkRuntimeException through the structured
  SparkError / JNI error-conversion path rather than as a generic native exception.
• New MergeRows / MergeInstruction / MergeOutputRow protobuf messages.
• CometOperatorSerde support for MergeRowsExec, version-gated for supported Spark versions
  (the operator does not exist before Spark 3.5).
• New spark.comet.exec.mergeRows.enabled configuration, disabled by default
  (experimental / opt-in).
• output_rows / output_batches / elapsed_compute metrics.
• Documentation updates covering operator support and compatibility/caveats.

Spark 4.x's eight per-clause row counters
(numTargetRowsCopied, numTargetRowsInserted, numTargetRowsUpdated,
numTargetRowsDeleted, etc.) are not currently exposed by the native operator.
This is documented as a known compatibility gap; Spark 3.5.x does not expose these
metrics either.

Performance

I benchmarked an Iceberg copy-on-write mixed-CDC MERGE workload across three execution paths:

Execution path	Best time (ms)	Avg time (ms)	Stdev (ms)	Rate (M rows/s)	Per row (ns)	Relative
Spark (Comet disabled)	1632	1687	50	0.6	1813.8	1.0x
Comet MergeRowsExec + JVM Iceberg writer	1241	1257	14	0.7	1378.7	1.3x
Comet MergeRowsExec + native Iceberg writer	1018	1054	38	0.9	1131.2	1.6x

Environment

• OpenJDK 64-Bit Server VM 17.0.20+8-LTS
• Linux 6.17.0-1022-azure
• AMD EPYC 7763 64-Core Processor

The JVM-writer comparison isolates the benefit of native MergeRowsExec: with the final
Iceberg write still performed by Spark's JVM writer, keeping merge processing in Comet improves
the benchmark from 1632 ms to 1241 ms best time, approximately 1.3x relative to the Spark
baseline.

When MergeRowsExec is combined with Comet's native Iceberg writer, the workload improves from
1632 ms to 1018 ms best time, approximately 1.6x relative to the Spark baseline.

The benchmark harness was used for performance validation only and is not included as part of
this PR.

How are these changes tested?

• Rust unit tests in merge_rows.rs covering:

  • matched / not-matched / not-matched-by-source dispatch
  • first-match-wins ordering
  • NULL handling
  • cardinality violations, including violations split across batches and NULL row IDs
  • Split output-row values
  • out-of-range row-ID ordinal rejection

• CometMergeRowsSuite (Scala), using Spark's
  InMemoryRowLevelOperationTableCatalog so the core MergeRowsExec contract is tested
  independently of Iceberg or another connector:

  • asserts native Comet engagement
  • verifies exact Comet/Spark result parity
  • verifies cardinality violations surface as
    SparkRuntimeException[MERGE_CARDINALITY_VIOLATION], matching Spark's exception type
    and error class

• End-to-end validation against Iceberg covering:

  • MATCHED
  • NOT MATCHED
  • NOT MATCHED BY SOURCE
  • update / delete / insert
  • multiple clauses and first-match-wins behavior
  • NULL conditions
  • cardinality violations

The end-to-end validation confirms native engagement and result/error parity with Spark, while
the benchmark above provides performance evidence both with Spark's JVM Iceberg writer and with
Comet's native Iceberg writer.

[unikdahal]

@peterxcli @andygrove @jordepic ready for review whenever you have time.

[jordepic]

Thanks for doing this @unikdahal ! I won't be the best source of truth since this is just the default merge operator as opposed to the iceberg stuff. My hope is that once this is in supporting the MERGE INTO operator in iceberg is pretty trivial.

I think @andygrove will need to provide some more feedback on how much the merge operator here can deviate from default spark functionality, I know the intent of comet is to be perfectly identical, but can deviate with some opt-in toggles.

[sunchao]

[P2] Evaluate later predicates only for the remaining rows

Could you narrow to remaining before evaluating each clause condition? The mask is currently applied after instr.condition.evaluate(&group_batch), so rows handled by an earlier clause still evaluate later predicates whenever another row remains. With ANSI enabled, two matched rows having s.d = 0 and 2, and clauses WHEN MATCHED AND s.d = 0 THEN UPDATE SET v = 111 followed by WHEN MATCHED AND 2 / s.d > 0 THEN UPDATE SET v = 222, the second predicate divides by the already-handled row's zero. I ran this MERGE on Spark 4.0.4 and got [1,111], [2,222]. A two-row probe of this exact run_group fails with DivideByZero, while the same function succeeds row-at-a-time. Spark's applyInstructions returns immediately after the first matching clause, so this is a failure of an otherwise valid MERGE, not just a different error ordering.

[unikdahal]

Fixed. run_group now physically shrinks the working batch to unclaimed rows after each instruction instead of masking a stable batch -- a claimed row is gone before the next clause's condition ever evaluates, not just masked out after. Added a regression test with your exact divide-by-zero shape.

[sunchao]

[P2] Normalize nested output types before stamping the batch

Could this use cast_and_stamp_schema, as ExpandStream::expand does? ExpandExec::build_schema widens nested nullability across the instructions, but this direct stamp does not reconcile each projected array with that widened type. For a target column payload STRUCT<n: BOOLEAN>, UPDATE SET payload = named_struct('n', s.d IS NULL) produces a non-nullable inner n, while Spark's appended carryover Keep references the target's nullable n. The derived schema is therefore nullable, and RecordBatch::try_new_with_options rejects the update array with column types must match schema types. I reproduced that error using the exact-head CreateNamedStruct, schema builder, and run_group. The corresponding Spark 4.0.4 MERGE succeeds. Normalizing the array before stamping also succeeds and preserves the value.

[unikdahal]

Fixed. project now uses cast_and_stamp_schema, same as ExpandStream::expand, to reconcile each projected column's nested nullability against the declared schema. Added a regression test with a non-nullable projected struct field against a nullable declared one.

[sunchao]

Both existing [P2] findings still apply to this head and need to be addressed before merge: later predicates evaluate already-handled rows and nested output types are not normalized before stamping. This approval does not mark either finding resolved.

The focused exact-head native probes reproduced both failures while the corresponding Spark 4.0.4 MERGE statements succeeded. GitHub CI is still marked action_required, so there is no successful full CI result for this head.

Withdrawing this approval because I approved the wrong PR in error. The existing P1/P2 findings remain unresolved and still need to be addressed before merge.

[sunchao]

Two follow-up P2 findings on 7f378fb6. All 15 existing native MergeRows tests passed. The subquery finding was checked with a successful Spark 4.1.3 baseline and an expression-discovery probe compiled from the new class against cached Comet dependencies; a full Comet JVM run was blocked by Maven dependency resolution. The memory finding was reproduced using this revision's exact check_cardinality function and a 16 MiB memory pool.

[sunchao]

[P2] Retain MERGE assignment subqueries

Could you preserve the instruction expressions as fields on this node, or fall back for scalar-subquery assignments? An assignment such as WHEN MATCHED THEN UPDATE SET amount = (SELECT max(amount) FROM source) is accepted by the scalar-subquery serializer, but Spark's expression traversal does not recurse into originalPlan. Consequently, CometNativeExec.prepareSubqueries and collectSubqueries cannot discover or register that subquery, leaving the native lookup without an entry (Subquery ... not found for plan ...). I verified that the corresponding Spark 4.1.3 MERGE succeeds with [1,3]; a probe compiled from this new class found one scalar subquery on the original MergeRowsExec, zero on CometMergeRowsExec, and zero returned by collectSubqueries.

[sunchao]

[P2] Account for allocated hash-table capacity

The fixed 16-byte charge per ID is not conservative immediately after HashSet growth. Calling this exact check_cardinality function with 917,505 distinct IDs in 4,096-row batches succeeded against a 16 MiB GreedyMemoryPool, while an allocator probe measured 18,874,384 bytes of live table allocation versus only 14,680,080 bytes reserved. Thus an 18 MiB persistent table is accepted by a 16 MiB pool. This uncharged capacity grows with partition size, so the overshoot is not bounded by the current batch and can bypass the memory budget, risking executor OOM. Could you account for allocated table capacity, for example using DataFusion's estimate_memory_size, instead of only the number of inserted IDs?

[unikdahal]

Thanks @sunchao for the detailed review. I rebased this PR onto the latest main and addressed the remaining review findings.

The current head includes the following fixes:

• Later predicates evaluating already-handled rows: run_group now physically narrows the working batch after each matched instruction, so rows claimed by an earlier clause are never evaluated by later predicates. The ANSI divide-by-zero case from the review is covered by a regression test.

• Nested output nullability: projected output now goes through cast_and_stamp_schema, matching the approach used by ExpandExec, so nested nullability/type widening is normalized before constructing the output batch. The reproduced nested-struct case is covered by a regression test.

• MERGE assignment subqueries: CometMergeRowsExec now retains the source/target predicates and the three instruction groups as actual SparkPlan expression fields instead of hiding them under originalPlan. This lets Catalyst/Comet discover and prepare scalar subqueries correctly while also preserving MATCHED / NOT MATCHED / NOT MATCHED BY SOURCE boundaries for plan equality and canonicalization. Added an end-to-end scalar-subquery assignment test and semantic equality/canonicalization coverage.

• Cardinality hash-table memory accounting: cardinality state now uses DataFusion's hash table implementation and reserves memory with the DataFusion memory pool before growing the table. Accounting includes hash-table capacity/allocation overhead rather than charging a fixed number of bytes per inserted ID, and reservation is rolled back if allocation fails. The exact 917,505-ID / 16 MiB reproduction from the review is covered, along with rehash-boundary and actual-allocation accounting tests.

I also added/strengthened coverage for native WriteDelta execution and Split update-as-delete+reinsert semantics, required native engagement in the cardinality-error test, validated the row-ID type/ordinal up front, and kept MERGE_CARDINALITY_VIOLATION mapped to Spark's structured runtime error.

Would appreciate another look when you have time. Thanks again for the precise reproductions, they helped tighten this substantially.

[andygrove]

Thanks for this, and thanks @sunchao for the earlier rounds. I went through the current head (e8af41e48) and all four of the previous P2 findings look genuinely addressed: run_group physically shrinks the working batch between instructions, project goes through cast_and_stamp_schema, the instruction groups are real case-class fields so prepareSubqueries can find scalar subqueries, and the cardinality set reserves against estimate_memory_size at rehash boundaries. I also checked the null-row-id handling against UnsafeRow.setNullAt semantics and it is correct, and confirmed the MergeRows.Instruction API (condition, outputs, ROW_ID) is identical in 3.5.9 and 4.1.0. Nice piece of work.

Relationship to #5122

The scoping is right. #5122 covers four operators (MergeRowsExec, ReplaceDataExec, WriteDeltaExec, and InsertOnlyMergeExec on 4.2+), and this takes only the first, independent of native writes, which is exactly what the issue proposed.

Two things worth naming:

• Accelerate row-level MERGE / UPDATE / DELETE plans (MergeRowsExec, ReplaceDataExec, WriteDeltaExec) #5122's motivating workload is Iceberg CDC upsert. Testing against Spark's own InMemoryRowLevelOperationTableCatalog is a good call for pinning the Spark-core contract, but it leaves the actual motivating case untested with the flag on. See point 7 below.
• Accelerate row-level MERGE / UPDATE / DELETE plans (MergeRowsExec, ReplaceDataExec, WriteDeltaExec) #5122 notes that full acceleration needs Accelerate DataSource V2 writes (AppendDataExec, OverwriteByExpressionExec, OverwritePartitionsDynamicExec) #5121 (DSv2 writes). Until that lands the writer stays on the JVM, so this moves the ColumnarToRow boundary up one node rather than removing it. That is the crux of point 5.

InsertOnlyMergeExec is correctly out of scope here, but on Spark 4.2 an insert-only MERGE plans that instead of MergeRowsExec, so this operator will silently not engage. A follow-up issue under #5122 would be better than expanding this PR.

@peterxcli asked to be tagged for review on the issue thread.

Findings

1. Rebase and get a green CI run. The branch conflicts with main in two places. operator.proto field 120 is now IcebergWrite iceberg_write = 120, so MergeRows needs to move to 121, and operators/mod.rs conflicts too. main also moved to DataFusion 55.0 and Arrow 59.2 in 75fdddc92. gh pr checks currently reports no checks at all on this head. A 2400-line change touching the planner, the proto, error mapping and three Spark profiles really needs a full build before it can be evaluated.

2. CometMergeRows.scala is byte-for-byte identical in spark-3.5/ and spark-4.x/. I diffed them and the only difference is the path. Same for ShimCometMergeRows.scala and the three-line ShimSparkErrorConverter case. That is 174 lines of real fallback logic that has to be kept in sync by hand across four Spark profiles. ShimCometWindowGroupLimit sets a precedent for duplication, but that is a 20-line shim, not a serde. The spark-4.1+ source root with the spark-none placeholder is the existing pattern for a shared "3.5 and later" root. Could the serde move there, leaving only the thin class-registration shim per version?

3. getSupportLevel returns Compatible(None) while three divergences are documented. The one that stands out is metrics. On Spark 4.1+, MergeRowsExec publishes eight per-clause counters (numTargetRowsCopied, numTargetRowsInserted, numTargetRowsUpdated, numTargetRowsDeleted, and the four matched / not-matched-by-source breakdowns). Enabling this operator replaces all of them with generic Comet metrics, and 4.1 is the default build profile. Should this be Incompatible(Some(...)) with a matching getIncompatibleReasons() so it surfaces on the compat page? Implementing the counters may not be much work either: the proto already has a per-instruction message and run_group already computes fire.true_count() per instruction, so a Context enum on MergeInstruction plus per-instruction counts would cover it. Also, the doc says "Spark 4.x's MergeRowsExec exposes eight metrics" but 4.0.1 has none. They arrived in 4.1 alongside the Context field on Keep.

4. Output row order diverges from Spark and is not on the compat page. run_group emits rows grouped by the instruction that produced them, and process_batch concatenates matched, then not-matched, then not-matched-by-source. Spark emits in input row order. The Rust doc comment argues this is safe because DistributionAndOrderingUtils places the required repartition and sort above MergeRows, which holds for a connector that declares one. A V2 table that declares neither ends up with a different physical row order in the written files than Spark produces. That belongs on compatibility/operators.md alongside the other two caveats, since it is what a user would actually notice from an unordered SELECT *.

5. Could you add some benchmark numbers? The rationale is that MERGE falls back today, but with the writer still on the JVM this moves the ColumnarToRow up one node rather than removing it, and process_batch adds a concat_batches full copy per input batch on top of a filter_record_batch per instruction. A before and after on a realistic CDC-shaped MERGE would tell us whether the operator earns the flag now, or whether the win really arrives with #5121.

6. Test gaps. WHEN NOT MATCHED BY SOURCE never executes end to end. It appears only in the equality and canonicalization test, which builds the plan but never runs the query, so not_matched_by_source_instructions is covered by Rust unit tests only. Copy-on-write WHEN MATCHED THEN DELETE is also uncovered at the SQL level, only the delta variant is tested, and multi-clause first-match-wins with conditions is likewise Rust-only. It would also be worth one test that the default (flag off) falls back cleanly and matches Spark, since that is the path every user is on today. Separately, the four withSQLConf(CometConf.COMET_ENABLED.key -> "true", ...) calls can drop the COMET_ENABLED entry, CometTestBase already enables Comet by default.

7. Iceberg MERGE coverage. The suite scaladoc points at CometIcebergWriteActionSuite for Iceberg coverage, but nothing there enables mergeRows, so with the flag on there is no Iceberg MERGE test anywhere. Given that is the workload #5122 is about, and given how the native Iceberg writer behaves differently depending on whether its child converted natively, one Iceberg MERGE test with the flag on would be worth adding.

8. Why is the output schema derived from the projections rather than output_types? merge.output_types is serialized on the wire but only read in the branch where every instruction is a Discard, which effectively never happens since RewriteMergeIntoTable appends a catch-all Keep. The live path goes through ExpandExec::build_schema. Since project already funnels every column through cast_and_stamp_schema, always building the schema from Spark's declared output_types would make the native output schema provably equal to what the JVM expects across FFI, and would delete the two-branch derivation. If the derived types are preferred in order to avoid a lossy safe cast on, say, a decimal, that reasoning is worth a comment. Otherwise the field is dead weight.

9. Comment density in merge_rows.rs. A good chunk of the 1320 lines is prose, and some of it is genuinely valuable, the UnsafeRow.setNullAt research and the cardinality-versus-instruction ordering caveat especially. But the twelve-line comment on the baseline field, the twelve lines on SEEN_HASH_TABLE_SLACK_BYTES, and several blocks that pre-argue with a hypothetical reviewer put this well above the density of the neighbouring operators, and that kind of prose rots first. A trim pass down to the non-obvious facts would help.

10. ShimCometMergeRows creates a new org.apache.comet.rules.shims package. Every sibling shim, including ShimCometWindowGroupLimit, ShimCometStreaming and ShimSubqueryBroadcast, lives in org.apache.comet.shims. Was the new package deliberate?

[unikdahal]

Thanks @andygrove for the detailed review. I’ve pushed another update addressing the feedback and also updated the PR description with the latest benchmark results, compatibility boundaries, and test coverage.

A few of the review points led to more substantial changes than I initially expected, so here is a summary of what changed and the reasoning behind the remaining decisions.

• Spark 4.1+ native MergeRows is now intentionally disabled.
  While looking into the missing per-action metrics, I found that this is more than a metrics-parity issue. Starting in Spark 4.1, the concrete MergeRowsExec contributes the action counters used to construct MergeSummary, which is then passed through the summary-aware V2 BatchWrite.commit path. Replacing MergeRowsExec without preserving that state could therefore alter the writer commit contract even if the produced rows are correct.

  For this PR, native MergeRowsExec is consequently limited to Spark 3.5.x and 4.0.x, while 4.1+ explicitly falls back to Spark's implementation. There is test coverage for that version boundary. Full 4.1+ support should preserve the metrics, action context, MergeSummary, and writer contract end-to-end rather than just exposing additional counters, so I’m keeping that out of this already-large change.

• output_types are now authoritative for the native output schema.
  Previously the native side could derive a structurally compatible schema independently. That is risky for nested types/nullability because Spark's serialized schema is the actual execution contract. The updated implementation validates the independently derived projection schema against Spark's output_types and uses the Spark declaration for the emitted schema. This keeps the native operator aligned with Spark rather than relying on two independently inferred schemas remaining equivalent.

• The SQL-level coverage has been expanded substantially.
  In addition to the lower-level native tests, the updated tests exercise MATCHED / NOT MATCHED / NOT MATCHED BY SOURCE behavior, matched DELETE/copy-on-write paths, conditional first-match-wins behavior, explicit native-disabled fallback, and the relevant delta/split update paths. I wanted these at the SQL level because several of the earlier bugs were not failures of an individual Rust expression—they were failures in the contract between Spark's MERGE semantics, planning, and native execution.

• Added real Iceberg MERGE coverage with native MergeRows enabled.
  The test verifies both the resulting table contents and that execution actually engages CometMergeRowsExec rather than accidentally passing because Spark fell back to the JVM path. This gives us coverage of the integration boundary that unit-level MergeRows tests cannot provide.

• The physical ordering difference is now documented explicitly.
  Spark's implementation happens to preserve input ordering in places where the native vectorized implementation does not guarantee it. MERGE semantics do not define that physical row order as part of the result contract, so I don’t think forcing the native operator to reproduce incidental ordering is desirable. It is, however, worth making the difference explicit so it is not mistaken for an unnoticed behavioral difference.

• The shim was moved under the existing org.apache.comet.shims layout, and I reduced some of the implementation comments while keeping the ones that explain non-obvious Spark compatibility/cardinality invariants.

On the duplicated Spark 3.5/4.0 serde sources, I looked more closely at whether these should be moved into a shared source root, but I decided to keep the version-local copies intentionally.

The relevant compatibility window is specifically 3.5 + 4.0: Spark 3.4 does not have the corresponding MergeRowsExec, while Spark 4.1+ now deliberately follows a different path because of the MergeSummary contract described above. Comet's compatibility code is already organized around Spark-version-specific source roots, and there isn't currently a shared source-set abstraction representing exactly "3.5 and 4.0 but not 3.4 or 4.1+".

Introducing a new build/source-layout convention solely to deduplicate these small adapters would add another compatibility abstraction without eliminating the underlying version boundary. I think keeping them version-local is simpler and more consistent with the surrounding Comet structure.

The real concern with duplicated compatibility sources is silent drift, so I addressed that directly: there is now a parity check covering the 3.5/4.0 MergeRows serde/shim implementations. As long as they are intended to be equivalent, an accidental change to only one side will fail the test. If Spark eventually requires the implementations to diverge, that difference will have to be made explicit rather than happening unnoticed.

I’ve also added the benchmark results to the PR description rather than duplicating them here.

The two larger pieces I’m intentionally keeping outside this PR are:

• native Spark 4.1+ MergeRows, because supporting it correctly requires preserving the complete MergeSummary / V2 writer contract rather than only adding metrics; and
• Spark 4.2 InsertOnlyMergeExec, which is a separate execution operator and should be treated independently under the broader MERGE work.

At this point I think expanding #5318 further would make the change harder to review without materially improving the 3.5/4.0 implementation being added here. The current scope is therefore: make native MergeRowsExec correct and well-covered where we can preserve Spark's contract, and explicitly fall back where we currently cannot.

Thanks again for the review several of these comments helped uncover compatibility contracts that go beyond simply producing the same output rows.