fix: restore columnar transitions under the native Iceberg write

[andygrove]

Which issue does this PR close?

Closes #5689.

Rationale for this change

CometIcebergWriteExec is tagged with ColumnarToRowTransition so that Spark does not wedge a
ColumnarToRow between the write and its Comet-native child. That trait does more than suppress
one insertion. In Spark's ApplyColumnarRulesAndInsertTransitions:

private def ensureOutputsRowBased(plan: SparkPlan): SparkPlan = {
  if (plan.supportsColumnar && !plan.supportsRowBased) {
    ColumnarToRowExec(ensureOutputsColumnar(plan))
  } else if (!plan.isInstanceOf[ColumnarToRowTransition]) {
    plan.withNewChildren(plan.children.map(insertTransitions(_, outputsColumnar)))
  } else {
    plan   // returned untouched -- the subtree is never visited
  }
}

CometIcebergWriteExec is row-based (supportsColumnar = false), so it lands in the third branch
and the entire subtree below the write skips the transition-insertion pass.

The Iceberg copy-on-write rewrite plan feeds a Spark-columnar BatchScan (IcebergCopyOnWriteScan)
into row-based joins and filters, so the ColumnarToRow that plan needs is never inserted and
every CoW DELETE / UPDATE / MERGE fails at runtime with:

java.lang.ClassCastException: class org.apache.spark.sql.vectorized.ColumnarBatch
  cannot be cast to class org.apache.spark.sql.catalyst.InternalRow

With AQE enabled the failure disappears, because each stage gets its own insertion pass when it
materialises. Every existing Comet Iceberg suite runs with AQE on, which is why this was never
caught here; it is the dominant failure in Iceberg's own spark-extensions suites (see #5649),
whose ExtensionsTestBase randomises AQE per session.

What changes are included in this PR?

• CometIcebergWriteExec: drop the ColumnarToRowTransition trait so Spark walks the write's
  subtree normally and inserts the transitions it needs. The comment explaining why the node is
  not a transition is kept and expanded.
• EliminateRedundantTransitions: strip the columnar-to-row transition Spark now inserts below
  the write, so doExecuteColumnar still sees the columnar child directly. This mirrors the
  existing ColumnarToRowExec(nativeWrite: CometNativeWriteExec) arm. The new
  stripColumnarToRow helper handles all three variants (ColumnarToRowExec and the two Comet
  ones), because transformUp has usually already rewritten the plain node by the time the
  parent arm sees it.

The strip is unconditional: CometIcebergNativeWrite.requiresNativeChildren = true already
guarantees the write's child was a CometNativeExec at conversion time. Guarding it on
child.isInstanceOf[CometPlan] is wrong -- under AQE the transition sits over an
AQEShuffleReadExec, a plain Spark node, and the guard leaves the transition in place (ten
AQE-on tests fail with requires a columnar (Comet native) child; got WholeStageCodegenExec).

How are these changes tested?

New test native acceleration: ReplaceData (CoW DELETE) with AQE disabled in
CometIcebergWriteActionSuite, plus an assertColumnarContract helper that walks the executed
plan and flags any row-based operator consuming a columnar-only child (the shape that produces the
ClassCastException at runtime rather than a planning error).

• Without the fix, the new test fails with exactly the ClassCastException from the issue, while
  the other 53 tests in the suite pass -- so it reproduces the bug rather than being vacuous.
• With the fix, 275 tests pass across CometIcebergWriteActionSuite (54),
  CometIcebergWriteDetectionSuite (46), CometIcebergRewriteActionSuite (5),
  CometIcebergSystemFunctionSuite (11), CometExecSuite (144) and
  RevertNativeForTransitionHeavyStagesSuite (15), on the default Spark 4.1 / Iceberg 1.11.0
  profile.

[andygrove]

cc @jordepic

[jordepic]

Thanks for tracking this down. I traced the same path through Spark and agree with the root cause. ensureOutputsRowBased has had the same three-branch shape since 3.4.3 (there it is insertTransitions directly), so the fix applies uniformly across the supported Spark versions.

On the question of whether dropping ColumnarToRowTransition affects the commit-message row that IcebergCommitExec collects: it does not. The row contract comes from supportsColumnar = false plus doExecute / executeCollect, and none of that changes. In Spark the trait is only consulted in ensureOutputsRowBased and in CachedBatchSerializer (for InMemoryRelation), so on the write it was purely a marker to suppress insertion. IcebergCommitExec.collectAndCommit still calls child.executeCollect() and deserialises the binary column exactly as before.

