sunchao opened a new issue, #5482: URL: https://github.com/apache/datafusion-comet/issues/5482
## 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. The reproduction above isolates the planner invariant; it does not claim a particular query-result error, cancellation pattern, or benchmark regression. Runtime validation belongs with the accompanying fix. -- 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]
