Spark Shuffle 模块③ - Sort Based Shuffle writeSpark Shuffle 模块③ - Sort Based Shuffle write

Spark Shuffle 模块③ - Sort Based Shuffle write

本文为 Spark 2.0 源码剖析，其他版本可能有所不同

自 Spark 1.2 起，Sort Based Shuffle 替代 Hash Based Shuffle 成为 Spark 默认的 Shuffle 策略。

Shuffle Map Task 会按照 key 相对应的 partition id 进行排序，对于属于同一个 partition 的 keys 可选的进行或不进行排序。因为对于不需要排序的操作来说，这个排序是有损性能的。对于那些需要 Sort 的操作，比如 sortByKey，这个排序是由 Reducer 完成的。整个过程的如图所示：

records 的插入、排序、spill 及最终的合并操作都封装在 SortShuffleWriter#write(records: Iterator[Product2[K, V]]): Unit 中，其核心流程如下（注意：下面的流程针对于一个 partition）：

我们自上而下进行剖析

插入 records 至 sorter

Step1：创建 sorter: ExternalSorter

创建 ExternalSorter 对象 sorter 分两种情况：

• 需要在 map 端做 combine 操作：要提供聚合（Aggregator）和可选的 key 排序器（keyOrdering）来构造 ExternalSorter 对象，下文称 “mapCombine ExternalSorter”
• 不需要再 map 端做 combine 操作：不需要 map 端聚合，也不需要 key 排序器来构造 ExternalSorter 对象，下文称 “非 mapCombine ExternalSorter”

这两种不同的 ExternalSorter 对象的不同行为会在下文中说明

Step2：调用 ExternalSorter#insertAll( records: Iterator... ) 插入一个 partition 数据

这一步也分为两种情况

需要在 map 端做 combine

核心流程如下：

以 WordCount 中的 reduceByKey 为例：

textFile.map(word => (word, 1)).reduceByKey( (p1, p2) => p1 + p2 )

首先得到以下几个函数变量（这一小节主要关注 insertAll 行为，暂不关注 maybeSpillCollection 方法）：

• createCombiner 用于创建某个 key 对应的初始聚合值。在这里为 (p: Int) => p，即 key 第一次出现的 record 的 value 即为初始聚合值
• mergeValue 用于在 map 端将同一个 partition 中 key 相同的 value 合并到一起。在这里，为 (p1: Int, p2: Int) => p1 + p2
• mergeCombiners 用于在 reduce 端将不同 partition 中 key 相同的 value 合并到一起。在这里，为 (p1: Int, p2: Int) => p1 + p2

之后，遍历 records 并调用 map.changeValue((getPartition(kv._1), kv._1), update) 进行 combine，这里的 map 是 PartitionedAppendOnlyMap 类型，是一个只能添加数据不能移除数据的 HashMap，该 HashMap 最终也是基于数组来实现，数组中的 a[i] = key i, a[i+1]=value i，当添加的元素超过一定的指标时，数组会进行 grow，容量翻倍，同时所有的 key、value 会进行 rehash 重新分配位置。combine 的过程如 update 定义，对于已经有聚合结果的新值进行合并，将没有聚合结果的新值作为初始的聚合结果值。

在每 insert 一条 record 之后都会调用 maybeSpillCollection(usingMap = true) 来判断是否需要进行 spill，这将在后面的内容中介绍。

不需要在 map 端做 combine

核心流程如下：

这里的 buffer 是 PartitionedPairBuffer 类型，是一个只支持添加的 PairBuffer，底层实现是一个数组，存储的 pair，即 key、value 对，key 的类型为 (Int, K)，即 Partition Id 和 key。同一个 k、v 对被放在数组相邻的位置，与 PartitionedAppendOnlyMap 相同。

每插入一条 record 会调用 maybeSpillCollection(usingMap = false) 方法，这同样将在后文中说明。

spill 内存 buffer 至文件

主要流程如下：

第一个值得关注的点是：map 和 buffer 是如何估算占用了多少内存？实现如下：

val extrapolatedDelta = bytesPerUpdate * (numUpdates - samples.last.numUpdates)
(samples.last.size + extrapolatedDelta).toLong

即：以 bytePerUpdate 作为最近平均每次更新时的大小，用当前的 update 次数减去最后一个Sample 的 update 次数，然后乘以 bytePerUpdate，结果加上最后一个 Sample 记录的大小

另一个重要的点是 maybeSpill 方法，该方法首先确定是否有必要进行 spill，如果有，则将当前内存集合中溢写到磁盘上。在执行溢写前会尝试申请更多内存。有两种情况会需要进行 spill：

1. 当前集合包含的 records 数超过 spark.shuffle.spill.numElementsForceSpillThreshold 指定值，该值默认大小为 Long.MaxValue
2. 这种情况复杂些，其判断流程如下：

主要是在进行 spill 之前会尝试申请更多的内存来存放 records 来避免 spill。

若需要进行 spill，则首先会进行 spill 操作，然后释放已 spill 的集合对应的内存，释放主要是释放 execution pool 内存以供其他 tasks 使用，并将 myMemoryThreshold 赋值为初始值 spark.shuffle.spill.initialMemoryThreshold 对应值，即初始值

spill 操作

spill 的操作要考虑到之后要对之后生成的 spill 文件做 merge，因为最终一个 Shuffle Map Task 只生成一个输出文件和 Index 文件。

如果是需要做 map 端 combine，spill 时会对 map 中的数据先按 partition id 进行排序，若也提供了 key comparator，则会对属于同一个 partition 的 records 按 key 进行排序。做完排序后，先进行序列化再写入磁盘文件。

如果是不需要做 map 端的 combine，则只需对 buffer 按 partition id 进行排序即可，不需要对同一partition 的 records 按 key 进行排序。排序后，同样先序列化，再写入磁盘文件。

之后做 merge 时，使用 SpillReader 来读取 spill 数据又要先反序列化，再做最终排序，再写入最终文件，这一过程是 shuffle 过程中消耗比较大的一部分。

合并数据

合并的核心流程如下，由 ExternalSorter#writePartitionedFile(...) 方法实现

其中，最关键的 merge 流程如下：

1. 为每个 spill 出来的文件生成一个 reader: SpillReader，得到 readers: Seq[SpillReader]（reader 读取 spilled 文件要先反序列化）
2. 将内存集合进行 buffered，得到 inMemBuffered
3. 针对每个 partition p，执行以下操作：
   1. 取出 inMemBuffered 中属于 partition p 的 records，得到迭代器 inMemIterator
   2. 使用 readers 取出属于 partition p 的 records 对应的迭代器 Seq 与上一步中得到的 inMemIterator 合并得到最终的迭代器序列
   3. 如果定义了 map 端聚合，则先对上一步得到的迭代器序列中的 records 进行聚合，若还定义了 key comparator，则使用该 comparator 对 records 进行排序
   4. 若没有定义 map 端聚合但定义了 key comparator，则不做聚合而直接对 records 进行排序
   5. 若既没有定义 map 端聚合也没有定义 key comparator，则直接返回总体 records 对应的迭代器

参考

• http://dataknocker.github.io/2014/07/23/spark-appendonlymap
• http://www.cnblogs.com/devos/p/4805526.html