andygrove commented on code in PR #6852:
URL: https://github.com/apache/datafusion-comet/pull/6852#discussion_r4265127983
##########
spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala:
##########
@@ -65,19 +65,39 @@ case class RevertNativeForTransitionHeavyStages(session:
SparkSession, wholePlan
case exchange: ShuffleExchangeLike =>
revertShuffleStageIfNeeded(exchange).getOrElse(plan)
case _ =>
- // Result stage: its output is collected as rows.
- revertStageIfNeeded(plan, outputColumnar = false).getOrElse(plan)
+ revertResultStageIfNeeded(plan).getOrElse(plan)
}
}
private def applyForNonAQE(plan: SparkPlan): SparkPlan = {
val withRevertedStages = plan.transformUp { case exchange:
ShuffleExchangeLike =>
revertShuffleStageIfNeeded(exchange).getOrElse(exchange)
}
- revertStageIfNeeded(withRevertedStages, outputColumnar = false)
- .getOrElse(withRevertedStages)
+ revertResultStageIfNeeded(withRevertedStages).getOrElse(withRevertedStages)
}
+ /**
+ * Reverts the result stage if needed, keeping the output format its
consumer reads. That is
+ * usually rows, but on Spark 4.0+, when the cache serializer accepts
columnar input, as Comet's
+ * does, Spark marks a cached `AdaptiveSparkPlanExec` columnar and AQE plans
its final stage to
+ * output columnar batches. Only then does Spark's transition insertion
leave a columnar-only
+ * root; when rows are wanted it puts a columnar-to-row transition on top.
+ *
+ * Spark leaves a root that supports both formats as it is in either case,
so such a root, for
+ * example `ArrowEvalPythonExec` over a columnar child on Spark 4.2, does
not say which format
+ * its consumer reads. No transition provides both, so that stage is
reverted only if the
+ * reverted root supports both formats too.
+ */
+ private def revertResultStageIfNeeded(plan: SparkPlan): Option[SparkPlan] =
+ if (plan.supportsColumnar && plan.supportsRowBased) {
Review Comment:
You're right. Without AQE, `CachedRDDBuilder` reads
`cachedPlan.supportsColumnar` from the plan this rule leaves, and every other
consumer reads rows, so the check bought nothing there and only kept the stage
native. In bc5bddc208 it applies only when `spark.sql.adaptive.enabled` is set,
read from the conf rather than from `wholePlan` as you suggested, and without
AQE a root that supports both formats is reverted to row output as on `main`.
The test now runs under AQE with `wholePlan` off and on, where both stages
behave as before, and without AQE, where both are reverted to row output.
Keying the check on `wholePlan` fails the preview case, and dropping the AQE
condition fails the case without AQE. I've updated the description to match.
--
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]
