diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala index 24b9b810fc5c..157b377eb278 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -2420,6 +2420,13 @@ class Dataset[T] private[sql]( RepartitionByExpression(partitionExprs.map(_.expr), logicalPlan, numPartitions = None) } + /** + * Returns the number of partitions of this Dataset. + * @group basic + * @since 2.2.0 + */ + def getNumPartitions: Int = rdd.getNumPartitions + /** * Returns a new Dataset that has exactly `numPartitions` partitions. * Similar to coalesce defined on an `RDD`, this operation results in a narrow dependency, e.g. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala index cb7b97906a7d..044a72b678a5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala @@ -269,6 +269,7 @@ class DataFrameSuite extends QueryTest with SharedSQLContext { checkAnswer( testData.select('key).repartition(10).select('key), testData.select('key).collect().toSeq) + assert(testData.select('key).repartition(10).select('key).getNumPartitions == 10) } test("coalesce") { @@ -1301,12 +1302,14 @@ class DataFrameSuite extends QueryTest with SharedSQLContext { test("distributeBy and localSort") { val original = testData.repartition(1) assert(original.rdd.partitions.length == 1) + assert(original.getNumPartitions == 1) val df = original.repartition(5, $"key") assert(df.rdd.partitions.length == 5) + assert(df.getNumPartitions == 5) checkAnswer(original.select(), df.select()) val df2 = original.repartition(10, $"key") - assert(df2.rdd.partitions.length == 10) + assert(df2.getNumPartitions == 10) checkAnswer(original.select(), df2.select()) // Group by the column we are distributed by. This should generate a plan with no exchange