[SPARK-28583][SQL] Subqueries should not call onUpdatePlan in Adaptive Query Execution

## What changes were proposed in this pull request?

Subqueries do not have their own execution id, thus when calling `AdaptiveSparkPlanExec.onUpdatePlan`, it will actually get the `QueryExecution` instance of the main query, which is wasteful and problematic. It could cause issues like stack overflow or dead locks in some circumstances.

This PR fixes this issue by making `AdaptiveSparkPlanExec` compare the `QueryExecution` object retrieved by current execution ID against the `QueryExecution` object from which this plan is created, and only update the UI when the two instances are the same.

## How was this patch tested?

Manual tests on TPC-DS queries.

Closes #25316 from maryannxue/aqe-updateplan-fix.

Authored-by: maryannxue <[email protected]>
Signed-off-by: herman <[email protected]>

Author: maryannxue

sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala

@@ -118,7 +118,7 @@ class QueryExecution(
   protected def preparations: Seq[Rule[SparkPlan]] = Seq(
     // `AdaptiveSparkPlanExec` is a leaf node. If inserted, all the following rules will be no-op
     // as the original plan is hidden behind `AdaptiveSparkPlanExec`.
-    InsertAdaptiveSparkPlan(sparkSession),
+    InsertAdaptiveSparkPlan(sparkSession, this),
     PlanSubqueries(sparkSession),
     EnsureRequirements(sparkSession.sessionState.conf),
     ApplyColumnarRulesAndInsertTransitions(sparkSession.sessionState.conf,

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala

@@ -61,7 +61,8 @@ case class AdaptiveSparkPlanExec(
     initialPlan: SparkPlan,
     @transient session: SparkSession,
     @transient subqueryMap: Map[Long, ExecSubqueryExpression],
-    @transient stageCache: TrieMap[SparkPlan, QueryStageExec])
+    @transient stageCache: TrieMap[SparkPlan, QueryStageExec],
+    @transient queryExecution: QueryExecution)
   extends LeafExecNode {
 
   @transient private val lock = new Object()
@@ -118,8 +119,15 @@ case class AdaptiveSparkPlanExec(
     if (isFinalPlan) {
       currentPhysicalPlan.execute()
     } else {
+      // Make sure we only update Spark UI if this plan's `QueryExecution` object matches the one
+      // retrieved by the `sparkContext`'s current execution ID. Note that sub-queries do not have
+      // their own execution IDs and therefore rely on the main query to update UI.
       val executionId = Option(
-        session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)).map(_.toLong)
+        session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)).flatMap { idStr =>
+        val id = idStr.toLong
+        val qe = SQLExecution.getQueryExecution(id)
+        if (qe.eq(queryExecution)) Some(id) else None
+      }
       var currentLogicalPlan = currentPhysicalPlan.logicalLink.get
       var result = createQueryStages(currentPhysicalPlan)
       val events = new LinkedBlockingQueue[StageMaterializationEvent]()
@@ -171,10 +179,11 @@ case class AdaptiveSparkPlanExec(
       currentPhysicalPlan = applyPhysicalRules(result.newPlan, queryStageOptimizerRules)
       currentPhysicalPlan.setTagValue(SparkPlan.LOGICAL_PLAN_TAG, currentLogicalPlan)
       isFinalPlan = true
+
+      val ret = currentPhysicalPlan.execute()
       logDebug(s"Final plan: $currentPhysicalPlan")
       executionId.foreach(onUpdatePlan)
-
-      currentPhysicalPlan.execute()
+      ret
     }
   }
 

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/InsertAdaptiveSparkPlan.scala

@@ -35,14 +35,18 @@ import org.apache.spark.sql.types.StructType
  *
  * Note that this rule is stateful and thus should not be reused across query executions.
  */
-case class InsertAdaptiveSparkPlan(session: SparkSession) extends Rule[SparkPlan] {
+case class InsertAdaptiveSparkPlan(
+    session: SparkSession,
+    queryExecution: QueryExecution) extends Rule[SparkPlan] {
 
   private val conf = session.sessionState.conf
 
   // Exchange-reuse is shared across the entire query, including sub-queries.
   private val stageCache = new TrieMap[SparkPlan, QueryStageExec]()
 
-  override def apply(plan: SparkPlan): SparkPlan = plan match {
+  override def apply(plan: SparkPlan): SparkPlan = applyInternal(plan, queryExecution)
+
+  private def applyInternal(plan: SparkPlan, qe: QueryExecution): SparkPlan = plan match {
     case _: ExecutedCommandExec => plan
     case _ if conf.adaptiveExecutionEnabled && supportAdaptive(plan) =>
       try {
@@ -54,7 +58,7 @@ case class InsertAdaptiveSparkPlan(session: SparkSession) extends Rule[SparkPlan
           session.sessionState.conf, subqueryMap)
         val newPlan = AdaptiveSparkPlanExec.applyPhysicalRules(plan, preparations)
         logDebug(s"Adaptive execution enabled for plan: $plan")
-        AdaptiveSparkPlanExec(newPlan, session, subqueryMap, stageCache)
+        AdaptiveSparkPlanExec(newPlan, session, subqueryMap, stageCache, qe)
       } catch {
         case SubqueryAdaptiveNotSupportedException(subquery) =>
           logWarning(s"${SQLConf.ADAPTIVE_EXECUTION_ENABLED.key} is enabled " +
@@ -120,7 +124,7 @@ case class InsertAdaptiveSparkPlan(session: SparkSession) extends Rule[SparkPlan
     val queryExec = new QueryExecution(session, plan)
     // Apply the same instance of this rule to sub-queries so that sub-queries all share the
     // same `stageCache` for Exchange reuse.
-    val adaptivePlan = this.apply(queryExec.sparkPlan)
+    val adaptivePlan = this.applyInternal(queryExec.sparkPlan, queryExec)
     if (!adaptivePlan.isInstanceOf[AdaptiveSparkPlanExec]) {
       throw SubqueryAdaptiveNotSupportedException(plan)
     }