RangePartitioner

概要

这篇文章是Partitioner博客的一部分，内容较多，单独介绍

水塘抽样算法

RangePartitioner基于水塘抽样算法实现，其目的在于从包含n个项目的集合S中选取k个样本，其中n为一很大或未知的数量，尤其适用于不能把所有n个项目都存放到内存的情况。算法如下

从S中抽取首k项放入「水塘」中
对于每一个S[j]项（j ≥ k）：
   随机产生一个范围0到j的整数r
   若 r < k 则把水塘中的第r项换成S[j]项

运行时机制简述

RangePartitioner运行时机制，可以概述为如何选取分区的分割符，如下

我们假设待处理的数据的key均为英文大写字母且在A-Z之间，Partition数为3，则RangePartitioner运行时会选出两个分隔符，假设为H和S(如何选择后续介绍)，则key在字母A-H之间的(<=H)，属于Partition 0，同理I-S属于Partition 1，T-Z属于Partition 2。

举个例子

val data = sc.parallelize(List("a c", "a b", "b c", "b d", "c d"), 2)
val range = data.flatMap(_.split(" ")).map((_, 1))
range.partitionBy(new RangePartitioner(3, range)).collect()

这是个简单的Wordcount例子，只是使用了partitionBy将Partitioner设置为RangePartitioner，通过打印log，我们可以发现上述代码选取的分隔符为b和c,则key为a和b的数据属于Partition 0(<=b)，c属于Partition 1(<=c)，d属于Partition 2，结果如下

了解了其运行时的机制，接下来查看源码，了解这个分割符是如何选择的。
#RangePartitioner源码
先回顾Partitioner的定义，其有两个抽象方法

接着查看RangePartitioner及其实现的父类的抽象方法

我们可以发现，这两个方法的实现均用到了变量rangeBounds，其类型为Array[K]，内容就是我们上面提到的分区分割符，查看对应代码

其算法就是水塘抽样算法(reservoir sampling)，所以如上图所示，138、139行先确定了总样本大小sampleSize，为每个Partition20条，最多不超过1e6条，每个Partition的样本数sampleSizePerPartition，总样本除以分区数，但是乘以了系数3.0，这么做是为了保证分区数少时能收集更多样本，接着140行调用了RangePartitioner的sketch方法，查看此方法

sketch方法很简短，根据rdd的id获得抽样的seed(用于reservoir sampling中产生随机数)，然后调用
SamplingUtils.reservoirSampleAndCount方法(reservoir sampling算法)，接下来我们看这个算法怎么实现的

逻辑很清晰，3个参数，原始数据(input)、分区样本数(k)、seed(用于产生随机数)，结合截图中的注释，主要分为两步，1、先获取大小为k(sampleSizePerPartition)的样本数，如果记录数小于sampleSizePerPartition直接返回样本，2、否则使用seed生成一个随机数生成器rand，继续遍历数据，同时每条数据对应生成一个随机数，随机数小于k，则把1步中的样本数据下标对应的数据替换掉，这就是完整的reservoir sampling算法的逻辑。reservoirSampleAndCount方法返回(样本，分区记录总数)，回到sketch方法252行，继续执行，256行求出总记录数numItems，sketch结束，返回结果**(总记录数，Array[(PartitionId, 对应分数记录数, 样本)])**

回到变量rangeBounds处(源码第三幅图140行)，接下来的代码如注释所示，如果fraction * 分区内记录数 > sampleSizePerPartition，则该分区会再进行一次抽样，否则计算权重weight，为分区记录总数/分区样本数，最后调用RangePartitioner的determineBounds方法求分区分隔符(第三张源码图的补充，如下)，参数为Array[(key, weight)]和分区数

查看determineBounds方法

逻辑如下，先将candidate(Array[(key, weight)])按照key排序，计算总权重sumWeights，除以分区数，得到每个分区的平均权重step，接下来while循环遍历已排序的candidate，累加其权重cumWeight，每当累加权重达到一个分区的平均权重step，就获取一个key作为分区间隔符，最后返回所有获取到的分隔符，determineBounds执行完毕，也就返回了变量rangeBounds(分区分隔符)。

至此，rangeBounds就是分区分隔符，以及如何计算的就理清了，再结合前一小节运行时机制简述中分隔符的作用，理解RangePartitioner的numPartitions和getPartition方法实现应该没有问题了。回到第二幅源码图，numPartitions方法，分隔符的长度加1就是分区数了，getPartition，先判断间隔数数，少于128，直接遍历比较key和分隔符，得到PartitionId，否则使用二分查找，并做了边界条件的判断，最后，根据升序还是降序返回PartitionId，整个RangePartitioner的源码至此基本看完了。
#总结
通过源码深入了解了RangePartitioner的实现机制，如下

1. 使用reservoir Sample抽样方法，对每个Partition进行抽样
2. 计算权重， 对数据多(大于sampleSizePerPartition)的分区再进行抽样
3. 由权重信息计算分区分隔符rangeBounds
4. 由rangeBounds计算分区数和key属于哪个分区

此外，RDD的transformation，sortBy、sortByKey，使用RangePartitioner实现。