andygrove commented on code in PR #5696:
URL: https://github.com/apache/datafusion-comet/pull/5696#discussion_r3941230551
##########
spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala:
##########
@@ -91,6 +91,18 @@ case class EliminateRedundantTransitions(session:
SparkSession)
// Write should be final operation in the plan
case ColumnarToRowExec(nativeWrite: CometNativeWriteExec) =>
nativeWrite
+ // `CometIcebergWriteExec` is row-based (it emits the serialised Iceberg
commit message) but
+ // consumes Arrow batches from its child over FFI, so Spark inserts a
columnar-to-row
+ // transition *underneath* it. Strip it so `doExecuteColumnar` sees the
columnar child
+ // directly; `CometIcebergNativeWrite.requiresNativeChildren` already
guarantees that child
+ // was Comet-native when the write was converted.
+ //
+ // The write deliberately does not tag itself as a
`ColumnarToRowTransition` to suppress the
+ // insertion: Spark leaves such a node untouched, so the whole subtree
below the write is
+ // never visited and the transitions the rest of that subtree needs are
never inserted
+ // (https://github.com/apache/datafusion-comet/issues/5689).
+ case w: CometIcebergWriteExec =>
+ stripColumnarToRow(w.child).map(child =>
w.withNewChildren(Seq(child))).getOrElse(w)
Review Comment:
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 `ColumnarVector`s where the FFI
adapter requires `CometVector`s. 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:
```scala
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.
--
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]
