simplify QueryStage

Author: cloud-fan

core/src/main/scala/org/apache/spark/MapOutputStatistics.scala

@@ -25,4 +25,3 @@ package org.apache.spark
  *   (may be inexact due to use of compressed map statuses)
  */
 private[spark] class MapOutputStatistics(val shuffleId: Int, val bytesByPartitionId: Array[Long])
-  extends Serializable

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

@@ -94,11 +94,12 @@ class QueryExecution(
    * row format conversions as needed.
    */
   protected def prepareForExecution(plan: SparkPlan): SparkPlan = {
-    if (sparkSession.sessionState.conf.adaptiveExecutionEnabled) {
-      adaptivePreparations.foldLeft(plan) { case (sp, rule) => rule.apply(sp)}
+    val rules = if (sparkSession.sessionState.conf.adaptiveExecutionEnabled) {
+      adaptivePreparations
     } else {
-      preparations.foldLeft(plan) { case (sp, rule) => rule.apply(sp)}
+      preparations
     }
+    rules.foldLeft(plan) { case (sp, rule) => rule.apply(sp)}
   }
 
   /** A sequence of rules that will be applied in order to the physical plan before execution. */
@@ -109,8 +110,7 @@ class QueryExecution(
     ReuseExchange(sparkSession.sessionState.conf),
     ReuseSubquery(sparkSession.sessionState.conf))
 
-  // With adaptive execution, whole stage codegen will be planned inside `QueryStage`, so we exclude
-  // `CollapseCodegenStages` here.
+  // With adaptive execution, whole stage codegen will be done inside `QueryStageExecutor`.
   protected def adaptivePreparations: Seq[Rule[SparkPlan]] = Seq(
     PlanSubqueries(sparkSession),
     EnsureRequirements(sparkSession.sessionState.conf),
@@ -119,7 +119,7 @@ class QueryExecution(
     // PlanQueryStage needs to be the last rule because it divides the plan into multiple sub-trees
     // by inserting leaf node QueryStageInput. Transforming the plan after applying this rule will
     // only transform node in a sub-tree.
-    PlanQueryStage)
+    PlanQueryStage(sparkSession.sessionState.conf))
 
   def simpleString: String = withRedaction {
     val concat = new StringConcat()

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

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.execution
 
 import org.apache.spark.annotation.DeveloperApi
-import org.apache.spark.sql.execution.adaptive.QueryStageInput
+import org.apache.spark.sql.execution.adaptive.QueryStage
 import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
 import org.apache.spark.sql.execution.metric.SQLMetricInfo
 import org.apache.spark.sql.internal.SQLConf
@@ -53,7 +53,7 @@ private[execution] object SparkPlanInfo {
   def fromSparkPlan(plan: SparkPlan): SparkPlanInfo = {
     val children = plan match {
       case ReusedExchangeExec(_, child) => child :: Nil
-      case i: QueryStageInput => i.childStage :: Nil
+      case stage: QueryStage => stage.finalPlan :: Nil
       case _ => plan.children ++ plan.subqueries
     }
     val metrics = plan.metrics.toSeq.map { case (key, metric) =>

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

@@ -20,6 +20,7 @@ package org.apache.spark.sql.execution.adaptive
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, Exchange, ReusedExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.internal.SQLConf
 
 /**
  * Divide the spark plan into multiple QueryStages. For each Exchange in the plan, it adds a
@@ -29,27 +30,27 @@ import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, Exchange,
  * is a leaf node. Transforming the plan after applying this rule will only transform node in a
  * sub-tree.
  */
-object PlanQueryStage extends Rule[SparkPlan] {
+case class PlanQueryStage(conf: SQLConf) extends Rule[SparkPlan] {
 
   def apply(plan: SparkPlan): SparkPlan = {
     val exchangeToQueryStage = new java.util.IdentityHashMap[Exchange, QueryStage]
     val newPlan = plan.transformUp {
       case e: ShuffleExchangeExec =>
         val queryStage = ShuffleQueryStage(e)
         exchangeToQueryStage.put(e, queryStage)
-        ShuffleQueryStageInput(queryStage, e.output)
+        ShuffleQueryStageReaderExec(queryStage, queryStage.output)
       case e: BroadcastExchangeExec =>
         val queryStage = BroadcastQueryStage(e)
         exchangeToQueryStage.put(e, queryStage)
-        BroadcastQueryStageInput(queryStage, e.output)
+        BroadcastQueryStageReaderExec(queryStage, queryStage.output)
       // The `ReusedExchangeExec` was added in the rule `ReuseExchange`, via transforming up the
       // query plan. This rule also transform up the query plan, so when we hit `ReusedExchangeExec`
       // here, the exchange being reused must already be hit before and there should be an entry
       // for it in `exchangeToQueryStage`.
       case e: ReusedExchangeExec =>
         exchangeToQueryStage.get(e.child) match {
-          case q: ShuffleQueryStage => ShuffleQueryStageInput(q, e.output)
-          case q: BroadcastQueryStage => BroadcastQueryStageInput(q, e.output)
+          case q: ShuffleQueryStage => ShuffleQueryStageReaderExec(q, e.output)
+          case q: BroadcastQueryStage => BroadcastQueryStageReaderExec(q, e.output)
         }
     }
     ResultQueryStage(newPlan)