apache / apache/datafusion-comet
CometExecRule overwrites direct AQE LogicalQueryStage links during replanning
- Dominant language
- Scala
- Stars
- 1.3k
- Forks
- 373
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 198
Description
## Describe the bug
`CometExecRule` can overwrite the current AQE logical-stage link on a reused native operator with an older link from `originalPlan`.
Spark can represent a native final aggregate above a shuffle stage as a `LogicalQueryStage`. During replanning, Spark's planner reuses that physical aggregate and gives it a direct `SparkPlan.LOGICAL_PLAN_TAG` pointing to the current logical-stage object. Comet's subsequent logical-link repair unconditionally restores `originalPlan.logicalLink`, or clears the tags when the original link is absent.
The original logical aggregate is now stored inside a logical-stage leaf, rather than appearing as a node in the active logical tree. Restoring that older link breaks the correspondence between the current physical root and the current logical stage. A new exchange above the aggregate can inherit the stale link, but Spark's identity-based logical-stage replacement cannot find that old logical node in the active tree.
The unconditional restoration is present on OSS `main` at [`a2c6bd4b930174cf81e0ad2857d2feb422e081c5`](https://github.com/apache/datafusion-comet/blob/a2c6bd4b930174cf81e0ad2857d2feb422e081c5/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala#L647-L678).
## Steps to reproduce
The following planner-level regression can be added inside `CometExecRuleSuite`, using its existing `createSparkPlan` and `applyCometExecRule` helpers. It creates a native final aggregate above a shuffle query stage, then passes a `LogicalQueryStage` through Spark's actual planner twice. AQE execution is disabled only to construct the initial plan; the test explicitly exercises its planner reuse sequence.
```scala
import org.apache.spark.sql.execution.adaptive.{LogicalQueryStage, ShuffleQueryStageExec}
import org.apache.spark.sql.internal.SQLConf
test("CometExecRule preserves the current direct AQE logical link") {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
CometConf.COMET_SHUFFLE_MODE.key -> "native",
CometConf.COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST.key -> "Range") {
val plan = createSparkPlan(
spark,
"SELECT id % 3 AS k, SUM(id) AS total FROM range(0, 100, 1, 2) GROUP BY id % 3")
val aggregate = applyCometExecRule(plan).asInstanceOf[CometHashAggregateExec]
val shuffle = aggregate.child.asInstanceOf[CometShuffleExchangeExec]
val originalLogicalPlan = aggregate.originalPlan.logicalLink.get
var current: SparkPlan = aggregate.withNewChildren(
Seq(ShuffleQueryStageExec(0, shuffle, shuffle.canonicalized)))
(1 to 2).foreach { _ =>
val logicalStage = LogicalQueryStage(originalLogicalPlan, current)
val replanned = spark.sessionState.planner.plan(logicalStage).next()
assert(replanned eq current)
assert(replanned.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).exists(_ eq logicalStage))
current = applyCometExecRule(replanned)
assert(current.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).exists(_ eq logicalStage))
}
}
}
```
The final assertion specifies the required invariant: applying Comet's rule must not replace the direct link that Spark just assigned. The current repair branch instead restores `originalLogicalPlan`. The same preservation requirement applies when `originalPlan` has an inherited link or no logical link at all.
## Expected behavior
Preserve an existing direct `LogicalQueryStage` link on a `CometExec` when the rule revisits it. The tag must continue to reference the exact current logical-stage object after repeated replanning.
Keep the existing repair and clearing behavior for ordinary direct links and inherited links. In particular, merely inheriting a `LogicalQueryStage` link from an ancestor must not trigger preservation. Leave shuffle and broadcast exchange link handling unchanged, including the empty-link invariant from #323.
## Additional context
This follows public Spark planner behavior: [`LogicalQueryStageStrategy`](https://github.com/apache/spark/blob/e221b56be7b6d9e48e107fc4d1cf0c15f02700f8/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStageStrategy.scala#L64-L65) returns the existing physical plan, and [`SparkStrategies.plan`](https://github.com/apache/spark/blob/e221b56be7b6d9e48e107fc4d1cf0c15f02700f8/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala#L78-L87) sets its direct logical link. [`AdaptiveSparkPlanExec.replaceWithQueryStagesInLogicalPlan`](https://github.com/apache/spark/blob/e221b56be7b6d9e48e107fc4d1cf0c15f02700f8/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala#L759-L782) replaces logical nodes by object identity.
A narrow fix is to skip logical-link repair only when `getTagValue(SparkPlan.LOGICAL_PLAN_TAG)` is a `LogicalQueryStage`. No Spark changes, configuration changes, or changes to native aggregate execution are needed.
Reproduced locally on OSS Spark 4.1.3 / Scala 2.13.17 with JDK 17 and a native library built from this OSS checkout. With the production rule unchanged, the direct-link regression fails at the assertion that Comet preserves the current logical-stage object. A generated-data SQL regression also returns the correct result but finishes with zero native broadcast hash joins instead of the expected two. With the narrow preservation guard, both regressions pass, along with the rest of `CometExecRuleSuite` and the existing shuffle logical-link test (33 tests total).
This establishes an OSS planner and native broadcast adaptation regression. It does not claim a particular cancellation pattern or benchmark result.
Contributor guide
Research direction
Start in spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala at the logical-link repair branch, then add the regression to CometExecRuleSuite using its existing helpers. Run the planner-level test and the suite; done means the exact current LogicalQueryStage link survives repeated replanning while ordinary, inherited, and exchange-link behavior remains unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- scala
- Domain
- backend, testing
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100