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

gatorsmile · 2017-02-14T21:20:52Z

What changes were proposed in this pull request?

Observed by @felixcheung in #16739, when users use the shuffle-enabled repartition API, they expect the partition they got should be the exact number they provided, even if they call shuffle-disabled coalesce later.

Currently, CollapseRepartition rule does not consider whether shuffle is enabled or not. Thus, we got the following unexpected result.

    val df = spark.range(0, 10000, 1, 5)
    val df2 = df.repartition(10)
    assert(df2.coalesce(13).rdd.getNumPartitions == 5)
    assert(df2.coalesce(7).rdd.getNumPartitions == 5)
    assert(df2.coalesce(3).rdd.getNumPartitions == 3)

This PR is to fix the issue. We preserve shuffle-enabled Repartition.

How was this patch tested?

Added a test case

dongjoon-hyun · 2017-02-14T22:45:12Z

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

    // Case 1
-    case Repartition(numPartitions, shuffle, Repartition(_, _, child)) =>
+    case Repartition(numPartitions, shuffle, Repartition(_, shuffleChild, child))
+        if shuffle == shuffleChild || shuffle =>


Oh, thank you for fixing this!

SparkQA · 2017-02-15T00:14:43Z

Test build #72894 has finished for PR 16933 at commit 7b4a9dd.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

felixcheung · 2017-02-15T07:55:41Z

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

        case plan => SubqueryAlias(alias, plan, None)
      }

+      def coalesce(num: Integer): LogicalPlan =


does it conflict with sql coalesce by having it here?

They are used for the test cases in catalyst package, in which Dataset APIs are not available. Thus, that is why we add these DSL for test cases

felixcheung · 2017-02-15T07:56:52Z

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

