[SPARK-9703] [SQL] Refactor EnsureRequirements to avoid certain unnecessary shuffles

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

@@ -95,6 +95,22 @@ sealed trait Partitioning {
   def guarantees(other: Partitioning): Boolean
 }
 
+object Partitioning {
+  def allCompatible(partitionings: Seq[Partitioning]): Boolean = {
+    // Note: this assumes transitivity
+    partitionings.sliding(2).map {
+      case Seq(a) => true
+      case Seq(a, b) =>
+        if (a.numPartitions != b.numPartitions) {
+          assert(!a.guarantees(b) && !b.guarantees(a))
+          false
+        } else {
+          a.guarantees(b) && b.guarantees(a)
+        }
+    }.forall(_ == true)
+  }
+}
+
 case class UnknownPartitioning(numPartitions: Int) extends Partitioning {
   override def satisfies(required: Distribution): Boolean = required match {
     case UnspecifiedDistribution => true

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

@@ -197,66 +197,108 @@ case class Exchange(newPartitioning: Partitioning, child: SparkPlan) extends Una
  * of input data meets the
  * [[org.apache.spark.sql.catalyst.plans.physical.Distribution Distribution]] requirements for
  * each operator by inserting [[Exchange]] Operators where required.  Also ensure that the
- * required input partition ordering requirements are met.
+ * input partition ordering requirements are met.
  */
 private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[SparkPlan] {
   // TODO: Determine the number of partitions.
-  def numPartitions: Int = sqlContext.conf.numShufflePartitions
+  private def numPartitions: Int = sqlContext.conf.numShufflePartitions
 
-  def apply(plan: SparkPlan): SparkPlan = plan.transformUp {
-    case operator: SparkPlan =>
-      // Adds Exchange or Sort operators as required
-      def addOperatorsIfNecessary(
-          partitioning: Partitioning,
-          rowOrdering: Seq[SortOrder],
-          child: SparkPlan): SparkPlan = {
-
-        def addShuffleIfNecessary(child: SparkPlan): SparkPlan = {
-          if (!child.outputPartitioning.guarantees(partitioning)) {
-            Exchange(partitioning, child)
-          } else {
-            child
-          }
-        }
+  /**
+   * Given a required distribution, returns a partitioning that satisfies that distribution.
+   */
+  private def canonicalPartitioning(requiredDistribution: Distribution): Partitioning = {
+    requiredDistribution match {
+      case AllTuples => SinglePartition
+      case ClusteredDistribution(clustering) => HashPartitioning(clustering, numPartitions)
+      case OrderedDistribution(ordering) => RangePartitioning(ordering, numPartitions)
+      case dist => sys.error(s"Do not know how to satisfy distribution $dist")
+    }
+  }
 
-        def addSortIfNecessary(child: SparkPlan): SparkPlan = {
+  /**
+   * Return true if all of the operator's children satisfy their output distribution requirements.
+   */
+  private def childPartitioningsSatisfyDistributionRequirements(operator: SparkPlan): Boolean = {
+    operator.children.zip(operator.requiredChildDistribution).forall {
+      case (child, distribution) => child.outputPartitioning.satisfies(distribution)
+    }
+  }
 
-          if (rowOrdering.nonEmpty) {
-            // If child.outputOrdering is [a, b] and rowOrdering is [a], we do not need to sort.
-            val minSize = Seq(rowOrdering.size, child.outputOrdering.size).min
-            if (minSize == 0 || rowOrdering.take(minSize) != child.outputOrdering.take(minSize)) {
-              sqlContext.planner.BasicOperators.getSortOperator(rowOrdering, global = false, child)
-            } else {
+  /**
+   * Given an operator, check whether the operator requires its children to have compatible
+   * output partitionings and add Exchanges to fix any detected incompatibilities.
+   */
+  private def ensureChildPartitioningsAreCompatible(operator: SparkPlan): SparkPlan = {
+    if (operator.requiresChildPartitioningsToBeCompatible) {
+      if (!Partitioning.allCompatible(operator.children.map(_.outputPartitioning))) {
+        val newChildren = operator.children.zip(operator.requiredChildDistribution).map {
+          case (child, requiredDistribution) =>
+            val targetPartitioning = canonicalPartitioning(requiredDistribution)
+            if (child.outputPartitioning.guarantees(targetPartitioning)) {
               child
+            } else {
+              Exchange(targetPartitioning, child)
             }
-          } else {
-            child
-          }
         }
-
-        addSortIfNecessary(addShuffleIfNecessary(child))
+        val newOperator = operator.withNewChildren(newChildren)
+        assert(childPartitioningsSatisfyDistributionRequirements(newOperator))
+        newOperator
+      } else {
+        operator
       }
+    } else {
+      operator
+    }
+  }
 
-      val requirements =
-        (operator.requiredChildDistribution, operator.requiredChildOrdering, operator.children)
+  private def ensureDistributionAndOrdering(operator: SparkPlan): SparkPlan = {
 
-      val fixedChildren = requirements.zipped.map {
-        case (AllTuples, rowOrdering, child) =>
-          addOperatorsIfNecessary(SinglePartition, rowOrdering, child)
-        case (ClusteredDistribution(clustering), rowOrdering, child) =>
-          addOperatorsIfNecessary(HashPartitioning(clustering, numPartitions), rowOrdering, child)
-        case (OrderedDistribution(ordering), rowOrdering, child) =>
-          addOperatorsIfNecessary(RangePartitioning(ordering, numPartitions), rowOrdering, child)
+    def addShuffleIfNecessary(child: SparkPlan, requiredDistribution: Distribution): SparkPlan = {
+      // A pre-condition of ensureDistributionAndOrdering is that joins' children have compatible
+      // partitionings. Thus, we only need to check whether the output partitionings satisfy
+      // the required distribution. In the case where the children are all compatible, then they
+      // will either all satisfy the required distribution or will all fail to satisfy it, since
+      // A.guarantees(B) implies that A and B satisfy the same set of distributions.
+      // Therefore, if all children are compatible then either all or none of them will shuffled to
+      // ensure that the distribution requirements are met.
+      //
+      // Note that this reasoning implicitly assumes that operators which require compatible
+      // child partitionings have equivalent required distributions for those children.
+      if (child.outputPartitioning.satisfies(requiredDistribution)) {
+        child
+      } else {
+        Exchange(canonicalPartitioning(requiredDistribution), child)
+      }
+    }
 
-        case (UnspecifiedDistribution, Seq(), child) =>
+    def addSortIfNecessary(child: SparkPlan, requiredOrdering: Seq[SortOrder]): SparkPlan = {
+      if (requiredOrdering.nonEmpty) {
+        // If child.outputOrdering is [a, b] and requiredOrdering is [a], we do not need to sort.
+        val minSize = Seq(requiredOrdering.size, child.outputOrdering.size).min
+        if (minSize == 0 || requiredOrdering.take(minSize) != child.outputOrdering.take(minSize)) {
+          sqlContext.planner.BasicOperators.getSortOperator(requiredOrdering, global = false, child)
+        } else {
           child
-        case (UnspecifiedDistribution, rowOrdering, child) =>
-          sqlContext.planner.BasicOperators.getSortOperator(rowOrdering, global = false, child)
-
-        case (dist, ordering, _) =>
-          sys.error(s"Don't know how to ensure $dist with ordering $ordering")
+        }
+      } else {
+        child
       }
+    }
+
+    val children = operator.children
+    val requiredChildDistribution = operator.requiredChildDistribution
+    val requiredChildOrdering = operator.requiredChildOrdering
+    assert(children.length == requiredChildDistribution.length)
+    assert(children.length == requiredChildOrdering.length)
+    val newChildren = (children, requiredChildDistribution, requiredChildOrdering).zipped.map {
+      case (child, requiredDistribution, requiredOrdering) =>
+        addSortIfNecessary(addShuffleIfNecessary(child, requiredDistribution), requiredOrdering)
+    }
+    operator.withNewChildren(newChildren)
+  }
 
-      operator.withNewChildren(fixedChildren)
+  def apply(plan: SparkPlan): SparkPlan = plan.transformUp {
+    case operator: SparkPlan =>
+      ensureDistributionAndOrdering(ensureChildPartitioningsAreCompatible(operator))
   }
 }

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

@@ -109,6 +109,12 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
   /** Specifies sort order for each partition requirements on the input data for this operator. */
   def requiredChildOrdering: Seq[Seq[SortOrder]] = Seq.fill(children.size)(Nil)
 
+  /**
+   * Specifies whether this operator requires all of its children to have [[outputPartitioning]]s
+   * that are compatible with each other.
+   */
+  def requiresChildPartitioningsToBeCompatible: Boolean = false
+
   /** Specifies whether this operator outputs UnsafeRows */
   def outputsUnsafeRows: Boolean = false
 

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

@@ -245,6 +245,11 @@ case class Repartition(numPartitions: Int, shuffle: Boolean, child: SparkPlan)
   extends UnaryNode {
   override def output: Seq[Attribute] = child.output
 
+  override def outputPartitioning: Partitioning = {
+    if (numPartitions == 1) SinglePartition
+    else UnknownPartitioning(numPartitions)
+  }
+
   protected override def doExecute(): RDD[InternalRow] = {
     child.execute().map(_.copy()).coalesce(numPartitions, shuffle)
   }

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

@@ -42,6 +42,8 @@ case class LeftSemiJoinHash(
   override def requiredChildDistribution: Seq[Distribution] =
     ClusteredDistribution(leftKeys) :: ClusteredDistribution(rightKeys) :: Nil
 
+  override def requiresChildPartitioningsToBeCompatible: Boolean = true
+
   protected override def doExecute(): RDD[InternalRow] = {
     right.execute().zipPartitions(left.execute()) { (buildIter, streamIter) =>
       if (condition.isEmpty) {

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

@@ -46,6 +46,8 @@ case class ShuffledHashJoin(
   override def requiredChildDistribution: Seq[Distribution] =
     ClusteredDistribution(leftKeys) :: ClusteredDistribution(rightKeys) :: Nil
 
+  override def requiresChildPartitioningsToBeCompatible: Boolean = true
+
   protected override def doExecute(): RDD[InternalRow] = {
     buildPlan.execute().zipPartitions(streamedPlan.execute()) { (buildIter, streamIter) =>
       val hashed = HashedRelation(buildIter, buildSideKeyGenerator)

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

@@ -44,6 +44,8 @@ case class ShuffledHashOuterJoin(
   override def requiredChildDistribution: Seq[Distribution] =
     ClusteredDistribution(leftKeys) :: ClusteredDistribution(rightKeys) :: Nil
 
+  override def requiresChildPartitioningsToBeCompatible: Boolean = true
+
   override def outputPartitioning: Partitioning = joinType match {
     case LeftOuter => left.outputPartitioning
     case RightOuter => right.outputPartitioning

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

@@ -48,6 +48,8 @@ case class SortMergeJoin(
   override def requiredChildDistribution: Seq[Distribution] =
     ClusteredDistribution(leftKeys) :: ClusteredDistribution(rightKeys) :: Nil
 
+  override def requiresChildPartitioningsToBeCompatible: Boolean = true
+
   override def outputOrdering: Seq[SortOrder] = requiredOrders(leftKeys)
 
   override def requiredChildOrdering: Seq[Seq[SortOrder]] =

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

@@ -21,6 +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._
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
 import org.apache.spark.sql.catalyst.rules.Rule
 
 /**
@@ -33,6 +34,8 @@ case class ConvertToUnsafe(child: SparkPlan) extends UnaryNode {
   require(UnsafeProjection.canSupport(child.schema), s"Cannot convert ${child.schema} to Unsafe")
 
   override def output: Seq[Attribute] = child.output
+  override def outputPartitioning: Partitioning = child.outputPartitioning
+  override def outputOrdering: Seq[SortOrder] = child.outputOrdering
   override def outputsUnsafeRows: Boolean = true
   override def canProcessUnsafeRows: Boolean = false
   override def canProcessSafeRows: Boolean = true
@@ -51,6 +54,8 @@ case class ConvertToUnsafe(child: SparkPlan) extends UnaryNode {
 @DeveloperApi
 case class ConvertToSafe(child: SparkPlan) extends UnaryNode {
   override def output: Seq[Attribute] = child.output
+  override def outputPartitioning: Partitioning = child.outputPartitioning
+  override def outputOrdering: Seq[SortOrder] = child.outputOrdering
   override def outputsUnsafeRows: Boolean = false
   override def canProcessUnsafeRows: Boolean = true
   override def canProcessSafeRows: Boolean = false