概述
Partitioner是shuffle过程中key重分区时的策略，即计算key决定k-v属于哪个分区，Transformation是宽依赖的算子时，父RDD和子RDD之间会进行shuffle操作，shuffle涉及到网络开销，由于父RDD和子RDD中的partition是多对多关系，所以容易造成partition中数据分配不均匀，导致数据的倾斜。
Shuffle
在MapReduce框架中，Shuffle是连接Map和和Reduce之间的桥梁，Map的输出要用到Reduce中必须经过Shuffle这个环节，Spark作为MapReduce框架的一种实现，自然也实现了shuffle的逻辑。
Shuffle描述的是一个过程，表现多对多的依赖关系，是Map和Reduce两个阶段的纽带，每个Reduce Task都会从Map Task产生的数据里读取其中的一片数据，如果Map Task的数目为m，Reduce Task数目为n (暂时认为Task的数目为partition的数目，因为partition的数目在不同的rdd中partition的个数也许会不同)，极端情况下可能会触发m*n个数据复制通道。因为Shuffle通常分为两部分：Map阶段的数据准备和Reduce阶段的数据复制。Map阶段需根据Reduce阶段的Task数量来决定每个Map Task输出的数据分片数目。
Shuffle的写操作是将MapTask操作后的数据写入到磁盘中，Spark中shuffle输出的ShuffleMapTask会为每个ReduceTask创建对应的bucket，ShuffleMapTask产生的结果会根据设置的partitioner中不同的key值得到对应的bucketId，然后写入到相应的bucket（同一个bucket存储的key值应该一致）中去。Shuffle的读操作由Socket或Netty完成，读出数据后如何切分partition暂不讨论。==BucketId其实也就是下游RDD的partitionId(或许不是Id，但是可以标识分区)，但是bucketId针对的是单个MapTask，partition针对的是整个RDD，bucket进行聚合后其实就是下游RDD中的partition。每个ShuffleMapTask的输出结果可能包含所有的ReduceTask需要的数据，所以每个ShuffleMapTask创建bucket的数目是和ReduceTask数目相等的。
Shuffle实现分为HashShuffleManager和SortShuffleManager，也可以自定义。综上，Partitioner在shuffle阶段发挥作用，依据Partitioner的逻辑计算key，得出对应的k-v属于哪个分区。
Partitioner定义
抽象类Partitioner定义了两个抽象方法numPartitions和getPartition。getPartition方法根据输入的k-v对的key值返回一个Int型数据。
abstract class Partitioner extends Serializable {
def numPartitions: Int
def getPartition(key: Any): Int
}
该抽象类伴生了一个Partitioner对象如下，主要包含defaultPartitioner函数，该函数定义了Partitioner的默认选择策略。如果设置了spark.default.parallelism，则使用该值作为默认partitions，否则使用上游RDD中partitions最大的数作为默认partitions。过滤出上游RDD中包含partitioner的RDD，选择包含有最大partitions并且isEligible的RDD，将该RDD中的partitioner设置为分区策略，否则返回一个带有默认partitions数的HashPartitioner作为Partitioner。Partition个数应该和partition个数最多的上游RDD一致，不然可能会导致OOM异常。
object Partitioner {
def defaultPartitioner(rdd: RDD[_], others: RDD[_]*): Partitioner = {
val rdds = (Seq(rdd) ++ others)
val hasPartitioner = rdds.filter(_.partitioner.exists(_.numPartitions > 0))
val hasMaxPartitioner: Option[RDD[_]] = if (hasPartitioner.nonEmpty) {
Some(hasPartitioner.maxBy(_.partitions.length))
} else {
None
}
val defaultNumPartitions = if (rdd.context.conf.contains("spark.default.parallelism")) {
rdd.context.defaultParallelism
} else {
rdds.map(_.partitions.length).max
}
if (hasMaxPartitioner.nonEmpty && (isEligiblePartitioner(hasMaxPartitioner.get, rdds) ||
defaultNumPartitions < hasMaxPartitioner.get.getNumPartitions)) {
hasMaxPartitioner.get.partitioner.get
} else {
new HashPartitioner(defaultNumPartitions)
}
}
private def isEligiblePartitioner(
hasMaxPartitioner: RDD[_],
rdds: Seq[RDD[_]]): Boolean = {
val maxPartitions = rdds.map(_.partitions.length).max
log10(maxPartitions) - log10(hasMaxPartitioner.getNumPartitions) < 1
}
}
Partitioner实现类
Partitioner主要有两个实现类：HashPartitioner和RangePartitioner，上面提到过HashPartitioner是默认的重分区方式。
HashPartitioner
numPartitions方法返回传入的分区数，getPartition方法使用key的hashCode值对分区数取模得到PartitionId，写入到对应的bucket中。因为Arrays的hashCodes值并不依赖于arrays的内容，导致hash函数将无法根据key值进行重分区，所以HashPartitioner不支持RDD[Array[_]] 或 RDD[(Array[_], _)]。
class HashPartitioner(partitions: Int) extends Partitioner {
require(partitions >= 0, s"Number of partitions ($partitions) cannot be negative.")
def numPartitions: Int = partitions
def getPartition(key: Any): Int = key match {
case null => 0
case _ => Utils.nonNegativeMod(key.hashCode, numPartitions)
}
override def equals(other: Any): Boolean = other match {
case h: HashPartitioner =>
h.numPartitions == numPartitions
case _ =>
false
}
override def hashCode: Int = numPartitions
}
RangePartitioner
RangePartitioner继承了Partitoner，Partitioner中定义了两个抽象方法numPartitions和getPartition，前者是重分区后的partition数，后者是获取某个key值重新分区后的partitionId，我们着重关注RangePartitioner类中这两个函数的实现。但我们先对rangeBounds这一函数进行分析，通过分析源码可以发现，该函数首先设置了一个样本大小(默认为每个partition包含20条数据)，再对父RDD中的每个partition需要抽取的样本数进行计算，调用工具类中的reservoir sampling算法对每个partition进行分别抽样。每个分区的记录数*fraction如果大于该partition中设定的样本数(这是由于不同的分区中包含的数据量不同，数据量较大的分区中抽样数将会大于平均值)，则用imbalancedPartitions存储，并重新抽样以确保每个分区中都有足够数量的样本。最后计算权重——分区记录总数/分区样本数，并调用determineBounds方法求分区分隔符。rangeBounds函数返回了一个数组，该数组中存储的是重分区的分隔符，简而言之，该方法的作用是尽可能均衡地将原有的数据进行重分区，并返回用于重分区的分隔符。例如数据为"a c", “a b”, “b c”, “b d”, “c d”，通过以上函数可能会将b、c、d作为分隔符，所有key为a的将会被重分区到同一个partition中(c和d类似)。
class RangePartitioner[K : Ordering : ClassTag, V](
partitions: Int,
rdd: RDD[_ <: Product2[K, V]],
private var ascending: Boolean = true,
val samplePointsPerPartitionHint: Int = 20)
extends Partitioner {
def this(partitions: Int, rdd: RDD[_ <: Product2[K, V]], ascending: Boolean) = {
this(partitions, rdd, ascending, samplePointsPerPartitionHint = 20)
}
private var rangeBounds: Array[K] = {
if (partitions <= 1) {
Array.empty
} else {
val sampleSize = math.min(samplePointsPerPartitionHint.toDouble * partitions, 1e6)
val sampleSizePerPartition = math.ceil(3.0 * sampleSize / rdd.partitions.length).toInt
val (numItems, sketched) = RangePartitioner.sketch(rdd.map(_._1), sampleSizePerPartition)
if (numItems == 0L) {
Array.empty
} else {
val fraction = math.min(sampleSize / math.max(numItems, 1L), 1.0)
val candidates = ArrayBuffer.empty[(K, Float)]
val imbalancedPartitions = mutable.Set.empty[Int]
sketched.foreach { case (idx, n, sample) =>
if (fraction * n > sampleSizePerPartition) {
imbalancedPartitions += idx
} else {
val weight = (n.toDouble / sample.length).toFloat
for (key <- sample) {
candidates += ((key, weight))
}
}
}
if (imbalancedPartitions.nonEmpty) {
val imbalanced = new PartitionPruningRDD(rdd.map(_._1), imbalancedPartitions.contains)
val seed = byteswap32(-rdd.id - 1)
val reSampled = imbalanced.sample(withReplacement = false, fraction, seed).collect()
val weight = (1.0 / fraction).toFloat
candidates ++= reSampled.map(x => (x, weight))
}
RangePartitioner.determineBounds(candidates, math.min(partitions, candidates.size))
}
}
}
}
接下来关注RangePartitioner类中两个重要函数——numPartitions和getPartition的实现。首先我们观察numPartitions，它的返回值为数组rangeBounds的长度加上1，这很容易理解，如果我们重分区的分隔符有n个，那么我们重分区后的partition便是n+1个。getPartition先判断间隔符的个数，如果小于128则直接遍历比较key和分隔符得到PartitionId，否则使用二分查找，并对边界条件进行了判断，最后根据构造RangePartitioner时传入的ascending参数确定是升序或降序返回PartitionId。
def numPartitions: Int = rangeBounds.length + 1
def getPartition(key: Any): Int = {
val k = key.asInstanceOf[K]
var partition = 0
if (rangeBounds.length <= 128) {
while (partition < rangeBounds.length && ordering.gt(k, rangeBounds(partition))) {
partition += 1
}
} else {
partition = binarySearch(rangeBounds, k)
if (partition < 0) {
partition = -partition-1
}
if (partition > rangeBounds.length) {
partition = rangeBounds.length
}
}
if (ascending) {
partition
} else {
rangeBounds.length - partition
}
}
总结RangePartitioner的过程为以下几步：
- 使用reservoir Sample方法对每个Partition进行分别抽样
- 对数据量大(大于sampleSizePerPartition)的分区进行重新抽样
- 由权重信息计算出分区分隔符rangeBounds
- 由rangeBounds计算分区数和key的所属分区
总结
综上所述，Partitioner主要的作用是在shuffle过程中对数据的partition进行重新分区，其主要实现的函数为：
- 获得重新分区的分区个数；
- 针对某个k-v对根据其中的key，将它按特定的方法重新分区。
Partitioner的具体实现类HashPartitioner和RangePartitioner对以上两个方法进行了重写。HashPartitioner主要是通过对原partition中数据的key值进行hash后，根据key的hash值将放入bucket中，再将不同partition中的bucket合并实现重新分区；RangePartitioner是先根据所有partition中数据的分布情况，尽可能均匀地构造出重分区的分隔符，再将数据的key值根据分隔符进行重新分区。HashPartitioner是大部分transformation的默认实现，sortBy、sortByKey使用RangePartitioner实现，也可以自定义Partitioner。以下是自定义Partitioner的一个例子，该Partitioner以key的长度进行分区：
class CustomPartitioner(partitions: Int) extends Partitioner {
def numPartitions: Int = partitions
def getPartition(key: Any): Int = {
val k = key.asInstanceOf[String]
return k.length() & (partitions -1)
}
}
具体的调用如下：
val data = sc.parallelize(List("a c", "a b", "bb c", "bb d", "c d"), 2)
data.flatMap(_.split(" ")).map((_, 1))
.partitionBy(new CustomPartitioner(2)).reduceByKey(_ + _)
.collect()
我们可以看到，Partitioner并不作用于reduceByKey而是作用于map，是因为它是对shuffle之前的partition进行重分区策略的定义，在reduceByKey过程中，会调用当前RDD的dep(父RDD)中定义的Partitioner进行重分区。