Move transform into BroadcastMode

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala

@@ -80,12 +80,14 @@ case class OrderedDistribution(ordering: Seq[SortOrder]) extends Distribution {
  * Marker trait to identify the shape in which tuples are broadcasted. Typical examples of this are
  * identity (tuples remain unchanged) or hashed (tuples are converted into some hash index).
  */
-trait BroadcastMode
+trait BroadcastMode extends (Array[InternalRow] => Any)
 
 /**
  * IdentityBroadcastMode requires that rows are broadcasted in their original form.
  */
-case object IdentityBroadcastMode extends BroadcastMode
+case object IdentityBroadcastMode extends BroadcastMode {
+  def apply(rows: Array[InternalRow]): Array[InternalRow] = rows
+}
 
 /**
   * Represents data where tuples are broadcasted to every node. It is quite common that the

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

@@ -31,7 +31,6 @@ import org.apache.spark.util.ThreadUtils
  */
 case class Broadcast(
     mode: BroadcastMode,
-    transform: Array[InternalRow] => Any,
     child: SparkPlan) extends UnaryNode {
 
   override def output: Seq[Attribute] = child.output
@@ -60,7 +59,7 @@ case class Broadcast(
         }.collect()
 
         // Construct and broadcast the relation.
-        sparkContext.broadcast(transform(input))
+        sparkContext.broadcast(mode(input))
       }
     }(Broadcast.executionContext)
   }

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

@@ -388,19 +388,6 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
     withCoordinator
   }
 