I also walked the AQE path. InsertAdaptiveSparkPlan wraps the child of the V2CommandExec, so the AdaptiveSparkPlanExec root is the IcebergWriteExec and its supportsColumnar is false. The final stage therefore runs insertTransitions(_, outputsColumnar = false), which puts a ColumnarToRowExec between the write and its child and the new arm strips it again. One small correction to the description: at that point the write's child is the CometSinkPlaceHolder that CometExecRule wraps around the ShuffleQueryStageExec, which is a CometPlan, so a CometPlan guard would have matched too. The unconditional strip is still the right call, since requiresNativeChildren already guarantees the child type.

[jordepic]

The comment above says the child was guaranteed Comet-native at conversion time, but RevertNativeForTransitionHeavyStages runs between conversion and this rule (it is first in postColumnarTransitions). That rule counts ColumnarToRowTransition nodes in the stage, and the write itself used to count as one, so the new ColumnarToRowExec underneath simply takes its place and the count is unchanged. So no behaviour change from this PR there.

While looking at that interaction I noticed something pre-existing that this PR does not introduce but sits right next to: CometIcebergWriteExec.originalPlan is child, and revertToSpark replaces every CometExec with originalPlan.withNewChildren(children). If the write's stage ever exceeds maxTransitions (default 2, reachable with spark.comet.sparkToColumnar.enabled and, say, a row-based Union of two Spark-columnar scans directly under the write), the write node disappears and IcebergCommitExec would try to deserialise data rows as commit messages. CometNativeWriteExec has the same originalPlan = child. Does that deserve a tracking issue? It feels like a separate fix, but it is the same transition-accounting area this PR touches, so I wanted to raise it here rather than lose it.

[andygrove]

Confirmed on both counts.

On the transition accounting: agreed, the count is unchanged. RevertNativeForTransitionHeavyStages runs first in postColumnarTransitions and countTransitions counts ColumnarToRowTransition nodes, so the ColumnarToRowExec Spark now inserts below the write takes the place the write itself used to occupy. No behaviour change from this PR there.

On originalPlan: I read revertToSpark again and it is worse than a disappearing node. The arm is

case cometExec: CometExec =>
  if (cometExec.originalPlan.children.size == cometExec.children.size) {
    cometExec.originalPlan.withNewChildren(cometExec.children)
  } else { ... }

For CometIcebergWriteExec, originalPlan is child and children is Seq(child), so when the child is itself unary the size check passes and the result is child.withNewChildren(Seq(child)): the write is gone and the child is duplicated beneath itself. CometNativeWriteExec has the same originalPlan = child. IcebergCommitExec would then deserialise data rows as commit messages, as you say.

This is pre-existing and orthogonal to the transition fix, so I have left it out of this PR. It needs its own issue and fix; the shape of the fix is probably that a write exec should not report its child as originalPlan, or that revertToSpark should decline to revert a node whose originalPlan is one of its own children.

[andygrove]

Filed as #5719, with you credited for spotting it.

I probed revertToSpark directly before writing it up, and both branches of the arity check are broken, not just the warning one:

• leaf or multi-child child (size mismatch): returns the child as-is, write gone
• unary child (both sizes 1): child.withNewChildren(Seq(child)), so the write is gone and the child is duplicated

