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.

[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]

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.