Always use nullSafe version partitioning when creating Exchanges.

Author: yhuai

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

@@ -98,6 +98,8 @@ sealed trait Partitioning {
    * Returns true iff we can say that the partitioning scheme of this [[Partitioning]]
    * guarantees the same partitioning scheme described by `other`.
    *
+   * If a [[Partitioning]] supports `nullSafe` setting, the nullSafe version of this
+   * [[Partitioning]] should always `guarantees` its nullUnsafe version.
    * For example, HashPartitioning(expressions = 'a, numPartitions = 10, nullSafe = true)
    * guarantees HashPartitioning(expressions = 'a, numPartitions = 10, nullSafe = false).
    * However, HashPartitioning(expressions = 'a, numPartitions = 10, nullSafe = false) does not

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

@@ -213,7 +213,12 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
 
         def addShuffleIfNecessary(child: SparkPlan): SparkPlan = {
           if (!child.outputPartitioning.guarantees(partitioning)) {
-            Exchange(partitioning, child)
+            // If the child's outputPartitioning does not guarantees partitioning,
+            // we need to add an Exchange operator. At here, we always use
+            // the nullSafe version of the given partitioning because the nullSafe
+            // version always guarantees the nullUnsafe version of the partitioning and
+            // we do not have any special handling for nullUnsafe partitioning for now.
+            Exchange(partitioning.withNullSafeSetting(newNullSafe = true), child)
           } else {
             child
           }

sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala

@@ -210,6 +210,54 @@ class PlannerSuite extends SparkFunSuite with SQLTestUtils {
           }
           i += 1
         }
+
+        {
+          val numExchanges: Int = sql(
+            s"""
+                |SELECT small.key, count(*)
+                |FROM
+                |  normal JOIN small ON (normal.key = small.key)
+                |  JOIN tiny ON (small.key = tiny.key)
+                |GROUP BY
+                |  small.key
+              """.stripMargin
+          ).queryExecution.executedPlan.collect {
+            case exchange: Exchange => exchange
+          }.length
+          assert(numExchanges === 3)
+        }
+
+        {
+          val numExchanges: Int = sql(
+            s"""
+                |SELECT normal.key, count(*)
+                |FROM
+                |  normal LEFT OUTER JOIN small ON (normal.key = small.key)
+                |  JOIN tiny ON (small.key = tiny.key)
+                |GROUP BY
+                |  normal.key
+              """.stripMargin
+          ).queryExecution.executedPlan.collect {
+            case exchange: Exchange => exchange
+          }.length
+          assert(numExchanges === 3)
+        }
+
+        {
+          val numExchanges: Int = sql(
+            s"""
+                |SELECT small.key, count(*)
+                |FROM
+                |  normal LEFT OUTER JOIN small ON (normal.key = small.key)
+                |  JOIN tiny ON (small.key = tiny.key)
+                |GROUP BY
+                |  small.key
+              """.stripMargin
+          ).queryExecution.executedPlan.collect {
+            case exchange: Exchange => exchange
+          }.length
+          assert(numExchanges === 4)
+        }
       }
     }
   }