=== INPUT ===
CometIcebergWrite [iceberg_commit_message#6], , ICEBERG_WRITER_UNPARTITIONED
+- CometFilter [_1#4], (isnotnull(_1#4) AND (_1#4 > 2))
   +- CometNativeScan parquet [_1#4] ...

=== REVERTED ===
Filter (isnotnull(_1#4) AND (_1#4 > 2))
+- Filter (isnotnull(_1#4) AND (_1#4 > 2))
   +- ColumnarToRow
      +- FileScan parquet [_1#4] ...

The duplication is harmless for an idempotent Filter but not for a projection with monotonically_increasing_id(), a sample, or a limit.

[jordepic]

This helper is a nice general guard for exactly the class of bug the issue describes (a Comet rewrite hiding part of the subtree from Spark's insertion pass), and capturePlans already records qe.executedPlan for every write in the suite. Would it be worth calling assertColumnarContract from capturePlans (or captureWrite) so all the existing tests check the contract too, rather than only this one? That would also make it cheap to cover the UPDATE and MERGE shapes from the issue with AQE off, since MERGE in particular puts a different join and MergeRows between the columnar CoW scan and the write.

[andygrove]

Done in bc6429f. captureWrite now runs assertColumnarContract over every captured plan, and the explicit call in the CoW DELETE test is gone as redundant.

Also added the other two CoW shapes with AQE off. One thing that came out of it: the partitioned MERGE does engage natively, unlike the unpartitioned one the existing test pins. Partitioning puts a CometColumnarExchange (REBALANCE_PARTITIONS_BY_COL) between MergeRowsExec and the write, so the write's own child is Comet-native and requiresNativeChildren is satisfied even though MergeRowsExec stays JVM. Its plan is the strongest of the three for this issue, since the subtree needs transitions in three separate places:

CometIcebergWrite
+- CometColumnarExchange hashpartitioning(region, 10), REBALANCE_PARTITIONS_BY_COL
   +- MergeRowsExec
      +- *(4) CometColumnarToRow
         +- CometSortMergeJoin FullOuter
            :- CometSort
            :  +- CometColumnarExchange
            :     +- *(1) Project
            :        +- *(1) ColumnarToRow
            :           +- BatchScan IcebergCopyOnWriteScan
            ...

The test asserts native engagement rather than fallback as a result.

[andygrove]

Thanks for tracing it through — agreed on the trait being purely a marker on the write, and on the commit-message contract being unaffected.

One thing I want to push back on: the placeholder is not there by the time transitions get inserted. CometExecRule._apply strips CometSinkPlaceHolder and CometScanWrapper at the end of its own pass, and that pass is preColumnarTransitions, so insertTransitions never sees a placeholder. I did try the CometPlan guard first, and it took out ten AQE-on tests across CometIcebergWriteActionSuite, CometIcebergWriteDetectionSuite and CometIcebergSystemFunctionSuite, every one of them with CometIcebergWriteExec requires a columnar (Comet native) child; got WholeStageCodegenExec — the guard rejected the transition's child, so the ColumnarToRow survived and CollapseCodegenStages wrapped it.

Instrumenting the rule on the CoW DELETE test prints write child=ColumnarToRowExec stripped=Some(AQEShuffleReadExec) aqe=true, so under AQE the child is an AQEShuffleReadExec, which is a plain Spark node. That also explains why the plain ColumnarToRowExec form has to stay in stripColumnarToRow: the hasCometNativeChild arm above never fires on it, because QueryStageExec is a LeafExecNode and the op.exists(...) walk cannot see the Comet exchange inside the stage. I pushed a comment recording that, since it was not obvious to me either.

The TPC-H check that went red on the first run failed with Network is unreachable during setup; the push above re-runs it.

[sunchao]

Correctness

Review of 585ffe8e against authoritative base 719cba11.

Removing ColumnarToRowTransition fixes the traversal problem: on the maintained Spark 3.5 and 4.0 sources, that marker causes transition insertion to return the writer without visiting its subtree. Keeping supportsColumnar = false preserves its row-shaped commit-message output. Spark can now insert the inner transitions needed by mixed row/columnar CoW plans, and the new rule removes the writer's immediate input transition. IcebergCommitExec still collects and deserializes the same binary commit messages. This change does not alter expression evaluation, null handling, ANSI behavior, numeric boundaries or serialization formats.

There is one P2 in the input-boundary cleanup. A child admitted through CometScanWrapper can become CometSparkToColumnarExec when conversion wrappers are removed. The existing bottom-up cancellation processes that child before the new writer arm and removes the required Arrow conversion. A row source then violates the writer's columnar requirement. A Spark-columnar source can still report supportsColumnar while no longer supplying the required CometVectors. The inline comment describes the reachable configuration and the boundary that needs preserving.

I checked jordepic's existing review, both unresolved threads and the author's new explanation in issue comment 5546337821. The AQE explanation is consistent with the maintained source: AQEShuffleReadExec delegates columnar support to its child, but the ordinary exists traversal stops at QueryStageExec, so a plain C2R can legitimately reach the writer's cleanup. Preserving that plain-transition case is necessary. It does not address the separate Arrow-bridge cancellation described here. This finding is distinct from the transition-heavy fallback issue: the reproduced row-input plan has one transition, below the default threshold of two. The suggestion to extend the contract checks across UPDATE/MERGE is already covered by the existing feedback. The current MERGE test explicitly expects JVM fallback when MergeRowsExec breaks native conversion.

Validation reuses the prior matched component probe with an explicit source-equivalence check. The sole change from 992eb6c7098fa80717375eb901232adf77ba8812 to current HEAD 585ffe8ee6d5f0d61a33d0695e4efb6285f21342 is the helper's Scaladoc. Removing that one documentation block leaves the entire rule file byte-identical, and every other source, dependency and tree entry is unchanged. The retained probe compiled the unmodified pre-PR and preceding-HEAD cleanup rules with the insertion rule extracted from maintained Spark 4.0 at 03f28fc4318024830a2ee8da7e83c42e0994d37a, actual Spark 4.0.4 plan/transition runtime classes, and small Comet node doubles. It reproduced the lost Arrow conversion with a real RangeExec, demonstrated the intended repair in a CoW-shaped subtree, and passed native-input, all-three-transition-variant, unrelated-transition and AQE QueryStageExec wrapper controls. No build or probe was rerun for this documentation-only delta. This is component planning evidence, not native/Iceberg execution. Maintained Spark 3.5 at 5947fd6e74a1b2b04e4f83b7a659b02a9a2bac8b has the same insertion decisions. Both maintained pins were reverified, and the new AQE comment was traced against both branches. Maintained 3.4/4.1 sources remain unavailable.

The new DELETE test checks native engagement, one snapshot commit, a columnar child and the resulting IDs. Its walker traverses subqueries and adaptive-stage contents on the maintained sources. I did not independently run that SQL regression or the claimed 275-test suite result. The fresh current-head check read at 21:45 UTC reports 34 successful, 34 running and seven skipped checks, with no recorded failures. Native builds and relevant Spark/Iceberg pipelines remained in progress. The author attributes an earlier TPC-H failure to setup networking. I did not independently audit that earlier failure, and it is not counted as a successful test. The current authored and incremental diffs pass whitespace checks. The previous isolated-index apply check against authoritative base 719cba11076e4237d6030925c49b1c1ffcac6f8e is retained as prior-head evidence. The new comment is at the same unchanged helper location, but no new merge/apply or merged-runtime test is claimed.

Performance

The change restores a necessary planning traversal and removes the writer's redundant row conversion before execution. It adds a constant amount of matching and allocation at each native Iceberg writer in the existing cleanup pass. It does not add a data scan, shuffle, per-row loop or serialization stage, and the default-disabled native-write path does not acquire the writer-specific work. The required Arrow conversion in the finding cannot be canceled merely because the writer advertises row output. No performance measurement was made, and this correctness repair does not introduce a new expression or newly default-enabled execution feature that calls for a matched microbenchmark.

Design

Separating the writer's row output from its Arrow input is the right approach. A marker that stops recursion conflates those two contracts and hides transitions throughout the subtree. The proposed immediate-child cleanup is appropriately scoped, but it needs to protect that boundary before bottom-up cancellation can erase the Arrow producer. Checking only supportsColumnar afterward would not cover Spark-columnar vectors. The conversion-time CometNativeExec check also cannot establish the final child type because CometExecRule removes its scan and sink wrappers before transition insertion. AQE stage wrappers retain their columnar capabilities and need to remain valid inputs.

Abstraction & complexity

The three-variant unwrapping helper is small and justified because the cleanup pass can rewrite a plain transition before reaching its parent. A focused writer-boundary treatment can preserve this approach without introducing a generic traversal framework. The remaining complexity is the ordering contract between wrapper removal, Spark insertion and bottom-up cancellation. It should be pinned with both a row-to-Arrow input and a Spark-columnar-to-Arrow input, since the latter satisfies the boolean columnar check while still requiring a representation conversion.

[sunchao]

Correctness

[P2] Preserve the Arrow bridge before visiting the writer

Could this protect the writer's input before the bottom-up cancellation runs, while retaining the plain C2R handling needed by AQE? A direct leaf converted by CometSparkToColumnarExec.createExec is initially a CometScanWrapper, which extends CometNativeExec and passes requiresNativeChildren. CometExecRule then removes the wrapper, leaving the Arrow bridge directly below the writer. With this PR, Spark inserts ColumnarToRowExec(CometSparkToColumnarExec(source)), and the existing arm at lines 81–88 processes it before this writer arm. For a row source it removes both transitions, so this arm receives a bare row child and cannot restore the bridge. The writer then throws CometIcebergWriteExec requires a columnar (Comet native) child.

This affects a direct RangeExec/RDD source with Spark-to-columnar conversion enabled and that operator allowed, for example an unpartitioned identity append without an intervening sort/exchange. The matched component probe with a real RangeExec retains the bridge before the PR and loses it in the proposed implementation. The new documentation commit leaves that executable code unchanged. Both plans have only one transition, so this is separate from the transition-heavy fallback concern. A Spark-columnar source also loses its Arrow bridge, leaving non-CometVector input for the FFI adapter even though supportsColumnar remains true. Please preserve that conversion and cover both input representations in the regression tests.

[andygrove]

You are right, and I reproduced it before fixing it rather than taking the probe on trust. Fixed in bc6429f.

The repro is two plan-level tests in CometIcebergWriteDetectionSuite that build CometIcebergWriteExec -> CometSparkToColumnarExec -> source, run ApplyColumnarRulesAndInsertTransitions and then this rule, and assert the write still sits on the bridge. Against the previous commit both fail:

- row source keeps its Arrow bridge under the native Iceberg write *** FAILED ***
- Spark-columnar source keeps its Arrow bridge under the native Iceberg write *** FAILED ***

Both source representations are covered because the cancellation at lines 81-88 damages them differently: over a row source it drops the CometSparkToColumnarExec outright and the write is left with a row child; over a Spark-columnar source it keeps a ColumnarToRowExec, which this arm then strips, leaving the write reading Spark ColumnarVectors where the FFI adapter requires CometVectors. The second case would not even be caught by a supportsColumnar check on the child, as you note.

The fix is the ordering rather than a guard. The strip is now its own pass ahead of the transformUp, so the write's input transition is removed before anything else can consume it:

val eliminatedPlan = stripIcebergWriteInputTransition(plan) transformUp { ... }

A guard inside transformUp cannot work, because transformUp visits children first: whatever the write's arm checks, the child has already been rewritten by the time it looks. Running before the traversal also means the pre-pass sees exactly what Spark's insertion pass produced, which keeps the AQE AQEShuffleReadExec case working unchanged.

Suite results on Spark 4.1 / Iceberg 1.11.0: CometIcebergWriteActionSuite 56, CometIcebergWriteDetectionSuite 48, CometIcebergRewriteActionSuite 5, CometIcebergSystemFunctionSuite 11, RevertNativeForTransitionHeavyStagesSuite 15, CometExecSuite 144, all passing.

[sunchao]

Correctness

Re-reviewed bc6429faec5339c84815311b9187ae6a6e28ad84 against assigned base 719cba11076e4237d6030925c49b1c1ffcac6f8e, focusing on the update since 585ffe8e. The previous P2 is addressed. Stripping the writer's immediate input transition before the bottom-up cancellation pass preserves CometSparkToColumnarExec for both row and Spark-columnar producers. The generic cancellation rules can no longer remove that Arrow bridge before the writer is visited.

This agrees with the transition insertion semantics in the maintained Spark 3.5 and 4.0 branches: removing the marker lets Spark visit the writer's subtree and insert conversions needed by row consumers; the separate cleanup removes only the conversion at the native writer's input. Plain Spark, Comet, and native Comet C2R forms remain supported, including plain transitions over an AQE stage. The binary commit-message row contract is unchanged. The new global contract assertion also covers the existing write tests. The pre-existing transition-heavy fallback problem is now tracked separately in #5719; this update does not change that earlier fallback pass.

A fresh matched planning probe compiled the exact previous and updated rules with maintained Spark 4.0 insertion source and verified Spark 4.0.4 runtime libraries. It reproduced prior bridge loss and current preservation, with 32 prior and 41 current assertions passing. Controls covered all three transition forms, actual Range and QueryStage nodes, inner row/columnar transitions, native and JVM writer inputs, unrelated cancellation, and repeated application. Comet nodes were doubles; this was not a native Iceberg execution test.

The Spark 4.1/JDK 17 scans CI job explicitly passed both new bridge tests and AQE-disabled native UPDATE/MERGE tests: 485 passed, zero failed, four canceled and one ignored. The exec job passed 812 tests with zero failures and three ignored. Both executed merge e4df7f5 against a1f8bd3e; all four authored PR files match the reviewed head, but this is not execution against the assigned base. The snapshot has 67 successful and nine skipped checks. Maintained Spark 3.4/4.1 source was unavailable locally, so no local source-compatibility claim is made for those versions. I found no remaining PR-introduced correctness issue in this scope.

Performance

The new pre-pass adds one O(N) planning traversal, including plans without an Iceberg writer. It introduces no per-row or per-batch work and preserves the existing bridge instead of changing its conversion cost. The native execution and dependency trees are unchanged in this update. The planning probe is correctness evidence, not a performance measurement; no runtime speedup or measured disabled-path overhead is claimed. I found no material performance issue requiring a change.

Design

Separating writer-input cleanup from generic cancellation makes the ordering requirement explicit and fixes the demonstrated failure without changing cancellation for other consumers. Retaining the plain Spark transition case is appropriate for AQE inputs; restricting it to a Comet node type would exclude a supported plan shape. The added UPDATE/MERGE execution coverage and shared contract assertion address the earlier coverage requests.

Abstraction & complexity

The private helper keeps one writer-specific rule together and leaves the existing bottom-up rule behavior intact. It introduces no new plan node, configuration, protocol, or public abstraction. Its scope is proportionate to the ordering constraint; I have no additional actionable finding here.

[andygrove]

Merged. Thanks @sunchao! I'll keep iterating on this experimental feature