andygrove commented on code in PR #6852:
URL: https://github.com/apache/datafusion-comet/pull/6852#discussion_r4260562971
##########
spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala:
##########
@@ -65,19 +65,27 @@ 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.
+ */
+ private def revertResultStageIfNeeded(plan: SparkPlan): Option[SparkPlan] =
+ revertStageIfNeeded(plan, outputColumnar = plan.supportsColumnar &&
!plan.supportsRowBased)
Review Comment:
You're right, the root alone can't tell us which format was requested when
it supports both. There was a second way to hit it, too. If the reverted root
still supported both formats, `revertStageIfNeeded` put a `ColumnarToRowExec`
over it, and a columnar consumer failed the same way.
In 7552039264, a root that supports both formats is reverted only when the
reverted root supports both as well, and it goes back with no transition on
top. Otherwise the stage stays native, since no transition gives both formats.
The doc comment now says that such a root leaves the requested format unknown.
PySpark can't run in our JVM suites, so the new test uses a stand-in
operator that reports its formats the way `ArrowEvalPythonExec` does on 4.2. It
covers both cases with AQE on and off, and it fails against the previous commit.
--
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]
