Spark源码系列:DataFrame repartition、coalesce 对比

/**
* Returns a new [[DataFrame]] that has exactly `numPartitions` partitions.
* @group dfops
* @since 1.3.0
*/
def repartition(numPartitions: Int): DataFrame = withPlan {
Repartition(numPartitions, shuffle = true, logicalPlan)
}

/**
* Returns a new [[DataFrame]] partitioned by the given partitioning expressions preserving
* the existing number of partitions. The resulting DataFrame is hash partitioned.
*
* This is the same operation as "DISTRIBUTE BY" in SQL (Hive QL).
*
* @group dfops
* @since 1.6.0
*/
@scala.annotation.varargs
def repartition(partitionExprs: Column*): DataFrame = withPlan {
RepartitionByExpression(partitionExprs.map(_.expr), logicalPlan, numPartitions = None)
}

/**
* Returns a new [[DataFrame]] partitioned by the given partitioning expressions into
* `numPartitions`. The resulting DataFrame is hash partitioned.
*
* This is the same operation as "DISTRIBUTE BY" in SQL (Hive QL).
*
* @group dfops
* @since 1.6.0
*/
@scala.annotation.varargs
def repartition(numPartitions: Int, partitionExprs: Column*): DataFrame = withPlan {
RepartitionByExpression(partitionExprs.map(_.expr), logicalPlan, Some(numPartitions))
}

/**
* Returns a new [[DataFrame]] that has exactly `numPartitions` partitions.
* Similar to coalesce defined on an [[RDD]], this operation results in a narrow dependency, e.g.
* if you go from 1000 partitions to 100 partitions, there will not be a shuffle, instead each of
* the 100 new partitions will claim 10 of the current partitions.
* @group rdd
* @since 1.4.0
*/
def coalesce(numPartitions: Int): DataFrame = withPlan {
Repartition(numPartitions, shuffle = false, logicalPlan)
}

/**
* 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 {
override def output: Seq[Attribute] = child.output
}

/**
* This method repartitions data using [[Expression]]s into `numPartitions`, and receives
* information about the number of partitions during execution. Used when a specific ordering or
* distribution is expected by the consumer of the query result. Use [[Repartition]] for RDD-like
* `coalesce` and `repartition`.
* If `numPartitions` is not specified, the number of partitions will be the number set by
* `spark.sql.shuffle.partitions`.
*/
case class RepartitionByExpression(
partitionExpressions: Seq[Expression],
child: LogicalPlan,
numPartitions: Option[Int] = None) extends RedistributeData {
numPartitions match {
case Some(n) => require(n > 0, "numPartitions must be greater than 0.")
case None => // Ok
}
} 

// 获取一个测试的DataFrame 里面包含一个user字段
val testDataFrame: DataFrame = readMysqlTable(sqlContext, "MYSQLTABLE", proPath)
// 获得10个分区的DataFrame
testDataFrame.repartition(10) 

// 获取一个测试的DataFrame 里面包含一个user字段
val testDataFrame: DataFrame = readMysqlTable(sqlContext, "MYSQLTABLE", proPath)
// 根据 user 字段进行分区，分区数量由 spark.sql.shuffle.partition 决定
testDataFrame.repartition($"user")

// 获取一个测试的DataFrame 里面包含一个user字段
val testDataFrame: DataFrame = readMysqlTable(sqlContext, "MYSQLTABLE", proPath)
// 根据 user 字段进行分区，将获得10个分区的DataFrame，此方法有时候在join的时候可以极大的提高效率，但是得注意出现数据倾斜的问题
testDataFrame.repartition(10,$"user") 

val testDataFrame1: DataFrame = readMysqlTable(sqlContext, "MYSQLTABLE", proPath)
val testDataFrame2=testDataFrame1.repartition(10)
// 不会触发shuffle
testDataFrame2.coalesce(5)
// 触发shuffle 返回一个100分区的DataFrame
testDataFrame2.coalesce(100)