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

R/pkg/inst/tests/testthat/test_sparkSQL.R

@@ -2558,8 +2558,8 @@ test_that("coalesce, repartition, numPartitions", {
 
   df2 <- repartition(df1, 10)
   expect_equal(getNumPartitions(df2), 10)
-  expect_equal(getNumPartitions(coalesce(df2, 13)), 5)
-  expect_equal(getNumPartitions(coalesce(df2, 7)), 5)
+  expect_equal(getNumPartitions(coalesce(df2, 13)), 10)
+  expect_equal(getNumPartitions(coalesce(df2, 7)), 7)
   expect_equal(getNumPartitions(coalesce(df2, 3)), 3)
 })
 

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

@@ -370,6 +370,9 @@ package object dsl {
 
       def as(alias: String): LogicalPlan = SubqueryAlias(alias, logicalPlan, 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

@@ -562,27 +562,43 @@ 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]].
- * 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.
+ * Combines adjacent [[RepartitionOperation]] operators
  */
 object CollapseRepartition extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transformUp {
-    // Case 1
-    case Repartition(numPartitions, shuffle, Repartition(_, _, child)) =>
-      Repartition(numPartitions, shuffle, child)
-    // Case 2
-    case RepartitionByExpression(exprs, RepartitionByExpression(_, child, _), numPartitions) =>
-      RepartitionByExpression(exprs, child, numPartitions)
-    // Case 3
-    case Repartition(numPartitions, _, r: RepartitionByExpression) =>
-      r.copy(numPartitions = numPartitions)
-    // Case 3
-    case RepartitionByExpression(exprs, Repartition(_, _, child), numPartitions) =>
-      RepartitionByExpression(exprs, child, numPartitions)
+    // Case 1: When a Repartition has a child of Repartition or RepartitionByExpression,
+    // we can collapse it with the child based on the type of shuffle and the specified number
+    // of partitions.
+    case r @ Repartition(_, _, child: Repartition) =>
+      collapseRepartition(r, child)
+    case r @ Repartition(_, _, child: RepartitionByExpression) =>
+      collapseRepartition(r, child)
+    // Case 2: When a RepartitionByExpression has a child of Repartition or RepartitionByExpression
+    // we can remove the child.
+    case r @ RepartitionByExpression(_, child: RepartitionByExpression, _) =>
+      r.copy(child = child.child)
+    case r @ RepartitionByExpression(_, child: Repartition, _) =>
+      r.copy(child = child.child)
+  }
+
+  /**
+   * Collapses the [[Repartition]] with its child [[RepartitionOperation]], if possible.
+   * - Case 1 the top [[Repartition]] does not enable shuffle (i.e., coalesce API):
+   *   If the last numPartitions is bigger, returns the child node; otherwise, keep unchanged.
+   * - Case 2 the top [[Repartition]] enables shuffle (i.e., repartition API):
+   *   returns the child node with the last numPartitions.
+   */
+  private def collapseRepartition(r: Repartition, child: RepartitionOperation): LogicalPlan = {
+    (r.shuffle, child.shuffle) match {
+      case (false, true) => child match {
+        case c: Repartition => if (r.numPartitions >= c.numPartitions) c else r
+        case c: RepartitionByExpression => if (r.numPartitions >= c.numPartitions) c else r
+      }
+      case _ => child match {
+        case child: Repartition => child.copy(numPartitions = r.numPartitions, shuffle = r.shuffle)
+        case child: RepartitionByExpression => child.copy(numPartitions = r.numPartitions)
+      }
+    }
   }
 }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala

@@ -835,16 +835,24 @@ case class Distinct(child: LogicalPlan) extends UnaryNode {
   override def output: Seq[Attribute] = child.output
 }
 
+/**
+ * A base interface for [[RepartitionByExpression]] and [[Repartition]]
+ */
+abstract class RepartitionOperation extends UnaryNode {
+  def shuffle: Boolean
+  def numPartitions: Int
+  override def output: Seq[Attribute] = child.output
+}
+
 /**
  * Returns a new RDD that has exactly `numPartitions` partitions. Differs from
  * [[RepartitionByExpression]] as this method is called directly by DataFrame's, because the user
  * asked for `coalesce` or `repartition`. [[RepartitionByExpression]] is used when the consumer
  * of the output requires some specific ordering or distribution of the data.
  */
 case class Repartition(numPartitions: Int, shuffle: Boolean, child: LogicalPlan)
