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

[sam-1112]

Which issue does this PR close?

Closes #5719.

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. 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.

How are these changes tested?

Direct revertToSpark coverage in RevertNativeForTransitionHeavyStagesSuite:

• Iceberg write over a leaf, over SparkToColumnar, and over a unary Comet child (no duplicated Filter)
• stacked SparkToColumnar(C2R) under a native write
• native Parquet write with and without WriteFilesExec
• invalid originalPlan == child aliases leave the stage unchanged

End-to-end, with AQE on and off:

• Iceberg INSERT with transitionRevert.enabled=true and maxTransitions=0 restores IcebergWriteExec and still writes the rows
• Parquet writes (including a union of row sources) restore DataWritingCommandExec → WriteFilesExec and write the expected row counts

./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.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.