[SPARK-9860][SQL] Join: Determine the join strategy (broadcast join or shuffle join) at runtime

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

@@ -27,6 +27,7 @@ import org.apache.spark.sql.catalyst.analysis.UnsupportedOperationChecker
 import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, ReturnAnswer}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.DateTimeUtils
+import org.apache.spark.sql.execution.adaptive.QueryFragmentTransformer
 import org.apache.spark.sql.execution.command.{DescribeTableCommand, ExecutedCommandExec}
 import org.apache.spark.sql.execution.exchange.{EnsureRequirements, ReuseExchange}
 import org.apache.spark.sql.internal.SQLConf
@@ -79,7 +80,12 @@ class QueryExecution(val sparkSession: SparkSession, val logical: LogicalPlan) {
 
   // executedPlan should not be used to initialize any SparkPlan. It should be
   // only used for execution.
-  lazy val executedPlan: SparkPlan = prepareForExecution(sparkPlan)
+  lazy val executedPlan: SparkPlan =
+    if (sparkSession.sessionState.conf.adaptiveExecution2Enabled) {
+      prepareForAdaptiveExecution(sparkPlan)
+    } else {
+      prepareForExecution(sparkPlan)
+    }
 
   /** Internal version of the RDD. Avoids copies and has no schema */
   lazy val toRdd: RDD[InternalRow] = executedPlan.execute()
@@ -100,6 +106,17 @@ class QueryExecution(val sparkSession: SparkSession, val logical: LogicalPlan) {
     CollapseCodegenStages(sparkSession.sessionState.conf),
     ReuseExchange(sparkSession.sessionState.conf))
 
+  protected def prepareForAdaptiveExecution(plan: SparkPlan): SparkPlan = {
+    preparationsForAdaptive.foldLeft(plan) { case (sp, rule) => rule.apply(sp) }
+  }
+
+  protected def preparationsForAdaptive: Seq[Rule[SparkPlan]] = Seq(
+    python.ExtractPythonUDFs,
+    PlanSubqueries(sparkSession),
+    EnsureRequirements(sparkSession.sessionState.conf),
+    ReuseExchange(sparkSession.sessionState.conf),
+    QueryFragmentTransformer(sparkSession.sessionState.conf))
+
   protected def stringOrError[A](f: => A): String =
     try f.toString catch { case e: Throwable => e.toString }
 

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

@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.adaptive
+
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.execution.{LeafExecNode, ShuffledRowRDD, SparkPlan}
+
+/**
+ * FragmentInput is the leaf node of parent fragment that connect with an child fragment.
+ */
+case class FragmentInput(@transient childFragment: QueryFragment) extends LeafExecNode {
+
+  private[this] var optimized: Boolean = false
+
+  private[this] var inputPlan: SparkPlan = null
+
+  private[this] var shuffledRdd: ShuffledRowRDD = null
+
+  override def output: Seq[Attribute] = inputPlan.output
+
+  private[sql] def setOptimized() = {
+    this.optimized = true
+  }
+
+  private[sql] def isOptimized(): Boolean = this.optimized
+
+  private[sql] def setShuffleRdd(shuffledRdd: ShuffledRowRDD) = {
+    this.shuffledRdd = shuffledRdd
+  }
+
+  private[sql] def setInputPlan(inputPlan: SparkPlan) = {
+    this.inputPlan = inputPlan
+  }
+
+  override protected def doExecute(): RDD[InternalRow] = {
+    if (shuffledRdd != null) {
+      shuffledRdd
+    } else {
+      inputPlan.execute()
+    }
+  }
+
+  override def simpleString: String = "FragmentInput"
+
+  override def innerChildren: Seq[SparkPlan] = inputPlan :: Nil
+}