sunchao commented on code in PR #7357:
URL: https://github.com/apache/datafusion-comet/pull/7357#discussion_r4433047286
##########
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:
Restored the two Spark 4.1.3 links and added the identity check in
`replaceWithQueryStagesInLogicalPlan`. If the linked node is absent from the
current logical tree, that transform leaves the tree unchanged. I also added a
note at the guard: Spark checks the direct tag first, so once `logicalLink` has
matched a stage, a present direct tag identifies that same stage.
[a750c1bf8b](https://github.com/apache/datafusion-comet/pull/7357/commits/a750c1bf8b2d141a78fedcde99cbb6c93f38c2ce)
##########
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:
Yes, the difference is intentional. When restoring from `originalPlan`, a
missing source tag means the corresponding tag on the Comet node is stale and
must be cleared. A fallback plan or rebuilt projection can already have a valid
link of its own, so those callers keep it when the source has none. All three
paths now call `copyLogicalLinks` with an explicit `clearMissing` argument. The
full planner-rule suite and the AQE broadcast regression passed on Spark 3.5.9
and 4.1.3.
[a750c1bf8b](https://github.com/apache/datafusion-comet/pull/7357/commits/a750c1bf8b2d141a78fedcde99cbb6c93f38c2ce)
--
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]
