fix: restore Spark write execs when reverting transition-heavy stages

[sam-1112]

Which issue does this PR close?

Closes #5719.
Closes #6152

Rationale for this change

RevertNativeForTransitionHeavyStages.revertToSpark treated every CometExec as a like-for-like swap of originalPlan. CometIcebergWriteExec and CometNativeWriteExec reported their own child as originalPlan, so reverting a transition-heavy write stage erased the write node:

• a leaf or multi-child input dropped the write entirely
• a unary input also duplicated the child (child.withNewChildren(Seq(child)))

After that, IcebergCommitExec no longer received iceberg_commit_message rows. On Spark 3.5 the reverted native Parquet plan is just Range: the write reports success and no output directory is created. The same aliasing is never a valid Spark restore for any operator.

This is pre-existing and independent of #5696. Part of #5649.

What changes are included in this PR?

This implements both fixes from #5719:

• Give the write execs a real Spark original plan:
  • CometIcebergWriteExec.originalPlan is the IcebergWriteExec it replaced
  • CometNativeWriteExec.originalPlan is the DataWritingCommandExec it replaced
• Restore through CometExec.sparkFallback instead of grafting originalPlan onto itself.
• Native Parquet writes override sparkFallback so a reverted WriteFilesExec (when present) keeps wrapping the restored input.
• Reject operators whose originalPlan is one of their own children before rewriting. If restore is invalid, skip reversion for the whole stage rather than emitting a broken write plan.
• Re-apply transformStageDown after rewriting a node so stacked transitions such as SparkToColumnar(C2R(...)) unwrap fully.
• Stop at a stage boundary when that rewrite is itself the boundary. Unwrapping a transition that sits directly on a shuffle used to keep walking into the next stage. transformStageUp and insertTransitions do not cross the exchange, so those transitions were never restored (RevertNativeForTransitionHeavyStages strips transitions in the stage below when AQE is off #6152).
• Remove the two CometExecRule arms that unwrapped a second Spark write already sitting on a native write. They only matched the old originalPlan = child link (see below).
• Document the restore contract in adding_a_new_operator.md and iceberg-writes.md, and the Arrow input requirement for reverted native shuffle stages in native_shuffle.md.
• Preserve the stage's required output format after reversion. Wrap a columnar fallback root in ColumnarToRowExec when the result stage must return rows.
• Preserve the native shuffle's Arrow input contract. When its map stage reverts, use CometSparkToColumnarExec to convert Spark rows or columnar batches into Arrow-backed Comet batches.

AQE re-plan of native writes

This changes how AQE re-plans a native write, not only how fallback restores it.

CometExecRule copies originalPlan.logicalLink onto every CometExec. On main the write execs used their child as originalPlan, so they picked up that child's link. When the child was a shuffle stage, AQE folded the write into the stage's LogicalQueryStage and the re-plan came back double-wrapped: a Spark write operator around the native write that was already there. Two arms in CometExecRule existed only to unwrap that shape:

• Spark 3.x: DataWritingCommandExec(_, WriteFilesExec) whose child is CometNativeWriteExec
• IcebergWriteExec whose child is CometIcebergWriteExec

originalPlan is now the DataWritingCommandExec or IcebergWriteExec that was replaced, so AQE re-plans the write as that operator. A shuffle directly under a native Iceberg write stays in the child stage and is not wrapped again. Those arms no longer match, the suites stay green without them, and they are removed. The comments that remain describe this contract rather than the old re-plan.

How are these changes tested?

Plan restoration coverage in RevertNativeForTransitionHeavyStagesSuite:

• Iceberg writes over a leaf, SparkToColumnar, and a unary Comet child, without duplicating the child
• Stacked SparkToColumnar(C2R) transitions under a native write
• Native Parquet writes with and without WriteFilesExec
• Invalid originalPlan == child aliases leave the stage unchanged
• Stripping transitions above a shuffle preserves the stage below it
• A reverted vectorized Spark scan retains the columnar-to-row transition required by a row-output result stage
• Map-stage fallback supplies Arrow-backed input to native shuffle for both row-based and columnar fallback roots, exercising both AQE and non-AQE rule paths

Query and write regression coverage, with AQE on and off:

• DISTRIBUTE BY returns the same rows as Spark and retains the result-stage columnar-to-row transition. The AQE-on case explicitly checks for a coalesced read over a Comet shuffle.
• Iceberg INSERT restores IcebergWriteExec and writes the expected rows.
• Partitioned copy-on-write DELETE restores the writer, commits once, preserves a shuffle under the write, and matches the rows and partition directories of a sibling table written natively.
• A forced transition below the Parquet writer restores DataWritingCommandExec → WriteFilesExec on both the Spark 3.x and Spark 4.x write paths, then executes the restored plan.
• Range and union writes compare all written rows against a Spark-written reference. The query-pipeline cases verify fallback on Spark 3.x; on Spark 4.x, these inputs remain native because they do not introduce a C2R transition.

./mvnw test -Dtest=none -Dsuites="org.apache.comet.rules.RevertNativeForTransitionHeavyStagesSuite"
./mvnw test -Dtest=none -Dsuites="org.apache.comet.CometIcebergWriteActionSuite transition-heavy fallback preserves Iceberg writes"
./mvnw test -Dtest=none -Dsuites="org.apache.comet.CometIcebergWriteActionSuite transition-heavy fallback preserves partitioned CoW"
./mvnw test -Dtest=none -Dsuites="org.apache.comet.parquet.CometParquetWriterSuite transition-heavy fallback"

[sunchao]

Correctness

The previous fallback could discard a native writer because its originalPlan was the writer's input, rather than the Spark write operator. For a unary input, restoring children could also reconstruct the wrong shape. This change keeps the original DataWritingCommandExec or IcebergWriteExec, restores the Parquet WriteFilesExec wrapper when present, and removes stacked row/columnar transitions before rebuilding the stage. It also rejects missing, aliased, or arity-incompatible fallback plans and leaves an unsafe stage unreverted.

The restored Parquet shape agrees with the maintained Spark 3.5 and 4.0 sources: DataWritingCommandExec owns command execution, and FileFormatWriter uses WriteFilesExec.executeWrite for planned writes. Retaining the original command and wrapper preserves their write metadata and save-mode handling. For Iceberg, carrying the original write node and stable commit-message output preserves the relationship with the outer commit operator and its existing commit/abort path. The native writer execution methods are unchanged. I found no additional production correctness issue in the reviewed changes. Spark 3.4 and 4.1 source compatibility remains unverified because the required maintained branches were unavailable.

One [P2] test reliability issue remains, detailed inline: captureDataWritingCommand unregisters its listener before asynchronous callback completion is guaranteed, so three new regression tests can fail after a successful write. Waiting for the callback before unregistering fixes that ordering.

Validation

Reviewed all 10 changed files at 812b06c9 against 4e69a248, including the leaf/unary fallback cases, alias guard, wrapper restoration, and the added AQE-on/off write tests. A compiled Scala probe containing the exact new helper reproduced the delayed-callback failure, with passing immediate-delivery and wait-before-unregister controls. This was an isolated test-double check, not a Spark/JNI run. I did not execute the full Comet suites locally. At September 16, 03:50 UTC, CI and CodeQL were action_required with zero jobs. The successful label job provides no test coverage.

Performance

The added alias validation makes one extra linear walk of an eligible stage during driver-side planning. It runs only when transition-heavy reversion is enabled and the transition threshold is exceeded. That feature remains disabled by default. The transition-stripping recursion removes wrappers, and the changes add no per-row work or additional write execution. I found no material new hot-path overhead requiring a separate benchmark. No runtime speedup, memory improvement, or benchmark result is claimed by this review.

Design

Keeping the actual Spark writer as the fallback source is a sound way to preserve ownership of write and commit behavior. The generic bottom-up reconstruction handles ordinary Comet nodes, while the Parquet override reconstructs the extra WriteFilesExec layer that conversion previously removed. Checking aliases before rewriting children is necessary because child replacement could otherwise hide the old reference relationship. Catching only the dedicated fallback exception keeps this conservative decision local to an unsafe restoration without swallowing unrelated failures. The source changes and the plan-shape tests address the original defect directly. The callback wait is the concrete improvement needed for dependable regression coverage.

Abstraction & complexity

sparkFallback(newChildren) gives restoration a narrow operator-level extension point instead of embedding write-specific reconstruction in the stage rule. Its shared null/alias/arity checks and the single Parquet specialization earn their complexity because these operators have different restoration shapes. The dedicated exception makes the rule's refusal to revert explicit. I found no additional abstraction or simplification issue that warrants a separate finding.

[sunchao]

Correctness

[P2] Wait for the write callback before unregistering the listener

A successful parquet(path) call does not guarantee that QueryExecutionListener.onSuccess has run: Spark posts the SQL execution-end event to the shared asynchronous listener queue. If that queue delivers the event after this finally block, the listener has already been removed and captured stays null, so the three new Parquet reversion tests can fail with expected a captured parquet write plan even though the write succeeded. I verified this delivery path in the maintained Spark 3.5 and 4.0 branches and reproduced the failing schedule with the exact helper in isolated Scala test doubles. Immediate delivery and a wait-before-unregister control both pass. Please wait with a bounded latch/future while the listener remains registered, then unregister in finally, so completion and visibility of the captured plan are guaranteed.

[sam-1112]

Fixed. The helper now waits on a bounded CountDownLatch from onSuccess while the listener is still registered, then unregisters in finally.

[sunchao]

Re-reviewed 6d7be018. The previous P2 listener race is addressed: the helper stores the plan before releasing the latch, waits up to ten seconds while the listener remains registered, and unregisters in finally. The production changes are unchanged from the previous review. I found no new or remaining P1/P2 issues.

The exact old helper still reproduces the delayed-callback failure. The exact updated helper passed immediate and delayed delivery, visibility, unrelated-event filtering, write failure, real ten-second timeout, interruption, and cleanup checks in a compiled Scala probe. These are isolated test doubles, not Spark/JNI execution. At September 16, 05:28 UTC, CI and CodeQL remain action_required with zero jobs.

[sunchao]

Rechecked 76da1090bcab against 58ab5f618e1e, including the full ten-file change and the base merge since 6d7be0184a31. No new or remaining verified P1/P2 findings. Preserving the existing approval.

The authored edits are unchanged. Parquet fallback retains the Spark write command and restores its WriteFilesExec wrapper when present. Iceberg fallback preserves the original writer, output attributes and commit boundary. Null/alias/arity guards, stacked transition removal and AQE stage boundaries remain intact. The inherited plan-cache and scan changes introduce no additional fallback edits.

The listener-race fix is also unchanged: it publishes the captured plan before signaling completion and waits while the listener is registered. The earlier eleven-scenario isolated probe remains historical evidence. I did not rerun it or execute Spark/JNI/write-IO tests on this head. Source checks used the maintained Spark 3.5/4.0 branches. Versions 3.4/4.1 remain unavailable.

At 2026-09-17 04:37:19 UTC, current-head CI and CodeQL require approval with zero jobs. Only labeling passed.

[sunchao]

Rechecked 08bec891 against 4bf725f6 and the last reviewed head 76da1090. One new [P1] build blocker is detailed inline: the rebase retained two private Parquet test helpers that now conflict with protected methods in CometParquetWriterTestBase. Both fail with weaker-access errors in isolated compiler checks on Scala 2.12.18 and 2.13.17. Removing the duplicates makes the control compile.

The production write-restoration changes, whole-stage guards and prior listener-race fix remain intact. I also checked the inherited empty-relation and transition integrations against the maintained Spark 3.5/4.0 sources. Spark 3.4/4.1 source coverage remains unavailable.

The earlier eleven-scenario listener probe is historical evidence and was not rerun. No full Comet, Spark/JNI or write-IO suite was executed. Current Comet CI and CodeQL require approval with zero jobs. Label success provides no product validation.

[sunchao]

Correctness

[P1] Remove the duplicate helpers after the base-class extraction

Could you remove this retained captureWritePlan and assertNoCometNativeWriteExec at line 1123, and use the inherited implementations? The rebased suite now extends CometParquetWriterTestBase, which already defines both methods as protected. These private declarations override them with weaker access, so Spark test compilation fails before any suite can run. I reproduced both compiler errors with the exact signatures and inheritance on Scala 2.12.18 and 2.13.17. Removing the duplicate declarations makes the isolated control compile. Their bodies already match the inherited implementations.

[sunchao]

Rechecked 86e518f3 against 5fdc9619. The P1 compilation blocker is fixed: the suite now uses the inherited protected helpers. I reran the original failure case and the updated inheritance/call-site control on Scala 2.12.18 and 2.13.17. The old declarations fail with both weaker-access errors, and the updated controls compile.

The listener-race fix and write-restoration behavior remain intact. The base update preserves the Parquet command/wrapper and Iceberg writer/commit boundaries. No new or remaining verified P1/P2 findings.

These are isolated compiler checks with trivial bodies and a SparkPlan stand-in, plus source validation against maintained Spark 3.5/4.0. The earlier eleven-scenario listener probe was not rerun. Spark 3.4/4.1 source coverage remains unavailable, and no full Comet, Spark/JNI or write-IO suite was executed. At September 19, 02:01 UTC, Comet CI and CodeQL require approval with zero jobs. Only labeling passed.

[sunchao]

Rechecked fbe1d17c against b7f35b6a. No new or remaining verified P1/P2 findings. Preserving the existing approval.

All ten authored files are unchanged from 86e518f3. The eighteen changed files match the update from main. The registry and native length changes introduce no additional interaction with write restoration. Parquet fallback still restores the original command and optional WriteFilesExec wrapper. Iceberg fallback retains the original writer, output attributes and commit boundary.

The P1 helper-visibility fix and P2 listener-race fix remain intact. I verified source identity with the earlier isolated Scala compiler and listener evidence. Those probes were not rerun. Source validation used the maintained Spark 3.5/4.0 branches. Spark 3.4/4.1 source coverage remains unavailable, and no full Comet, Spark/JNI or write-IO suite was executed.

At September 19, 12:24 UTC, Comet CI and CodeQL require approval with zero jobs. Only labeling passed, which provides no product validation.

[andygrove]

I ran this merged into current main on 3.4, 3.5, 4.0, 4.1 and 4.2, and the fix holds. With main's originalPlan = child restored, the new unit tests and both Iceberg end-to-end tests fail on 4.1, and all eight end-to-end tests fail on 3.5. It also fixes something worse than #5719 describes. On 3.5, main's reverted plan for a native Parquet write is just Range, so the write reports success and no output directory is ever created.

This also changes how AQE re-plans native writes, not only how they revert. CometExecRule copies originalPlan.logicalLink onto every CometExec, so on main the write execs picked up their child's link. When that child was a shuffle stage, AQE folded the write into the stage's LogicalQueryStage, and the re-plan came back double-wrapped. That is what the two unwrap arms in CometExecRule.scala:441-459 handle. With this PR both arms stop firing, and the suites stay green. That looks right to me, but could we say so in the PR description? Could we also either remove those arms or update their comments so they stop describing a re-plan that no longer happens?

docs/source/contributor-guide/iceberg-writes.md:374-376 still says the write execs' originalPlan is their child. The branch predates that page, so it needs a merge from main first. Could that pitfall point at sparkFallback, with a sentence in adding_a_new_operator.md about when an operator needs to override it? I'm adding run-iceberg-tests as well, since this touches every AQE re-plan of a native Iceberg write that has a shuffle directly below it.

[andygrove]

With AQE off this still descends into the stage below an exchange. The boundary check only runs on the children of the transformed node. So when the stripped transition sits directly on a shuffle, the recursion strips the transitions inside the next stage down, and nothing puts them back, because transformStageUp and insertTransitions stop at the exchange. This is the problem in #6152, and here's a concrete repro: a copy-on-write DELETE ... WHERE id IN (SELECT ...) on a partitioned table, with spark.sql.adaptive.enabled=false and maxTransitions=0, fails the shuffle map stage with ColumnarBatch cannot be cast to InternalRow. Since this PR already rewrites transformStageDown, could it return transformed unchanged when it is a stage boundary, with if (isStageBoundary(transformed)) transformed else transformStageDown(transformed)(rule)? With that, the same DELETE commits once with the right rows and the same partition layout as the native write, and the rest of the suite still passes.

[andygrove]

The end-to-end case only covers an unpartitioned INSERT ... VALUES. There is no exchange under the write, and the restored IcebergWriteExec has no ReplaceDataDispatchInfo. Could we add a copy-on-write DELETE against a partitioned table, with AQE on and off, and compare its rows and partition directories with a sibling table written natively? That is the one shape where the restored node carries state the unit tests build by hand. With AQE off it's also the case that catches the boundary problem above.

[andygrove]

I would like to do a deeper review of this PR before it is merged.

[sam-1112]

I added the symmetric ColumnarToRowExec branch for row-output stages whose reverted root remains columnar. The DISTRIBUTE BY test now runs with AQE on and off, and the AQE-on case explicitly checks for a coalesced read over a Comet shuffle. There is also a scan-only regression test for a reverted vectorized Spark scan.

I also updated map-stage fallback to use CometSparkToColumnarExec before native shuffle, since Spark columnar batches do not satisfy its Arrow input contract. New cases cover both row-based and columnar fallback roots.

[andygrove]

@sam-1112 could you fix the conflicts?

[sam-1112]

@andygrove I have fixed conflicts, thanks.

[andygrove]

I ran this merged into current main on 3.5 and 4.1. RevertNativeForTransitionHeavyStagesSuite, the full CometParquetWriterSuite and CometIcebergWriteActionSuite, and the #6137 Iceberg DPP revert test all pass, so the conflict resolution holds. I also pushed nested array, map, struct and decimal columns through the new native shuffle bridge, with Spark's vectorized and row-based readers, and the rows match Spark. On 3.5 an AQE re-plan of a native Parquet write over an aggregate shuffle comes back with a single write node, and reverting that write restores DataWritingCommandExec and WriteFilesExec with the right rows.

I've left two comments inline. Neither is a regression against main. CI never started on the latest merge commit, so it needs a re-run before this goes in the queue.

[andygrove]

The result stage still assumes it has to produce rows, and that isn't true for a cached query on Spark 4.0+. When the cache serializer accepts columnar input, which Comet's ArrowCachedBatchSerializer does, Spark marks the cached AdaptiveSparkPlanExec columnar and AQE plans the final stage with outputsColumnar = true. With the Comet cache on, AQE on, transitionRevert.enabled=true, maxTransitions=0 and spark.comet.exec.filter.enabled=false, caching SELECT _2, s FROM (SELECT _2, sum(_1) AS s FROM tbl GROUP BY _2) WHERE s > 10 and collecting it fails with FilterExec has column support mismatch. main fails the same way, so it isn't a regression, but it's the same output-format restoration this PR is fixing. Could both result-stage call sites (here and line 81) pass outputColumnar = plan.supportsColumnar && !plan.supportsRowBased instead of false? The root already says which format the consumer wants. I tried that locally and the repro passes, along with all of RevertNativeForTransitionHeavyStagesSuite. The test would need a suite that installs the cache serializer the way CometInMemoryCacheSuite does. If you'd rather keep this PR to the write fix, could we file an issue for it instead?

[andygrove]

This says the rule rebuilds the Spark operator through sparkFallback, but revertToSpark handles CometLocalTopKExec, CometNativeScanExec and CometIcebergNativeScanExec itself before it gets there (RevertNativeForTransitionHeavyStages.scala:263-281). So sparkFallback isn't the whole contract yet, and on CometLocalTopKExec the default would return a second TakeOrderedAndProjectExec that applies the TopK twice. Could those three become sparkFallback overrides, with the local TopK returning its child and the two scans carrying their live DPP filters across? Then the rule only ever calls sparkFallback, and this section holds for the next operator whose live state differs from its originalPlan. If you'd rather not move the code, could the section name the exceptions?

[sam-1112]

@andygrove Thanks for the detailed review and validation.

I moved the three special cases into operator-level sparkFallback overrides. Local TopK returns the restored child, and both scans preserve their live DPP filters and the original logical link. The rule now calls sparkFallback uniformly, and I updated the documentation and added a direct Local TopK fallback test.

The full RevertNativeForTransitionHeavyStagesSuite and the Iceberg DPP revert test pass on Spark 3.5 and 4.1: 38 tests on each profile. The 3.5 run also used strict Scala warnings.

I opened #6687 to track the cached result-stage output-format issue separately.

[andygrove]

I re-ran this at ee75eef7ce merged into current main, and both points from my last review are covered. The three sparkFallback overrides are a faithful move of the old rule code, and now that the rule only calls sparkFallback, the new section in adding_a_new_operator.md holds. The tests guard the move too. With all three overrides disabled, the local TopK tests return three rows instead of five, the direct TopK test fails, and both DPP revert tests fail on an unexecutable SubqueryAdaptiveBroadcastExec. Thanks for filing #6687 for the cached result stage.

Since your last merge, #6689 changed native shuffle to read a CometNativeArrowSource child's Arrow stream directly, so the CometSparkToColumnarExec bridge for a reverted map stage now goes through that path. I pushed nested struct, array, map and decimal columns through it, 30,000 rows so they span several batches, with hash and range partitioning and both row and columnar Spark roots, and the rows match Spark. On the merged tree, the revert suite, CometParquetWriterSuite and CometIcebergWriteActionSuite pass on 4.1 and 3.5, along with the Iceberg DPP revert test and CometIcebergWriteDetectionSuite on 4.1.