parthchandra commented on code in PR #618:
URL: https://github.com/apache/datafusion-comet/pull/618#discussion_r1662916008
##########
spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala:
##########
@@ -192,6 +192,50 @@ class CometSparkSessionExtensions
}
}
+ /**
+ * CometQueryStagePrepRule gets called from AQE for the whole plan multiple
times as the plan is
+ * re-optimized after query stages complete. This is where we translate
Spark operators and
+ * expressions into Comet/DataFusion native versions.
+ */
+ case class CometQueryStagePrepRule(session: SparkSession) extends
Rule[SparkPlan] {
+
+ private val execRule = CometExecRule(session)
+
+ def apply(plan: SparkPlan): SparkPlan = {
+ val cometPlan = execRule.apply(plan)
+ if (CometConf.COMET_CBO_ENABLED.get()) {
+ // simple heuristic to avoid moving from Spark execution to Comet
execution just
+ // for the final sort
+ // in the future, we can make this check more generic and base it on
actual costs
+ cometPlan match {
+ case CometSortExec(_, _, _, e: CometShuffleExchangeExec, _)
+ if !e.child.supportsColumnar =>
+ // fall back for sort amd exchange operators
+ val fallbackReason = "avoid move to Comet just for sort"
+ plan.setTagValue(CometExplainInfo.CBO_FALLBACK, fallbackReason)
+ plan.children.head.setTagValue(CometExplainInfo.CBO_FALLBACK,
fallbackReason)
Review Comment:
Setting the fallbackReason for the plan node should be enough I think. No
need to set it for the child node.
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