-  /**
-   * Create a [[Broadcast]] operator for a given [[BroadcastMode]] and [[SparkPlan]].
-   */
-  private def createBroadcast(mode: BroadcastMode, plan: SparkPlan): Broadcast = mode match {
-    case IdentityBroadcastMode =>
-      Broadcast(mode, identity, plan)
-    case HashSetBroadcastMode(keys) =>
-      Broadcast(mode, HashSemiJoin.buildKeyHashSet(keys, plan, _), plan)
-    case HashedRelationBroadcastMode(canJoinKeyFitWithinLong, keys) =>
-      Broadcast(mode, HashedRelation(canJoinKeyFitWithinLong, keys, plan, _), plan)
-    case _ => sys.error(s"Unknown BroadcastMode: $mode")
-  }
-
   private def ensureDistributionAndOrdering(operator: SparkPlan): SparkPlan = {
     val requiredChildDistributions: Seq[Distribution] = operator.requiredChildDistribution
     val requiredChildOrderings: Seq[Seq[SortOrder]] = operator.requiredChildOrdering
@@ -415,11 +402,11 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
       case (child, BroadcastDistribution(m1)) =>
         child match {
           // The child is broadcasting the same variable: keep the child.
-          case Broadcast(m2, _, _) if m1 == m2 => child
+          case Broadcast(m2, _) if m1 == m2 => child
           // The child is broadcasting a different variable: replace the child.
-          case Broadcast(m2, _, src) => createBroadcast(m1, src)
+          case Broadcast(m2, src) => Broadcast(m1, src)
           // Create a broadcast on top of the child.
-          case _ => createBroadcast(m1, child)
+          case _ => Broadcast(m1, child)
         }
       case (child, distribution) =>
         Exchange(createPartitioning(distribution, defaultNumPreShufflePartitions), child)

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

@@ -53,7 +53,10 @@ case class BroadcastHashJoin(
   override def outputPartitioning: Partitioning = streamedPlan.outputPartitioning
 
   override def requiredChildDistribution: Seq[Distribution] = {
-    val mode = HashedRelationBroadcastMode(canJoinKeyFitWithinLong, rewriteKeyExpr(buildKeys))
+    val mode = HashedRelationBroadcastMode(
+      canJoinKeyFitWithinLong,
+      rewriteKeyExpr(buildKeys),
+      buildPlan.output)
     buildSide match {
       case BuildLeft =>
         BroadcastDistribution(mode) :: UnspecifiedDistribution :: Nil

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

@@ -48,7 +48,10 @@ case class BroadcastHashOuterJoin(
     "numOutputRows" -> SQLMetrics.createLongMetric(sparkContext, "number of output rows"))
 
   override def requiredChildDistribution: Seq[Distribution] = {
-    val mode = HashedRelationBroadcastMode(canJoinKeyFitWithinLong = false, buildKeys)
+    val mode = HashedRelationBroadcastMode(
+      canJoinKeyFitWithinLong = false,
+      buildKeys,
+      buildPlan.output)
     joinType match {
       case RightOuter =>
         BroadcastDistribution(mode) :: UnspecifiedDistribution :: Nil

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

@@ -41,9 +41,9 @@ case class BroadcastLeftSemiJoinHash(
 
   override def requiredChildDistribution: Seq[Distribution] = {
     val mode = if (condition.isEmpty) {
-      HashSetBroadcastMode(rightKeys)
+      HashSetBroadcastMode(rightKeys, right.output)
     } else {
-      HashedRelationBroadcastMode(canJoinKeyFitWithinLong = false, rightKeys)
+      HashedRelationBroadcastMode(canJoinKeyFitWithinLong = false, rightKeys, right.output)
     }
     UnspecifiedDistribution :: BroadcastDistribution(mode) :: Nil
   }

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

@@ -80,13 +80,6 @@ trait HashSemiJoin {
 }
 
 private[execution] object HashSemiJoin {
-  def buildKeyHashSet(
-    keys: Seq[Expression],
-    plan: SparkPlan,
-    rows: Array[InternalRow]): java.util.HashSet[InternalRow] = {
-    buildKeyHashSet(keys, plan.output, rows.iterator)
-  }
-
   def buildKeyHashSet(
     keys: Seq[Expression],
     attributes: Seq[Attribute],
@@ -110,4 +103,10 @@ private[execution] object HashSemiJoin {
 }
 
 /** HashSetBroadcastMode requires that the input rows are broadcasted as a set. */
-private[execution] case class HashSetBroadcastMode(keys: Seq[Expression]) extends BroadcastMode
+private[execution] case class HashSetBroadcastMode(
+    keys: Seq[Expression],
+    attributes: Seq[Attribute]) extends BroadcastMode {
+  def apply(rows: Array[InternalRow]): java.util.HashSet[InternalRow] = {
+    HashSemiJoin.buildKeyHashSet(keys, attributes, rows.iterator)
+  }
+}

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

@@ -209,19 +209,6 @@ private[execution] object HashedRelation {
       new GeneralHashedRelation(hashTable)
     }
   }
-
-  def apply(
-    canJoinKeyFitWithinLong: Boolean,
-    keys: Seq[Expression],
-    child: SparkPlan,
-    rows: Array[InternalRow]): HashedRelation = {
-    val generator = UnsafeProjection.create(keys, child.output)
-    if (canJoinKeyFitWithinLong) {
-      LongHashedRelation(rows.iterator, generator, rows.length)
-    } else {
-      HashedRelation(rows.iterator, generator, rows.length)
-    }
-  }
 }
 
 /**
@@ -691,5 +678,15 @@ private[joins] object LongHashedRelation {
 /** The HashedRelationBroadcastMode requires that rows are broadcasted as a HashedRelation. */
 private[execution] case class HashedRelationBroadcastMode(
     canJoinKeyFitWithinLong: Boolean,
-    keys: Seq[Expression]) extends BroadcastMode
+    keys: Seq[Expression],
+    attributes: Seq[Attribute]) extends BroadcastMode {
+  def apply(rows: Array[InternalRow]): HashedRelation = {
+    val generator = UnsafeProjection.create(keys, attributes)
+    if (canJoinKeyFitWithinLong) {
+      LongHashedRelation(rows.iterator, generator, rows.length)
+    } else {
+      HashedRelation(rows.iterator, generator, rows.length)
+    }
+  }
+}