viirya commented on code in PR #7357:
URL: https://github.com/apache/datafusion-comet/pull/7357#discussion_r4431010837
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1108,37 +1118,30 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
// Set up logical links
newPlan = newPlan.transform {
- case op: CometExec
- if op
- .getTagValue(SparkPlan.LOGICAL_PLAN_TAG)
- .exists(_.isInstanceOf[LogicalQueryStage]) =>
- // AQE replanning reuses this physical root and links it to the
current logical stage.
- // originalPlan can still point to a subtree hidden inside that
logical leaf, which
- // AQE cannot replace in the current logical plan. Only preserve a
direct stage link,
- // not a link inherited from an ancestor.
- // On the ordinary exchange path, the exchange itself is behind a
QueryStageExec
- // leaf and is not visited by this transform.
- // Spark 4.1.3 returns the existing root in
LogicalQueryStageStrategy and then calls
- // setLogicalLink from SparkStrategies.plan:
- // scalastyle:off line.size.limit
- //
https://github.com/apache/spark/blob/v4.1.3/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStageStrategy.scala#L64-L65
- //
https://github.com/apache/spark/blob/v4.1.3/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala#L78-L87
- // scalastyle:on line.size.limit
- op
case op: CometExec =>
- if (op.originalPlan.logicalLink.isEmpty) {
- op.unsetTagValue(SparkPlan.LOGICAL_PLAN_TAG)
- op.unsetTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG)
- } else {
- op.originalPlan.logicalLink.foreach(op.setLogicalLink)
+ // Spark links a reused physical root directly to the current AQE
stage. A descendant
Review Comment:
This drops the version-pinned links to `LogicalQueryStageStrategy` and
`SparkStrategies.plan` that were added during #5483 review. Could we keep them?
A link to `AdaptiveSparkPlanExec.replaceWithQueryStagesInLogicalPlan` would
also show why a stale inherited link makes the stage replacement a silent
no-op. Also, a short note that `getTagValue(LOGICAL_PLAN_TAG).isDefined` here
means the direct link is this stage would help.
##########
dev/diffs/3.5.9.diff:
##########
@@ -1706,20 +1706,51 @@ index 418ca3430bb..eb8267192f8 100644
withTempPath { path =>
val dir = path.getCanonicalPath
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/LogicalPlanTagInSparkPlanSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/LogicalPlanTagInSparkPlanSuite.scala
-index 743ec41dbe7..9f30d6c8e04 100644
+index 743ec41dbe7..aa549210fae 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/LogicalPlanTagInSparkPlanSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/LogicalPlanTagInSparkPlanSuite.scala
-@@ -53,6 +53,10 @@ class LogicalPlanTagInSparkPlanSuite extends
TPCDSQuerySuite with DisableAdaptiv
+@@ -22,6 +22,7 @@ import scala.reflect.ClassTag
+ import org.apache.spark.sql.TPCDSQuerySuite
+ import
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
Complete, Final}
+ import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Generate,
Join, LocalRelation, LogicalPlan, Range, Sample, Union, Window, WithCTE}
++import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometFilterExec,
CometProjectExec}
+ import org.apache.spark.sql.execution.adaptive.DisableAdaptiveExecutionSuite
+ import org.apache.spark.sql.execution.aggregate.{HashAggregateExec,
ObjectHashAggregateExec, SortAggregateExec}
+ import org.apache.spark.sql.execution.columnar.{InMemoryRelation,
InMemoryTableScanExec}
+@@ -47,12 +48,19 @@ class LogicalPlanTagInSparkPlanSuite extends
TPCDSQuerySuite with DisableAdaptiv
+ }
+
+ // A scan plan tree is a plan tree that has a leaf node under zero or more
Project/Filter nodes.
+- // We may add `ColumnarToRowExec` and `InputAdapter` above the scan node
after planning.
++ // A columnar-to-row transition and `InputAdapter` may be added above the
scan after planning.
+ @scala.annotation.tailrec
+ private def isScanPlanTree(plan: SparkPlan): Boolean = plan match {
case ColumnarToRowExec(i: InputAdapter) => isScanPlanTree(i.child)
++ case CometColumnarToRowExec(i: InputAdapter) => isScanPlanTree(i.child)
Review Comment:
Are the two `CometColumnarToRowExec(i: InputAdapter)` cases needed? Without
them, `checkLogicalPlanTag` falls through `case _` to the same
`CometProjectExec`. Dropping them and leaving Spark's two original comments
unchanged would shrink all five diffs. The same applies to the other four diffs.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1108,37 +1118,30 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
// Set up logical links
newPlan = newPlan.transform {
- case op: CometExec
- if op
- .getTagValue(SparkPlan.LOGICAL_PLAN_TAG)
- .exists(_.isInstanceOf[LogicalQueryStage]) =>
- // AQE replanning reuses this physical root and links it to the
current logical stage.
- // originalPlan can still point to a subtree hidden inside that
logical leaf, which
- // AQE cannot replace in the current logical plan. Only preserve a
direct stage link,
- // not a link inherited from an ancestor.
- // On the ordinary exchange path, the exchange itself is behind a
QueryStageExec
- // leaf and is not visited by this transform.
- // Spark 4.1.3 returns the existing root in
LogicalQueryStageStrategy and then calls
- // setLogicalLink from SparkStrategies.plan:
- // scalastyle:off line.size.limit
- //
https://github.com/apache/spark/blob/v4.1.3/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStageStrategy.scala#L64-L65
- //
https://github.com/apache/spark/blob/v4.1.3/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala#L78-L87
- // scalastyle:on line.size.limit
- op
case op: CometExec =>
- if (op.originalPlan.logicalLink.isEmpty) {
- op.unsetTagValue(SparkPlan.LOGICAL_PLAN_TAG)
- op.unsetTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG)
- } else {
- op.originalPlan.logicalLink.foreach(op.setLogicalLink)
+ // Spark links a reused physical root directly to the current AQE
stage. A descendant
+ // can inherit that stage if it wraps the descendant's original
logical node.
+ val keepStage = op.logicalLink.exists {
+ case stage: LogicalQueryStage =>
+ op.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).isDefined ||
+ op.originalPlan.logicalLink.exists(stageWraps(stage, _))
+ case _ => false
+ }
+ if (!keepStage) {
+ // setLogicalLink promotes an inherited link to a direct boundary
and propagates
+ // it into children, overwriting their current stage links.
Restore each tag locally.
+ def restore(tag: TreeNodeTag[LogicalPlan]): Unit =
Review Comment:
This is the third hand-written copy of the two tags. The other two are
`MakeDateEvaluation.copyLogicalLinks` and `limitCollectProjectionInput`. They
keep the target's tag when the source has none, but `restore` clears it. Is the
difference intentional? If it is, could the three share one helper that makes
that choice explicit?
--
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]
