sunchao commented on code in PR #7357:
URL: https://github.com/apache/datafusion-comet/pull/7357#discussion_r4515408177
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1108,37 +1131,29 @@ 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:
+ case op: CometExec =>
+ // Spark reuses the physical root and links it directly to the
current AQE stage.
+ // A descendant can inherit that stage if it wraps the descendant's
original node.
+ // AQE replaces the linked node by identity. A stale link leaves the
plan unchanged.
// 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
+ //
https://github.com/apache/spark/blob/v4.1.3/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala#L778-L783
// 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)
+ val keepStage = op.logicalLink.exists {
+ case stage: LogicalQueryStage =>
+ // logicalLink prefers the direct tag, so if present it points
to this stage.
+ op.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).isDefined ||
+ op.originalPlan.logicalLink.exists(stageWraps(stage, _))
+ case _ => false
+ }
+ if (!keepStage) {
+ copyLogicalLinks(op.originalPlan, op, clearMissing = true)
}
op
case op: CometShuffleExchangeExec =>
- // Original Spark shuffle exchange operator might have empty logical
link.
- // But the `setLogicalLink` call above on downstream operator of
- // `CometShuffleExchangeExec` will set its logical link to the
downstream
- // operators which cause AQE behavior to be incorrect. So we need to
unset
- // the logical link here.
+ // An exchange without an original logical link can inherit an
unrelated one from
+ // its parent. Clear it so AQE does not associate the exchange with
that parent.
if (op.originalPlan.logicalLink.isEmpty) {
op.unsetTagValue(SparkPlan.LOGICAL_PLAN_TAG)
op.unsetTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG)
Review Comment:
Good catch. On both exchanges, `setLogicalLink` makes the link direct on the
exchange and passes it down until it reaches a child with its own direct link.
A node-only copy would preserve the original tag kind and leave the children’s
tags alone, but it could still overwrite a newer AQE stage on a reused
exchange. Whether any exchange children rely on propagation is still open.
The scan fallbacks can also promote an inherited link on a leaf. The
transition and cache rules wrap an existing child, so child propagation matters
there too. I’m keeping those paths in the [existing AQE link
audit](https://github.com/apache/datafusion-comet/issues/6034). An exchange or
DPP plan that shows which node should own the link would tell us which
operation to use.
--
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]
