diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index a0ac7a93429267ffe697094f668f6373b3266693..dd670a9b3db2ac86dff00ef4e3bf21ee93fe8584 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -464,10 +464,10 @@ class DataFrame(object):
+---+-----+
|age| name|
+---+-----+
- | 5| Bob|
- | 5| Bob|
| 2|Alice|
+ | 5| Bob|
| 2|Alice|
+ | 5| Bob|
+---+-----+
>>> data.rdd.getNumPartitions()
7
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
index 84c9cc8c8e7f734d585f19b65a9077ae1eadb5a7..5181dcc786a3dffea028fa067e595e2aa34bd8a8 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
@@ -370,8 +370,11 @@ package object dsl {
case plan => SubqueryAlias(alias, plan)
}
- def distribute(exprs: Expression*): LogicalPlan =
- RepartitionByExpression(exprs, logicalPlan)
+ def repartition(num: Integer): LogicalPlan =
+ Repartition(num, shuffle = true, logicalPlan)
+
+ def distribute(exprs: Expression*)(n: Int = -1): LogicalPlan =
+ RepartitionByExpression(exprs, logicalPlan, numPartitions = if (n < 0) None else Some(n))
def analyze: LogicalPlan =
EliminateSubqueryAliases(analysis.SimpleAnalyzer.execute(logicalPlan))
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index 03d15eabdd8637c0554f38927bb8b7dc5db7882a..368e9a539630b62feca8be6a2ca6b991a4a35a7f 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -556,12 +556,27 @@ object CollapseProject extends Rule[LogicalPlan] {
}
/**
- * Combines adjacent [[Repartition]] operators by keeping only the last one.
+ * 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.
*/
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 = Some(numPartitions))
+ // Case 3
+ case RepartitionByExpression(exprs, Repartition(_, _, child), numPartitions) =>
+ RepartitionByExpression(exprs, child, numPartitions)
}
}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseRepartitionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseRepartitionSuite.scala
new file mode 100644
index 0000000000000000000000000000000000000000..8952c72fe42fec2018426050be90796d7a407894
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseRepartitionSuite.scala
@@ -0,0 +1,78 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+
+class CollapseRepartitionSuite extends PlanTest {
+ object Optimize extends RuleExecutor[LogicalPlan] {
+ val batches =
+ Batch("CollapseRepartition", FixedPoint(10),
+ CollapseRepartition) :: Nil
+ }
+
+ val testRelation = LocalRelation('a.int, 'b.int)
+
+ test("collapse two adjacent repartitions into one") {
+ val query = testRelation
+ .repartition(10)
+ .repartition(20)
+
+ val optimized = Optimize.execute(query.analyze)
+ val correctAnswer = testRelation.repartition(20).analyze
+
+ comparePlans(optimized, correctAnswer)
+ }
+
+ test("collapse repartition and repartitionBy into one") {
+ val query = testRelation
+ .repartition(10)
+ .distribute('a)(20)
+
+ val optimized = Optimize.execute(query.analyze)
+ val correctAnswer = testRelation.distribute('a)(20).analyze
+
+ comparePlans(optimized, correctAnswer)
+ }
+
+ test("collapse repartitionBy and repartition into one") {
+ val query = testRelation
+ .distribute('a)(20)
+ .repartition(10)
+
+ val optimized = Optimize.execute(query.analyze)
+ val correctAnswer = testRelation.distribute('a)(10).analyze
+
+ comparePlans(optimized, correctAnswer)
+ }
+
+ test("collapse two adjacent repartitionBys into one") {
+ val query = testRelation
+ .distribute('b)(10)
+ .distribute('a)(20)
+
+ val optimized = Optimize.execute(query.analyze)
+ val correctAnswer = testRelation.distribute('a)(20).analyze
+
+ comparePlans(optimized, correctAnswer)
+ }
+}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/PlanParserSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/PlanParserSuite.scala
index 456948d6455c9df98cf9d912ce7377e29a5456d2..fbe236e196268a0819562022b3c23746fb269cf7 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/PlanParserSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/PlanParserSuite.scala
@@ -151,9 +151,9 @@ class PlanParserSuite extends PlanTest {
("", basePlan),
(" order by a, b desc", basePlan.orderBy('a.asc, 'b.desc)),
(" sort by a, b desc", basePlan.sortBy('a.asc, 'b.desc)),
- (" distribute by a, b", basePlan.distribute('a, 'b)),
- (" distribute by a sort by b", basePlan.distribute('a).sortBy('b.asc)),
- (" cluster by a, b", basePlan.distribute('a, 'b).sortBy('a.asc, 'b.asc))
+ (" distribute by a, b", basePlan.distribute('a, 'b)()),
+ (" distribute by a sort by b", basePlan.distribute('a)().sortBy('b.asc)),
+ (" cluster by a, b", basePlan.distribute('a, 'b)().sortBy('a.asc, 'b.asc))
)
orderSortDistrClusterClauses.foreach {