[SPARK-19601] [SQL] Fix CollapseRepartition rule to preserve shuffle-enabled Repartition

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala

@@ -374,6 +374,9 @@ package object dsl {
         case plan => SubqueryAlias(alias, plan, None)
       }
 
+      def coalesce(num: Integer): LogicalPlan =
+        Repartition(num, shuffle = false, logicalPlan)
+
       def repartition(num: Integer): LogicalPlan =
         Repartition(num, shuffle = true, logicalPlan)
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -563,23 +563,27 @@ object CollapseProject extends Rule[LogicalPlan] {
 /**
  * Combines adjacent [[Repartition]] and [[RepartitionByExpression]] operator combinations
  * by keeping only the one.
- * 1. For adjacent [[Repartition]]s, collapse into the last [[Repartition]].
+ * 1. For adjacent [[Repartition]]s, collapse into the last [[Repartition]] if their shuffle types
+ *    are the same or the parent's shuffle is true.
  * 2. For adjacent [[RepartitionByExpression]]s, collapse into the last [[RepartitionByExpression]].
- * 3. For a combination of [[Repartition]] and [[RepartitionByExpression]], collapse as a single
- *    [[RepartitionByExpression]] with the expression and last number of partition.
+ * 3. When a shuffle-enabled [[Repartition]] is above a [[RepartitionByExpression]], collapse as a
+ *    single [[RepartitionByExpression]] with the expression and the last number of partition.
+ * 4. When a [[RepartitionByExpression]] is above a [[Repartition]], collapse as a single
+ *    [[RepartitionByExpression]] with the expression and the last number of partition.
  */
 object CollapseRepartition extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transformUp {
     // Case 1
-    case Repartition(numPartitions, shuffle, Repartition(_, _, child)) =>
+    case Repartition(numPartitions, shuffle, Repartition(_, shuffleChild, child))
+        if shuffle == shuffleChild || shuffle =>
       Repartition(numPartitions, shuffle, child)
     // Case 2
     case RepartitionByExpression(exprs, RepartitionByExpression(_, child, _), numPartitions) =>
       RepartitionByExpression(exprs, child, numPartitions)
     // Case 3
-    case Repartition(numPartitions, _, r: RepartitionByExpression) =>
+    case Repartition(numPartitions, shuffle, r: RepartitionByExpression) if shuffle =>
       r.copy(numPartitions = Some(numPartitions))
-    // Case 3
+    // Case 4
     case RepartitionByExpression(exprs, Repartition(_, _, child), numPartitions) =>
       RepartitionByExpression(exprs, child, numPartitions)
   }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseRepartitionSuite.scala

@@ -32,6 +32,18 @@ class CollapseRepartitionSuite extends PlanTest {
 
   val testRelation = LocalRelation('a.int, 'b.int)
 
+
+  test("collapse two adjacent coalesces into one") {
+    val query = testRelation
+      .coalesce(10)
+      .coalesce(20)
+
+    val optimized = Optimize.execute(query.analyze)
+    val correctAnswer = testRelation.coalesce(20).analyze
+
+    comparePlans(optimized, correctAnswer)
+  }
+
   test("collapse two adjacent repartitions into one") {
     val query = testRelation
       .repartition(10)
@@ -43,15 +55,44 @@ class CollapseRepartitionSuite extends PlanTest {
     comparePlans(optimized, correctAnswer)
   }
 
+  test("collapse one coalesce and one repartition into one") {
+    val query1 = testRelation
+      .coalesce(20)
+      .repartition(5)
+
+    val optimized1 = Optimize.execute(query1.analyze)
+    val correctAnswer1 = testRelation.repartition(5).analyze
+
+    comparePlans(optimized1, correctAnswer1)
+
+    val query2 = testRelation
+      .repartition(5)
+      .coalesce(20)
+
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer2 = testRelation.repartition(5).coalesce(20).analyze
+
+    comparePlans(optimized2, correctAnswer2)
+  }
+
   test("collapse repartition and repartitionBy into one") {
-    val query = testRelation
+    val query1 = testRelation
       .repartition(10)
       .distribute('a)(20)
 
-    val optimized = Optimize.execute(query.analyze)
-    val correctAnswer = testRelation.distribute('a)(20).analyze
+    val optimized1 = Optimize.execute(query1.analyze)
+    val correctAnswer1 = testRelation.distribute('a)(20).analyze
 
-    comparePlans(optimized, correctAnswer)
+    comparePlans(optimized1, correctAnswer1)
+
+    val query2 = testRelation
+      .coalesce(10)
+      .distribute('a)(20)
+
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer2 = testRelation.distribute('a)(20).analyze
+
+    comparePlans(optimized2, correctAnswer2)
   }
 
   test("collapse repartitionBy and repartition into one") {
@@ -65,6 +106,17 @@ class CollapseRepartitionSuite extends PlanTest {
     comparePlans(optimized, correctAnswer)
   }
 
+  test("do not collapse coalesce above repartitionBy into one") {
+    val query = testRelation
+      .distribute('a)(20)
+      .coalesce(10)
+
+    val optimized = Optimize.execute(query.analyze)
+    val correctAnswer = testRelation.distribute('a)(20).coalesce(10).analyze
+
+    comparePlans(optimized, correctAnswer)
+  }
+
   test("collapse two adjacent repartitionBys into one") {
     val query = testRelation
       .distribute('b)(10)

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

@@ -242,11 +242,12 @@ class PlannerSuite extends SharedSQLContext {
     val doubleRepartitioned = testData.repartition(10).repartition(20).coalesce(5)
     def countRepartitions(plan: LogicalPlan): Int = plan.collect { case r: Repartition => r }.length
     assert(countRepartitions(doubleRepartitioned.queryExecution.logical) === 3)
-    assert(countRepartitions(doubleRepartitioned.queryExecution.optimizedPlan) === 1)
+    assert(countRepartitions(doubleRepartitioned.queryExecution.optimizedPlan) === 2)
     doubleRepartitioned.queryExecution.optimizedPlan match {
-      case r: Repartition =>
-        assert(r.numPartitions === 5)
-        assert(r.shuffle === false)
+      case Repartition (numPartitions, shuffle, Repartition(_, shuffleChild, _)) =>
+        assert(numPartitions === 5)
+        assert(shuffle === false)
+        assert(shuffleChild === true)
     }
   }