CometExecRule overwrites direct AQE LogicalQueryStage links during replanning

[sunchao]

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.

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.

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 returns the existing physical plan, and SparkStrategies.plan sets its direct logical link. AdaptiveSparkPlanExec.replaceWithQueryStagesInLogicalPlan 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.