andygrove commented on code in PR #5957:
URL: https://github.com/apache/datafusion-comet/pull/5957#discussion_r4116433088
##########
spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala:
##########
@@ -146,16 +155,31 @@ case class RevertNativeForTransitionHeavyStages(session:
SparkSession, wholePlan
}
/**
- * Like `transformDown`, never descends stage-boundary children.
+ * Like `transformDown`, never descends stage-boundary children. If the rule
rewrites the
+ * current node, re-apply it to the result so stacked transitions such as
+ * `CometSparkToColumnarExec(CometNativeColumnarToRowExec(x))` are fully
unwrapped before
+ * children are visited. Spark's `transformDown` does not do this; leaving
the inner C2R in
+ * place later calls `CometNativeColumnarToRowExec.withNewChildren` with a
reverted row-based
+ * child, which asserts `child.supportsColumnar`.
+ *
+ * A rewrite can itself be the stage boundary. Unwrapping a transition that
sits directly on a
+ * shuffle yields that shuffle, and descending into it strips transitions in
the next stage.
+ * `transformStageUp` and `insertTransitions` do not cross the exchange, so
those transitions
+ * would not be restored (#6152). Return the boundary unchanged.
*/
private def transformStageDown(plan: SparkPlan)(
rule: PartialFunction[SparkPlan, SparkPlan]): SparkPlan = {
val transformed = rule.applyOrElse(plan, identity[SparkPlan])
- val newChildren = transformed.children.map { child =>
- if (isStageBoundary(child)) child else transformStageDown(child)(rule)
+ if (transformed ne plan) {
+ if (isStageBoundary(transformed)) transformed
Review Comment:
This stops the strip at the exchange, but if the stage's root is itself the
transition sitting on the exchange, `stripped` is the exchange. Then
`transformStageUp` and `insertTransitions` are called with a boundary as the
root. They only check children for boundaries, so they walk into the map stage.
With AQE off, `transitionRevert.enabled=true` and `maxTransitions=0`, `SELECT
_1, _2 FROM tbl DISTRIBUTE BY _2` plans as `CometColumnarToRow` over
`CometExchange` over `CometNativeScan`. The revert turns the scan into Spark's
`FileScan` under the still-native shuffle and drops the row transition at the
top, and the query fails with a `ClassCastException` casting Spark's
`OnHeapColumnVector` to a Comet vector. Could `revertToSpark` leave the stage
alone when `stripped` is a stage boundary? When I tried `if
(isStageBoundary(stripped))` throwing `InvalidSparkFallbackException`, the
query returned the right rows, the plan stayed native, and the rest of
`RevertNativeForTransitionHeavyStage
sSuite` passed. A test with that `DISTRIBUTE BY` query would cover it, and
then this PR could close #6152 as well.
--
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]
