plan twice

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

@@ -103,28 +103,27 @@ case class AdaptiveSparkPlanExec(
   // A list of physical plan rules to be applied before creation of query stages. The physical
   // plan should reach a final status of query stages (i.e., no more addition or removal of
   // Exchange nodes) after running these rules.
-  @transient private val queryStagePreparationRules: Seq[Rule[SparkPlan]] = Seq(
-    RemoveRedundantProjects,
-    // For cases like `df.repartition(a, b).select(c)`, there is no distribution requirement for
-    // the final plan, but we do need to respect the user-specified repartition. Here we ask
-    // `EnsureRequirements` to not optimize out the user-specified repartition-by-col to work
-    // around this case.
-    EnsureRequirements(optimizeOutRepartition = requiredDistribution.isDefined),
-    RemoveRedundantSorts,
-    DisableUnnecessaryBucketedScan
-  ) ++ context.session.sessionState.queryStagePrepRules
-
-  // This list rules are applied between queryStagePreparationRules and estimate physical plan cost
-  // so that we can support introduce extra shuffle
-  @transient private val optimizeSkewedJoinWithExtraShuffleRules: Seq[Rule[SparkPlan]] = {
-    val ensureExtraShuffleRule = if (forceOptimizeSkewedJoin) {
-      // Add the EnsureRequirements rule here and don't optimize out repartition so that we can
-      // ensure the output partitioning of OptimizeSkewedJoin is always expected.
-      Seq(EnsureRequirements(optimizeOutRepartition = requiredDistribution.isDefined))
+  private def queryStagePreparationRules(
+      optimizeSkewedJoin: Boolean = false): Seq[Rule[SparkPlan]] = {
+    val optimizeSkewedJoinRules = if (optimizeSkewedJoin) {
+      Seq(OptimizeSkewedJoin,
+        // Add the EnsureRequirements rule here and don't optimize out repartition so that we can
+        // ensure the output partitioning of OptimizeSkewedJoin is always expected.
+        EnsureRequirements(optimizeOutRepartition = requiredDistribution.isDefined))
     } else {
       Nil
     }
-    Seq(OptimizeSkewedJoin) ++ ensureExtraShuffleRule
+
+    Seq(
+      RemoveRedundantProjects,
+      // For cases like `df.repartition(a, b).select(c)`, there is no distribution requirement for
+      // the final plan, but we do need to respect the user-specified repartition. Here we ask
+      // `EnsureRequirements` to not optimize out the user-specified repartition-by-col to work
+      // around this case.
+      EnsureRequirements(optimizeOutRepartition = requiredDistribution.isDefined),
+      RemoveRedundantSorts,
+      DisableUnnecessaryBucketedScan
+    ) ++ optimizeSkewedJoinRules ++ context.session.sessionState.queryStagePrepRules
   }
 
   // A list of physical optimizer rules to be applied to a new stage before its execution. These
@@ -200,7 +199,7 @@ case class AdaptiveSparkPlanExec(
 
   @transient val initialPlan = context.session.withActive {
     applyPhysicalRules(
-      inputPlan, queryStagePreparationRules, Some((planChangeLogger, "AQE Preparations")))
+      inputPlan, queryStagePreparationRules(), Some((planChangeLogger, "AQE Preparations")))
   }
 
   @volatile private var currentPhysicalPlan = initialPlan
@@ -332,7 +331,7 @@ case class AdaptiveSparkPlanExec(
               }
             }._1
 
-        if (!newPhysicalPlan.fastEquals(currentPhysicalPlan)) {
+        if (newPhysicalPlan.ne(currentPhysicalPlan)) {
           logOnLevel(s"Plan changed from\n$currentPhysicalPlan\nto\n$newPhysicalPlan")
           cleanUpTempTags(newPhysicalPlan)
           currentPhysicalPlan = newPhysicalPlan
@@ -667,6 +666,25 @@ case class AdaptiveSparkPlanExec(
     logicalPlan
   }
 
+  private def isFinalStage(sparkPlan: SparkPlan): Boolean = {
+    sparkPlan match {
+      // avoid top level node is Exchange
+      case _: Exchange => false
+      case plan =>
+        // Plan is regarded as a final plan iff all shuffle nodes are wrapped inside query stage
+        // and all query stages are materialized.
+        plan.find {
+          case p if p.children.exists(
+            child => child.isInstanceOf[Exchange] || child.isInstanceOf[ReusedExchangeExec]) =>
+            p match {
+              case stage: QueryStageExec if stage.isMaterialized => false
+              case _ => true
+            }
+          case _ => false
+        }.isEmpty
+    }
+  }
+
   /**
    * Re-optimize and run physical planning on the current logical plan based on the latest stats.
    */
@@ -675,11 +693,25 @@ case class AdaptiveSparkPlanExec(
     logicalPlan.invalidateStatsCache()
     val optimized = optimizer.execute(logicalPlan)
     val sparkPlan = context.session.sessionState.planner.plan(ReturnAnswer(optimized)).next()
+
     val optimizedPhysicalPlan = applyPhysicalRules(
       sparkPlan,
-      preprocessingRules ++ queryStagePreparationRules,
+      preprocessingRules ++ queryStagePreparationRules(),
       Some((planChangeLogger, "AQE Replanning")))
 
+    val optimizedWithSkewedJoin = applyPhysicalRules(
+      sparkPlan,
+      preprocessingRules ++ queryStagePreparationRules(true),
+      Some((planChangeLogger, "AQE Replanning With Optimize Skewed Join")))
+
+    // respect the requiredDistribution for final stage
+    val validatedWithSkewedJoin =
+      checkDistribution(
+        optimizedWithSkewedJoin,
+        optimizedPhysicalPlan,
+        isFinalStage(optimizedWithSkewedJoin),
+        OptimizeSkewedJoin.ruleName)
+
     // When both enabling AQE and DPP, `PlanAdaptiveDynamicPruningFilters` rule will
     // add the `BroadcastExchangeExec` node manually in the DPP subquery,
     // not through `EnsureRequirements` rule. Therefore, when the DPP subquery is complicated
@@ -693,23 +725,10 @@ case class AdaptiveSparkPlanExec(
       case _ => plan
     }
 
-    val optimizedWithSkewedJoin = applyPhysicalRules(
-      optimizedPhysicalPlan,
-      optimizeSkewedJoinWithExtraShuffleRules,
-      Some((planChangeLogger, "AQE Optimize Skewed Join With Extra Shuffle"))
-    )
-    val validatedWithSkewedJoin =
-      checkDistribution(
-        optimizedWithSkewedJoin,
-        optimizedPhysicalPlan,
-        isFinalStage(optimizedWithSkewedJoin),
-        OptimizeSkewedJoin.ruleName)
-
-    // here are three reasons if validatedWithSkewedJoin is equal to optimizedPhysicalPlan:
+    // here are two reasons if validatedWithSkewedJoin is equal to optimizedPhysicalPlan:
     // 1. no skewed join optimized
-    // 2. optimize skewed join introduce extra shuffle and force optimize is disabled
-    // 3. optimize skewed join change final stage output partitioning
-    val newPhysicalPlans = if (validatedWithSkewedJoin.fastEquals(optimizedPhysicalPlan)) {
+    // 2. optimize skewed join doesn't satisfy requiredDistribution for final stage
+    val newPhysicalPlans = if (optimizedPhysicalPlan.fastEquals(validatedWithSkewedJoin)) {
       updateBroadcastExchange(optimizedPhysicalPlan) :: Nil
     } else {
       updateBroadcastExchange(optimizedPhysicalPlan) ::
@@ -718,25 +737,6 @@ case class AdaptiveSparkPlanExec(
     (newPhysicalPlans, optimized)
   }
 
-  private def isFinalStage(sparkPlan: SparkPlan): Boolean = {
-    sparkPlan match {
-      // avoid top level node is Exchange
-      case _: Exchange => false
-      case plan =>
-        // Plan is regarded as a final plan iff all shuffle nodes are wrapped inside query stage
-        // and all query stages are materialized.
-        plan.find {
-          case p if p.children.exists(
-            child => child.isInstanceOf[Exchange] || child.isInstanceOf[ReusedExchangeExec]) =>
-            p match {
-              case stage: QueryStageExec if stage.isMaterialized => false
-              case _ => true
-            }
-          case _ => false
-        }.isEmpty
-    }
-  }
-
   /**
    * Recursively set `TEMP_LOGICAL_PLAN_TAG` for the current `plan` node.
    */