- *    [[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


does shuffle type matter for Repartition in this case?

RepartitionByExpression always uses ShuffleExchange. Thus, it is like Repartition with enabled shuffle.

right, I was referring to shuffle on Repartition, but I see your point of RepartitionByExpression overriding it regardless

felixcheung · 2017-02-15T08:01:14Z

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

+  test("collapse two adjacent coalesces into one") {
+    val query = testRelation
+      .coalesce(10)
+      .coalesce(20)


hmm, I can see the argument.
but there are 2 adjacent coalesces like this shouldn't it take the smaller number? (since coalesce can't increase partition numbers)
whereas if there are 2 adjacent repartition it could take the last number

I think it would be better to respect the later input number, which is specified by users, for avoiding any surprise to users.

ok, agreed.

felixcheung · 2017-02-15T08:03:40Z

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

+      .coalesce(20)
+
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer2 = testRelation.repartition(5).coalesce(20).analyze


that might be the plan but the end result should be numPartitions == 5 correct? is there another suite we could add tests for repartition/coalesce like this?

Yeah. We can get rid of coalesce if the number of partitions is smaller than the child repartition

Actually, I can add some simple end-to-end test cases like what you did in the R side.

For improving this rule, we need to clean up the resolution of RepartitionByExpression at first. See the PR #16988

SparkQA · 2017-03-04T23:53:20Z

Test build #73910 has finished for PR 16933 at commit 0f95a6f.

This patch passes all tests.
This patch merges cleanly.
This patch adds the following public classes (experimental):
abstract class RepartitionOperation(numPartitions: Int) extends UnaryNode

SparkQA · 2017-03-05T00:14:09Z

Test build #73911 has finished for PR 16933 at commit 680c3af.

This patch passes all tests.
This patch merges cleanly.
This patch adds the following public classes (experimental):
abstract class RepartitionOperation extends UnaryNode

gatorsmile · 2017-03-06T06:43:15Z

cc @cloud-fan @hvanhovell

cloud-fan · 2017-03-06T07:26:50Z

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

+    // 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) =>


we can just write one case case r @ Repartition(_, _, child: RepartitionOperation)

cloud-fan · 2017-03-06T07:28:10Z

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

+   */
+  private def collapseRepartition(r: Repartition, child: RepartitionOperation): LogicalPlan = {
+    (r.shuffle, child.shuffle) match {
+      case (false, true) => child match {


why this pattern match? we can just call child.numPartitions

cloud-fan · 2017-03-06T23:54:19Z

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

+   * - 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 = {


can we inline this method?

SparkQA · 2017-03-07T00:51:29Z

Test build #74045 has finished for PR 16933 at commit 5453ad4.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2017-03-07T05:21:22Z

Test build #74064 has finished for PR 16933 at commit 4649af4.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

cloud-fan · 2017-03-07T20:37:18Z

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

+      .distribute('a)(20)
+
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer2 = testRelation.distribute('a)(20).analyze


this is same as correctAnswer1

cloud-fan · 2017-03-07T20:40:17Z

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

-    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


I not quite sure about this. Shall we optimize to relation.repartition(10)?

Here, I just followed what we did before. After more code reading, I think we can do it, since RoundRobinPartitioning looks cheaper.

case logical.Repartition(numPartitions, shuffle, child) => if (shuffle) { ShuffleExchange(RoundRobinPartitioning(numPartitions), planLater(child)) :: Nil } else { execution.CoalesceExec(numPartitions, planLater(child)) :: Nil } case logical.RepartitionByExpression(expressions, child, numPartitions) => exchange.ShuffleExchange(HashPartitioning( expressions, numPartitions), planLater(child)) :: Nil

My concern is, optimization should not change the result. relation.distributeBy('a, 10).repartition(10) should have same result of relation.repartition(10), instead of relation.distributeBy('a, 10). It's not about which one is cheaper, we should not surprise users.

cloud-fan · 2017-03-07T22:31:14Z

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

+
+    val query2 = testRelation
+      .coalesce(20)
+      .distribute('a)(30)


I'd like to make query2 as

testRelation .coalesce(30) .distribute('a)(20)

i.e. the numPartitions of coalesce is bigger than distribute

cloud-fan · 2017-03-07T22:32:16Z

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

+  }
+
+  test("repartition above repartitionBy") {
+    val query1 = testRelation


we can add a comment: // Always respects the top repartition amd removes useless distribute below repartition

cloud-fan · 2017-03-07T22:33:10Z

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

+  test("repartition above repartitionBy") {
+    val query1 = testRelation
      .distribute('a)(20)
      .repartition(10)


we can still pick the same numPartition pairs: 10, 20 and 30, 20

cloud-fan · 2017-03-07T22:33:50Z

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

+      .distribute('a)(20)
+
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer2 = testRelation.distribute('a)(20).analyze


it's same with correctAnswer1

gatorsmile · 2017-03-07T22:41:31Z

Let me re-write the whole test cases. Thanks!

SparkQA · 2017-03-08T00:07:18Z

Test build #74137 has finished for PR 16933 at commit 8306c49.

This patch fails PySpark unit tests.
This patch merges cleanly.
This patch adds no public classes.

cloud-fan · 2017-03-08T00:16:15Z

LGTM, pending test

SparkQA · 2017-03-08T00:40:27Z

Test build #74140 has finished for PR 16933 at commit d69c5a1.

This patch fails PySpark unit tests.
This patch merges cleanly.
This patch adds no public classes.

cloud-fan · 2017-03-08T00:52:49Z

retest this please

SparkQA · 2017-03-08T02:51:18Z

Test build #74157 has finished for PR 16933 at commit d69c5a1.

This patch fails PySpark unit tests.
This patch merges cleanly.
This patch adds no public classes.

cloud-fan · 2017-03-08T04:36:18Z

retest this please

SparkQA · 2017-03-08T07:06:44Z

Test build #74179 has finished for PR 16933 at commit d69c5a1.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

gatorsmile · 2017-03-08T17:35:38Z

Thanks! Merging to master.

felixcheung · 2017-03-09T06:45:29Z

awesome!

fix.

7b4a9dd

gatorsmile mentioned this pull request Feb 14, 2017

[SPARK-19399][SPARKR] Add R coalesce API for DataFrame and Column #16739

Closed

dongjoon-hyun reviewed Feb 14, 2017

View reviewed changes

felixcheung reviewed Feb 15, 2017

View reviewed changes

gatorsmile added 4 commits February 18, 2017 19:05

temp

7722781

Merge remote-tracking branch 'upstream/master' into CollapseRepartition

f9483cb

fix.

0f95a6f

clean

680c3af

cloud-fan reviewed Mar 6, 2017

View reviewed changes

address comments.

5453ad4

cloud-fan reviewed Mar 6, 2017

View reviewed changes

inline.

4649af4

cloud-fan reviewed Mar 7, 2017

View reviewed changes

fix

8306c49

cloud-fan reviewed Mar 7, 2017

View reviewed changes

gatorsmile added 2 commits March 7, 2017 14:49

fix

379a15a

style clean

d69c5a1

asfgit closed this in 9a6ac72 Mar 8, 2017

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

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

Uh oh!

Conversation

gatorsmile commented Feb 14, 2017 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

How was this patch tested?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SparkQA commented Feb 15, 2017

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gatorsmile Feb 19, 2017 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SparkQA commented Mar 4, 2017

Uh oh!

SparkQA commented Mar 5, 2017

Uh oh!

gatorsmile commented Mar 6, 2017

Uh oh!

cloud-fan Mar 6, 2017 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SparkQA commented Mar 7, 2017

Uh oh!

SparkQA commented Mar 7, 2017

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gatorsmile commented Mar 7, 2017

Uh oh!

SparkQA commented Mar 8, 2017

Uh oh!

cloud-fan commented Mar 8, 2017

gatorsmile commented Feb 14, 2017 •

edited

Loading

gatorsmile Feb 19, 2017 •

edited

Loading

cloud-fan Mar 6, 2017 •

edited

Loading