-  extends UnaryNode {
+  extends RepartitionOperation {
   require(numPartitions > 0, s"Number of partitions ($numPartitions) must be positive.")
-  override def output: Seq[Attribute] = child.output
 }
 
 /**
@@ -856,12 +864,12 @@ case class Repartition(numPartitions: Int, shuffle: Boolean, child: LogicalPlan)
 case class RepartitionByExpression(
     partitionExpressions: Seq[Expression],
     child: LogicalPlan,
-    numPartitions: Int) extends UnaryNode {
+    numPartitions: Int) extends RepartitionOperation {
 
   require(numPartitions > 0, s"Number of partitions ($numPartitions) must be positive.")
 
   override def maxRows: Option[Long] = child.maxRows
-  override def output: Seq[Attribute] = child.output
+  override def shuffle: Boolean = true
 }
 
 /**

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

@@ -32,47 +32,180 @@ class CollapseRepartitionSuite extends PlanTest {
 
   val testRelation = LocalRelation('a.int, 'b.int)
 
+
+  test("collapse two adjacent coalesces into one") {
+    // Always respects the top coalesces amd removes useless coalesce below coalesce
+    val query1 = testRelation
+      .coalesce(10)
+      .coalesce(20)
+    val query2 = testRelation
+      .coalesce(30)
+      .coalesce(20)
+
+    val optimized1 = Optimize.execute(query1.analyze)
+    val optimized2 = Optimize.execute(query2.analyze)
+
+    val correctAnswer = testRelation.coalesce(20).analyze
+
+    comparePlans(optimized1, correctAnswer)
+    comparePlans(optimized2, correctAnswer)
+  }
+
   test("collapse two adjacent repartitions into one") {
-    val query = testRelation
+    // Always respects the top repartition amd removes useless repartition below repartition
+    val query1 = testRelation
+      .repartition(10)
+      .repartition(20)
+    val query2 = testRelation
+      .repartition(30)
+      .repartition(20)
+
+    val optimized1 = Optimize.execute(query1.analyze)
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer = testRelation.repartition(20).analyze
+
+    comparePlans(optimized1, correctAnswer)
+    comparePlans(optimized2, correctAnswer)
+  }
+
+  test("coalesce above repartition") {
+    // Remove useless coalesce above repartition
+    val query1 = testRelation
       .repartition(10)
+      .coalesce(20)
+
+    val optimized1 = Optimize.execute(query1.analyze)
+    val correctAnswer1 = testRelation.repartition(10).analyze
+
+    comparePlans(optimized1, correctAnswer1)
+
+    // No change in this case
+    val query2 = testRelation
+      .repartition(30)
+      .coalesce(20)
+
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer2 = query2.analyze
+
+    comparePlans(optimized2, correctAnswer2)
+  }
+
+  test("repartition above coalesce") {
+    // Always respects the top repartition amd removes useless coalesce below repartition
+    val query1 = testRelation
+      .coalesce(10)
+      .repartition(20)
+    // Remove useless coalesce above repartition
+    val query2 = testRelation
+      .coalesce(30)
       .repartition(20)
 
-    val optimized = Optimize.execute(query.analyze)
+    val optimized1 = Optimize.execute(query1.analyze)
+    val optimized2 = Optimize.execute(query2.analyze)
+
     val correctAnswer = testRelation.repartition(20).analyze
 
-    comparePlans(optimized, correctAnswer)
+    comparePlans(optimized1, correctAnswer)
+    comparePlans(optimized2, correctAnswer)
   }
 
-  test("collapse repartition and repartitionBy into one") {
-    val query = testRelation
+  test("repartitionBy above repartition") {
+    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(optimized1, correctAnswer1)
+
+    val query2 = testRelation
+      .repartition(30)
+      .distribute('a)(20)
+
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer2 = testRelation.distribute('a)(20).analyze
 
-    comparePlans(optimized, correctAnswer)
+    comparePlans(optimized2, correctAnswer2)
   }
 
-  test("collapse repartitionBy and repartition into one") {
-    val query = testRelation
+  test("repartitionBy above coalesce") {
+    val query1 = testRelation
+      .coalesce(10)
+      .distribute('a)(20)
+
+    val optimized1 = Optimize.execute(query1.analyze)
+    val correctAnswer1 = testRelation.distribute('a)(20).analyze
+
+    comparePlans(optimized1, correctAnswer1)
+
+    val query2 = testRelation
+      .coalesce(20)
+      .distribute('a)(30)
+
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer2 = testRelation.distribute('a)(30).analyze
+
+    comparePlans(optimized2, correctAnswer2)
+  }
+
+  test("repartition above repartitionBy") {
+    val query1 = testRelation
       .distribute('a)(20)
       .repartition(10)
 
-    val optimized = Optimize.execute(query.analyze)
-    val correctAnswer = testRelation.distribute('a)(10).analyze
+    val optimized1 = Optimize.execute(query1.analyze)
+    val correctAnswer1 = testRelation.distribute('a)(10).analyze
+
+    comparePlans(optimized1, correctAnswer1)
+
+    val query2 = testRelation
+      .distribute('a)(20)
+      .repartition(30)
+
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer2 = testRelation.distribute('a)(30).analyze
 
-    comparePlans(optimized, correctAnswer)
+    comparePlans(optimized2, correctAnswer2)
+  }
+
+  test("coalesce above repartitionBy") {
+    val query1 = testRelation
+      .distribute('a)(20)
+      .coalesce(10)
+
+    val optimized1 = Optimize.execute(query1.analyze)
+    val correctAnswer1 = testRelation.distribute('a)(20).coalesce(10).analyze
+
+    comparePlans(optimized1, correctAnswer1)
+
+    val query2 = testRelation
+      .distribute('a)(20)
+      .coalesce(30)
+
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer2 = testRelation.distribute('a)(20).analyze
+
+    comparePlans(optimized2, correctAnswer2)
   }
 
   test("collapse two adjacent repartitionBys into one") {
-    val query = testRelation
+    val query1 = testRelation
       .distribute('b)(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(optimized1, correctAnswer1)
+
+    val query2 = testRelation
+      .distribute('b)(30)
+      .distribute('a)(20)
+
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer2 = testRelation.distribute('a)(20).analyze
 
-    comparePlans(optimized, correctAnswer)
+    comparePlans(optimized2, correctAnswer2)
   }
 }

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)
     }
   }