Fix Exchange and initial code gen attempt.

Author: hvanhovell

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

@@ -23,16 +23,26 @@ import org.apache.spark.broadcast
 import org.apache.spark.rdd.{EmptyRDD, RDD}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext
+import org.apache.spark.sql.execution.metric.SQLMetrics
 import org.apache.spark.util.ThreadUtils
 
 /**
  * A broadcaster collects transforms and broadcasts the result of an underlying spark plan.
  *
  * TODO whole stage codegen.
  */
-case class Broadcast(f: Iterable[InternalRow] => Any, child: SparkPlan) extends UnaryNode {
+case class Broadcast(
+    f: Iterable[InternalRow] => Any,
+    child: SparkPlan)
+  extends UnaryNode with CodegenSupport {
+
   override def output: Seq[Attribute] = child.output
 
+  override private[sql] lazy val metrics = Map(
+    "numRows" -> SQLMetrics.createLongMetric(sparkContext, "number of rows")
+  )
+
   val timeout: Duration = {
     val timeoutValue = sqlContext.conf.broadcastTimeout
     if (timeoutValue < 0) {
@@ -66,7 +76,11 @@ case class Broadcast(f: Iterable[InternalRow] => Any, child: SparkPlan) extends
     Await.result(future, timeout)
   }
 
+  override def upstream(): RDD[InternalRow] = {
+    child.asInstanceOf[CodegenSupport].upstream()
+  }
 
+  override def doProduce(ctx: CodegenContext): String = ""
 
   override protected def doPrepare(): Unit = {
     // Materialize the relation.
@@ -86,9 +100,9 @@ case class Broadcast(f: Iterable[InternalRow] => Any, child: SparkPlan) extends
 object Broadcast {
   def broadcastRelation[T](plan: SparkPlan): broadcast.Broadcast[T] = plan match {
     case builder: Broadcast => builder.broadcastRelation
-    case _ => sys.error("The given plan is not a Broadcaster")
+    case _ => sys.error("The given plan is not a Broadcast")
   }
 
   private[execution] val executionContext = ExecutionContext.fromExecutorService(
     ThreadUtils.newDaemonCachedThreadPool("build-broadcast", 128))
-}
+}

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

@@ -385,19 +385,24 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
     // Ensure that the operator's children satisfy their output distribution requirements:
     children = children.zip(requiredChildDistributions).map { case (child, distribution) =>
       distribution match {
-        case BroadcastDistribution(f) =>
-          Broadcast(f, child)
         case _ if child.outputPartitioning.satisfies(distribution) =>
           child
+        case BroadcastDistribution(f) =>
+          Broadcast(f, child)
         case _ =>
           Exchange(createPartitioning(distribution, defaultNumPreShufflePartitions), child)
       }
     }
 
     // If the operator has multiple children and specifies child output distributions (e.g. join),
     // then the children's output partitionings must be compatible:
+    def requireCompatiblePartitioning(distribution: Distribution): Boolean = distribution match {
+      case UnspecifiedDistribution => false
+      case BroadcastDistribution(_) => false
+      case _ => true
+    }
     if (children.length > 1
-        && requiredChildDistributions.toSet != Set(UnspecifiedDistribution)
+        && requiredChildDistributions.exists(requireCompatiblePartitioning)
         && !Partitioning.allCompatible(children.map(_.outputPartitioning))) {
 
       // First check if the existing partitions of the children all match. This means they are
@@ -434,8 +439,7 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
 
         children.zip(requiredChildDistributions).map {
           case (child, distribution) => {
-            val targetPartitioning =
-              createPartitioning(distribution, numPartitions)
+            val targetPartitioning = createPartitioning(distribution, numPartitions)
             if (child.outputPartitioning.guarantees(targetPartitioning)) {
               child
             } else {

sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashJoin.scala

@@ -42,15 +42,19 @@ case class BroadcastHashJoin(
     right: SparkPlan)
   extends BinaryNode with HashJoin with CodegenSupport {
 
+  val streamSideName = buildSide match {
+    case BuildLeft => "right"
+    case BuildRight => "left"
+  }
+
   override private[sql] lazy val metrics = Map(
-    "numLeftRows" -> SQLMetrics.createLongMetric(sparkContext, "number of left rows"),
-    "numRightRows" -> SQLMetrics.createLongMetric(sparkContext, "number of right rows"),
+    "numStreamRows" -> SQLMetrics.createLongMetric(sparkContext, s"number of $streamSideName rows"),
     "numOutputRows" -> SQLMetrics.createLongMetric(sparkContext, "number of output rows"))
 
   override def outputPartitioning: Partitioning = streamedPlan.outputPartitioning
 
   override def requiredChildDistribution: Seq[Distribution] = buildSide match {
-    case BuildLeft => longMetric("numLeftRows")
+    case BuildLeft =>
       BroadcastDistribution(buildRelation) :: UnspecifiedDistribution :: Nil
     case BuildRight =>
       UnspecifiedDistribution :: BroadcastDistribution(buildRelation) :: Nil
@@ -61,10 +65,7 @@ case class BroadcastHashJoin(
   }
 
   protected override def doExecute(): RDD[InternalRow] = {
-    val numStreamedRows = buildSide match {
-      case BuildLeft => longMetric("numRightRows")
-      case BuildRight => longMetric("numLeftRows")
-    }
+    val numStreamedRows = longMetric("numStreamRows")
     val numOutputRows = longMetric("numOutputRows")
 
     val broadcastRelation = Broadcast.broadcastRelation[HashedRelation](buildPlan)

sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashOuterJoin.scala

@@ -40,9 +40,18 @@ case class BroadcastHashOuterJoin(
     left: SparkPlan,
     right: SparkPlan) extends BinaryNode with HashOuterJoin {
 
+  private[this] def failOnWrongJoinType(jt: JoinType): Nothing = {
+    throw new IllegalArgumentException(s"HashOuterJoin should not take $jt as the JoinType")
+  }
+
+  val streamSideName = joinType match {
+    case RightOuter => "right"
+    case LeftOuter => "left"
+    case jt => failOnWrongJoinType(jt)
+  }
+
   override private[sql] lazy val metrics = Map(
-    "numLeftRows" -> SQLMetrics.createLongMetric(sparkContext, "number of left rows"),
-    "numRightRows" -> SQLMetrics.createLongMetric(sparkContext, "number of right rows"),
+    "numStreamRows" -> SQLMetrics.createLongMetric(sparkContext, s"number of $streamSideName rows"),
     "numOutputRows" -> SQLMetrics.createLongMetric(sparkContext, "number of output rows"))
 
   override def requiredChildDistribution: Seq[Distribution] = joinType match {
@@ -51,7 +60,7 @@ case class BroadcastHashOuterJoin(
     case LeftOuter =>
       UnspecifiedDistribution :: BroadcastDistribution(buildRelation) :: Nil
     case x =>
-      throw new IllegalArgumentException(s"HashOuterJoin should not take $x as the JoinType")
+      failOnWrongJoinType(x)
   }
 
   private val buildRelation: Iterable[InternalRow] => HashedRelation = { input =>
@@ -61,13 +70,7 @@ case class BroadcastHashOuterJoin(
   override def outputPartitioning: Partitioning = streamedPlan.outputPartitioning
 
   override def doExecute(): RDD[InternalRow] = {
-    val numStreamedRows = joinType match {
-      case RightOuter => longMetric("numRightRows")
-      case LeftOuter => longMetric("numLeftRows")
-      case x =>
-        throw new IllegalArgumentException(
-          s"HashOuterJoin should not take $x as the JoinType")
-    }
+    val numStreamedRows = longMetric("numStreamRows")
     val numOutputRows = longMetric("numOutputRows")
 
     val broadcastRelation = Broadcast.broadcastRelation[UnsafeHashedRelation](buildPlan)
@@ -102,8 +105,7 @@ case class BroadcastHashOuterJoin(
           })
 
         case x =>
-          throw new IllegalArgumentException(
-            s"BroadcastHashOuterJoin should not take $x as the JoinType")
+          failOnWrongJoinType(x)
       }
     }
   }

sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala

@@ -43,6 +43,7 @@ class BroadcastJoinSuite extends QueryTest with BeforeAndAfterAll {
     val conf = new SparkConf()
       .setMaster("local-cluster[2,1,1024]")
       .setAppName("testing")
+      .set("spark.sql.codegen.wholeStage", "false")
     val sc = new SparkContext(conf)
     sqlContext = new SQLContext(sc)
   }
@@ -62,7 +63,7 @@ class BroadcastJoinSuite extends QueryTest with BeforeAndAfterAll {
       // Comparison at the end is for broadcast left semi join
       val joinExpression = df1("key") === df2("key") && df1("value") > df2("value")
       val df3 = df1.join(broadcast(df2), joinExpression, joinType)
-      val plan = df3.queryExecution.sparkPlan
+      val plan = df3.queryExecution.executedPlan
       assert(plan.collect { case p: T => p }.size === 1)
       plan.executeCollect()
     }