comphead commented on code in PR #7357:
URL: https://github.com/apache/datafusion-comet/pull/7357#discussion_r4486202581
##########
spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala:
##########
@@ -1429,7 +1429,21 @@ class CometExecRuleSuite extends CometTestBase {
}
}
- test("CometExecRule repairs ordinary and inherited logical links from the
original plan") {
+ test("CometExecRule preserves an inherited AQE stage for its original
logical subtree") {
+ withAdaptiveAggregateConf {
+ val aggregate = createAdaptiveAggregate()
+ val original = aggregate.originalPlan.logicalLink.get
+ val stage = LogicalQueryStage(LogicalQueryStage(original, aggregate),
aggregate)
+ aggregate.unsetTagValue(SparkPlan.LOGICAL_PLAN_TAG)
+ aggregate.setTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG, stage)
+
+ val repaired = applyCometExecRule(aggregate)
+ assert(repaired.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).isEmpty)
+ assert(repaired.logicalLink.exists(_ eq stage))
+ }
+ }
Review Comment:
This test repeats the setup of the two neighbours
(`createAdaptiveAggregate`, unset the direct tag, set an inherited stage) and
only covers an original that has a direct tag. The shape that triggers the bug
is a descendant whose original has only an inherited tag, since
`SparkStrategies.plan` gives every non-root node of a planned subtree an
inherited link. The code reads `logicalLink` so I expect it to behave the same,
but nothing asserts it. Could this loop over the same three `originalTag`
values as the other two tests, or become one more case in "repairs stale
logical links" where the stage wraps the original instead of being a clone?
##########
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:
The helper comment says `setLogicalLink` promotes an inherited link and
rewrites descendants, but this branch and the identical
`CometBroadcastExchangeExec` branch below still call it. Could both use
`copyLogicalLinks(op.originalPlan, op, clearMissing = true)`, or do exchanges
need the propagation to children? The same `setLogicalLink` pattern remains in
`EliminateRedundantTransitions.scala` (around line 120),
`CometCacheBuildRule.scala` line 64 and the `sparkFallback` of
`CometNativeScanExec` and `CometIcebergNativeScanExec`. I could not build a
case where these get an inherited-only source, so this is a question about
consistency. Sharing the helper there would need it outside `private[rules]`.
--
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]
