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

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

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.

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.

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?

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.

Hi @sunchao @andygrove, when you get a chance, could you please take another look and trigger CI for the latest changes? Thanks!

Found one P2 issue at head a4541f6.

Valid MERGE fails for nested non-null fields

The new check in planner.rs:2206 rejects valid expressions when native metadata marks a nested field nullable.

For example, with target type ARRAY<STRUCT<n: INT NOT NULL>>:

WHEN MATCHED THEN UPDATE SET payload =
  array(named_struct('n', coalesce(s.id, 0)))

This succeeds on Spark 4.0.4, but enabling native MERGE fails with incompatible with Spark output type. Comet widens the nested nullability during expression conversion; the new guard rejects it before runtime schema normalization can handle it.

Allow safe normalization with runtime validation, or fall back during Spark planning.

Validation

• 23 native tests passed.
• 10 existing Spark 4.0 MERGE tests passed.
• Two added SQL probes reproduced this same defect.
• The four previous findings are fixed.
• Current CI has only a passing labeling check. Spark 3.5 and Iceberg-specific execution were not tested locally.

Checkout restored; nothing posted to GitHub.

@sunchao @andygrove would appreciate a review when you get a chance. Thanks!

Thanks for the detailed review, @andygrove really appreciate you taking the time to test this across versions and cases.

I've merged current main, so the conflict is resolved and have addressed your inline comments

On the source duplication: I looked at moving the shared implementation out of the 3.5/4.0 roots, but there isn't currently a common source root that maps cleanly to exactly those two versions. Introducing another Maven source root just for this felt like more build-layout complexity than this PR warrants.

So the remaining option is to keep the copies and enforce that they stay identical. I initially tried a direct diff in the CI workflow, but that felt a bit too one-off. Would you prefer a small reusable dev/ci source-parity check for this, or is an explicit CI diff sufficient?

Found one P2 issue at head edd1c6ba8.

[P2] Skip memory-consumer registration when cardinality checking is disabled — merge_rows.rs:248.

execute() always registers a memory consumer, including insert-only MERGEs with multiple clauses where row_id_ordinal=None. That consumer never reserves memory, but the default fair_unified pool counts it when dividing the task’s budget. Enabling native MergeRows can therefore make an otherwise successful upstream join fail.

I reproduced this with the actual native operator, fair pool, and LeftAnti hash join: the join returned 4,096 rows with 96 KiB, but wrapping it in MergeRows reduced its allowance to 48 KiB and caused ResourcesExhausted. The probe used unlimited mock Spark grants, isolating the failure to native memory accounting. Register the consumer conditionally or lazily when cardinality state is needed.

Validation:

• 23 existing native MergeRows tests passed with --no-default-features.
• The additional regression probe confirmed the issue.
• Earlier correctness findings appear addressed; no additional defects confirmed.
• JVM/Iceberg execution remained blocked by Maven dependency resolution. Current CI requires maintainer action and provides no build/test result.

Checkout restored and clean. Nothing posted to GitHub.

Thanks @andygrove , addressed the remaining points and pushed the updates. Also added the parity check and fixed the unnecessary memory consumer registration.

@andygrove whenever you get a chance, would really appreciate another look at this. I have a few small follow-ups ready on top once this lands.

Thanks @andygrove! Addressed all three comments:

• Cardinality memory: I’d prefer to keep the RoaringTreemap switch out of this PR, since getting the reservation sizing and duplicate-vs-OOM error ordering right deserves a focused change. Opened Use RoaringTreemap for MergeRows cardinality tracking #6608 for that, and added a note about the memory cost to the compatibility page.
• Spark 4.1+: Opened Support native MergeRowsExec on Spark 4.1+ with MergeSummary #6606 to track native MergeRowsExec support on 4.1+ with the MergeSummary contract. The config docs now clarify that the flag only takes effect on Spark 3.5 and 4.0, with links to the follow-up issue. I also updated the compatibility page and fixed the CometExecRule comment.
• Tests: Removed the redundant COMET_ENABLED=true entries. The explicit false baselines are still retained.