Add PartitioningCollection.

Author: yhuai

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

@@ -305,3 +305,43 @@ case class RangePartitioning(ordering: Seq[SortOrder], numPartitions: Int)
 
   override def toNullUnsafePartitioning: Partitioning = this
 }
+
+/**
+ * A collection of [[Partitioning]]s.
+ */
+case class PartitioningCollection(partitionings: Seq[Partitioning])
+  extends Expression with Partitioning with Unevaluable {
+
+  require(
+    partitionings.map(_.numPartitions).distinct.length == 1,
+    s"PartitioningCollection requires all of its partitionings have the same numPartitions.")
+
+  override def children: Seq[Expression] = partitionings.collect {
+    case expr: Expression => expr
+  }
+
+  override def nullable: Boolean = false
+
+  override def dataType: DataType = IntegerType
+
+  override val numPartitions = partitionings.map(_.numPartitions).distinct.head
+
+  override def satisfies(required: Distribution): Boolean =
+    partitionings.exists(_.satisfies(required))
+
+  override def compatibleWith(other: Partitioning): Boolean =
+    partitionings.exists(_.compatibleWith(other))
+
+  override def guarantees(other: Partitioning): Boolean =
+    partitionings.exists(_.guarantees(other))
+
+  override def keyExpressions: Seq[Expression] = partitionings.head.keyExpressions
+
+  override def toNullUnsafePartitioning: Partitioning = {
+    PartitioningCollection(partitionings.map(_.toNullUnsafePartitioning))
+  }
+
+  override def toString: String = {
+    partitionings.map(_.toString).mkString("(", " or ", ")")
+  }
+}

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

@@ -164,7 +164,20 @@ case class Exchange(newPartitioning: Partitioning, child: SparkPlan) extends Una
       // TODO: Handle BroadcastPartitioning.
     }
     def getPartitionKeyExtractor(): InternalRow => InternalRow = newPartitioning match {
-      case NullSafeHashPartitioning(expressions, _) => newMutableProjection(expressions, child.output)()
+      case NullSafeHashPartitioning(expressions, _) =>
+        // Since NullSafeHashPartitioning and NullUnsafeHashPartitioning may be used together
+        // for a join operator. We need to make sure they calculate the partition id with
+        // the same way.
+        val materalizeExpressions = newMutableProjection(expressions, child.output)()
+        val partitionExpressionSchema = expressions.map {
+          case ne: NamedExpression => ne.toAttribute
+          case expr => Alias(expr, "partitionExpr")().toAttribute
+        }
+        val partitionId = RowHashCode
+        val partitionIdExtractor =
+          newMutableProjection(partitionId :: Nil, partitionExpressionSchema)()
+        (row: InternalRow) => partitionIdExtractor(materalizeExpressions(row))
+        // newMutableProjection(expressions, child.output)()
       case NullUnsafeHashPartitioning(expressions, numPartition) =>
         // For NullUnsafeHashPartitioning, we do not want to send rows having any expression
         // in `expressions` evaluated as null to the same node.
@@ -261,7 +274,7 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
           child: SparkPlan): SparkPlan = {
 
         def addShuffleIfNecessary(child: SparkPlan): SparkPlan = {
-          if (child.outputPartitioning.guarantees(partitioning)) {
+          if (!child.outputPartitioning.guarantees(partitioning)) {
             Exchange(partitioning, child)
           } else {
             child

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

@@ -21,7 +21,7 @@ import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.Expression
-import org.apache.spark.sql.catalyst.plans.physical.{Distribution, NullUnsafeClusteredDistribution, NullSafeClusteredDistribution, Partitioning}
+import org.apache.spark.sql.catalyst.plans.physical._
 import org.apache.spark.sql.execution.{BinaryNode, SparkPlan}
 
 /**
@@ -38,7 +38,8 @@ case class ShuffledHashJoin(
     right: SparkPlan)
   extends BinaryNode with HashJoin {
 
-  override def outputPartitioning: Partitioning = left.outputPartitioning
+  override def outputPartitioning: Partitioning =
+    PartitioningCollection(Seq(left.outputPartitioning, right.outputPartitioning))
 
   override def requiredChildDistribution: Seq[Distribution] =
     NullSafeClusteredDistribution(leftKeys) :: NullSafeClusteredDistribution(rightKeys) :: Nil

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

@@ -48,9 +48,19 @@ case class ShuffledHashOuterJoin(
       NullUnsafeClusteredDistribution(rightKeys) :: Nil
 
   override def outputPartitioning: Partitioning = joinType match {
-    case LeftOuter => left.outputPartitioning
-    case RightOuter => right.outputPartitioning
-    case FullOuter => left.outputPartitioning.toNullUnsafePartitioning
+    case LeftOuter =>
+      val partitions =
+        Seq(left.outputPartitioning, right.outputPartitioning.toNullUnsafePartitioning)
+      PartitioningCollection(partitions)
+    case RightOuter =>
+      val partitions =
+        Seq(right.outputPartitioning, left.outputPartitioning.toNullUnsafePartitioning)
+      PartitioningCollection(partitions)
+    case FullOuter =>
+      val partitions =
+        Seq(left.outputPartitioning.toNullUnsafePartitioning,
+          right.outputPartitioning.toNullUnsafePartitioning)
+      PartitioningCollection(partitions)
     case x =>
       throw new IllegalArgumentException(s"HashOuterJoin should not take $x as the JoinType")
   }

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

@@ -40,7 +40,8 @@ case class SortMergeJoin(
 
   override def output: Seq[Attribute] = left.output ++ right.output
 
-  override def outputPartitioning: Partitioning = left.outputPartitioning
+  override def outputPartitioning: Partitioning =
+    PartitioningCollection(Seq(left.outputPartitioning, right.outputPartitioning))
 
   override def requiredChildDistribution: Seq[Distribution] =
     NullUnsafeClusteredDistribution(leftKeys) :: NullUnsafeClusteredDistribution(rightKeys) :: Nil