viirya commented on code in PR #6852:
URL: https://github.com/apache/datafusion-comet/pull/6852#discussion_r4263106035
##########
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:
Does the non-AQE path need this check? Without AQE, I think the consumer
adapts to whatever root the rule leaves. `CachedRDDBuilder` picks
`executeColumnar()` or `execute()` from `cachedPlan.supportsColumnar` on the
final plan, after this rule has run, and the other consumers read rows. Only
AQE fixes `AdaptiveSparkPlanExec.supportsColumnar = true` before it plans the
final stage, so that is the only place where a row-only reverted root breaks
the consumer.
So a non-AQE dual-format root could still be reverted the way `main` does
it. Keeping it native there gives up the revert for no correctness gain. It
also differs from "Without AQE, Spark always plans for row output, so that path
behaves as before" in the description, and the `AQE=false` case of the new test
pins that behavior.
Would it work to apply the dual-format check only when
`adaptiveExecutionEnabled` is set? Keying it on the AQE conf rather than on
`wholePlan` would keep the plan-only preview under AQE consistent with what AQE
executes. The `AQE=false` case of the test could then check that the stage is
reverted.
--
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]
