sam-1112 opened a new pull request, #5957:
URL: https://github.com/apache/datafusion-comet/pull/5957
## 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
```bash
./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"
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
