Spark源码分析之：Shuffle
这一篇我们来分析Spark2.1的Shuffle流程。
其实ShuffleDependency从SparkContext初始化就已经被DAGScheduler划分好了，本文主要探讨在Task运行过程中的ShufleWrite和ShuffleRead。
要从Task运行开始说起，就要知道Task在哪里运行的。我们普遍认为Executor是负责执行Task的，但是我们发现Executor其实就是一个类
private[spark] class Executor(){}而在一个Application提交后，用JPS命令查看，会发现有AplicationMaster进程和CoarseGrainedExecutorBackend进程。没错，后者就是真实管理运行Task的进程。
CoarseGrainedExecutorBackend
做为一个进程，CoarseGrainedExecutorBackend负责与ApplicationMaster进行RPC通信，其中有一个receive方法重写于RpcEndpoint，负责接收消息，当收到启动task的命令时，就会调用Executor的launchTask方法，此方法负责启动task。
Executor类中还有一个TaskRunner内部类，是一个实现了Runnable接口的线程。启动task也就是将序列化task的描述信息封装成一个TaskRunner，将其放在线程池中运行。
//Executor类
def launchTask(
context: ExecutorBackend,
taskId: Long,
attemptNumber: Int,
taskName: String,
serializedTask: ByteBuffer): Unit = {
// Runnable 接口的对象.
val tr = new TaskRunner(context, taskId = taskId, attemptNumber = attemptNumber, taskName,
serializedTask)
runningTasks.put(taskId, tr)
// 在线程池中执行 task
threadPool.execute(tr)
}由此我们看出，其实负责执行task的还是Executor，而与ApplicationMaster通信的却是CoarseGrainedExecutorBackend，也就是CoarseGrainedExecutorBackend是一个中间代理商。
TaskRunner的run方法是task的具体实现逻辑。
@volatile var task: Task[Any] = _
override def run(): Unit = {
//...
// 1.更新 task 的状态
execBackend.statusUpdate(taskId, TaskState.RUNNING, EMPTY_BYTE_BUFFER)
// 2.把任务相关的数据反序列化出来
val (taskFiles, taskJars, taskProps, taskBytes) = Task.deserializeWithDependencies(serializedTask)
//3.反序列化确定task属于哪个类型
task = ser.deserialize[Task[Any]](taskBytes, Thread.currentThread.getContextClassLoader)
// 4.开始运行 task
val res = task.run(
taskAttemptId = taskId,
attemptNumber = attemptNumber,
metricsSystem = env.metricsSystem)上述代码第3步，确定task是哪个类型，Task是个抽象类，有以下实现类：
主要实现为：ShuffleMapTask和ResultTask。
Spark中的Stage分为两种，ShuffleMapStage和ResultStage，ResultStage就是finallStage，也就是action算子所在的Stage，前面的所有Stage都是ShuffleMapStage。ShuffleMapStage对应的Task就是ShuffleMapTask，而ResultStage对应的Task就是ResultTask。
ShufleWriter
也就是说，首先运行的task是ShuffleMapTask。那么上面第4步的task.run会调用ShuffleMapTask的runtask方法。
//ShuffleMapTask类
override def runTask(context: TaskContext): MapStatus = {
//...
var writer: ShuffleWriter[Any, Any] = null
try {
val manager = SparkEnv.get.shuffleManager
//1.获取ShuffleWriter，传入一个shuffleHandle
writer = manager.getWriter[Any, Any](dep.shuffleHandle, partitionId, context)
//2.写数据
writer.write(rdd.iterator(partition, context).asInstanceOf[Iterator[_ <: Product2[Any, Any]]])
writer.stop(success = true).get
} catch {
//...
}如上代码中所知，ShuffleWriter为ShuffleManager.getWriter所得。
getWriter
ShuffleManager为一个trait，其实现类只有1个，为SortShuffleManager。
并且传入了一个shuffleHandle，点进去发现在shuffleDependency方法下被初始化：
val shuffleHandle: ShuffleHandle = _rdd.context.env.shuffleManager.registerShuffle(
shuffleId, _rdd.partitions.length, this)点开此SortShuffleManager的registerShuffle方法：
override def registerShuffle[K, V, C](
shuffleId: Int,
numMaps: Int,
dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
if (SortShuffleWriter.shouldBypassMergeSort(SparkEnv.get.conf, dependency)) {
// If there are fewer than spark.shuffle.sort.bypassMergeThreshold partitions and we don't
// need map-side aggregation, then write numPartitions files directly and just concatenate
// them at the end. This avoids doing serialization and deserialization twice to merge
// together the spilled files, which would happen with the normal code path. The downside is
// having multiple files open at a time and thus more memory allocated to buffers.
new BypassMergeSortShuffleHandle[K, V](
shuffleId, numMaps, dependency.asInstanceOf[ShuffleDependency[K, V, V]])
} else if (SortShuffleManager.canUseSerializedShuffle(dependency)) {
// Otherwise, try to buffer map outputs in a serialized form, since this is more efficient:
new SerializedShuffleHandle[K, V](
shuffleId, numMaps, dependency.asInstanceOf[ShuffleDependency[K, V, V]])
} else {
// Otherwise, buffer map outputs in a deserialized form:
new BaseShuffleHandle(shuffleId, numMaps, dependency)
}
}此方法用于注册shuffle，也就是选择使用何种shuffle。
先看第一种：BypassMergeSortShuffleHandle
private[spark] object SortShuffleWriter {
def shouldBypassMergeSort(conf: SparkConf, dep: ShuffleDependency[_, _, _]): Boolean = {
// We cannot bypass sorting if we need to do map-side aggregation.
//如果需要在map端聚合，如reduceByKey算子，则不能使用。
if (dep.mapSideCombine) {
require(dep.aggregator.isDefined, "Map-side combine without Aggregator specified!")
false
} else {//获取"spark.shuffle.sort.bypassMergeThreshold"参数的值，默认为200
val bypassMergeThreshold: Int = conf.getInt("spark.shuffle.sort.bypassMergeThreshold", 200)
//分区数 <= 200使用
dep.partitioner.numPartitions <= bypassMergeThreshold
}
}
}使用此Shuffle方式有两个条件：
- 在map端不能有聚合；
- 分区数<=200；
第二种：canUseSerializedShuffle
def canUseSerializedShuffle(dependency: ShuffleDependency[_, _, _]): Boolean = {
val shufId = dependency.shuffleId
val numPartitions = dependency.partitioner.numPartitions
//需要支持序列化重定向
if (!dependency.serializer.supportsRelocationOfSerializedObjects) {
log.debug(s"Can't use serialized shuffle for shuffle $shufId because the serializer, " +
s"${dependency.serializer.getClass.getName}, does not support object relocation")
false
}
//不能在map端聚合
else if (dependency.aggregator.isDefined) {
log.debug(
s"Can't use serialized shuffle for shuffle $shufId because an aggregator is defined")
false
}
//分区数大于16777216(2^24)
else if (numPartitions > MAX_SHUFFLE_OUTPUT_PARTITIONS_FOR_SERIALIZED_MODE) {
log.debug(s"Can't use serialized shuffle for shuffle $shufId because it has more than " +
s"$MAX_SHUFFLE_OUTPUT_PARTITIONS_FOR_SERIALIZED_MODE partitions")
false
} else {
log.debug(s"Can use serialized shuffle for shuffle $shufId")
true
}
}使用此Shuffle方式有三个条件：
- 支持序列化重定向，Serializer可以对已经序列化的对象进行排序，这种排序起到的效果和先对数据排序再序列化一致。支持relocation的Serializer是KryoSerializer和SparkSQL的custom serializers。
- 不能在map端进行聚合。
- 分区数不能大于最大分区数2^24个。
如果不符合以上两种，则选择BaseShuffleHandle
接下来回过头来看getWriter方法：
override def getWriter[K, V](
handle: ShuffleHandle,
mapId: Int,
context: TaskContext): ShuffleWriter[K, V] = {
numMapsForShuffle.putIfAbsent(
handle.shuffleId, handle.asInstanceOf[BaseShuffleHandle[_, _, _]].numMaps)
val env = SparkEnv.get
// 根据不同的 Handle, 创建不同的 ShuffleWriter
handle match {
case unsafeShuffleHandle: SerializedShuffleHandle[K@unchecked, V@unchecked] =>
new UnsafeShuffleWriter(
env.blockManager,
shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver],
context.taskMemoryManager(),
unsafeShuffleHandle,
mapId,
context,
env.conf)
case bypassMergeSortHandle: BypassMergeSortShuffleHandle[K@unchecked, V@unchecked] =>
new BypassMergeSortShuffleWriter(
env.blockManager,
shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver],
bypassMergeSortHandle,
mapId,
context,
env.conf)
case other: BaseShuffleHandle[K@unchecked, V@unchecked, _] =>
new SortShuffleWriter(shuffleBlockResolver, other, mapId, context)
}
}据上所知,Handle对应的ShuffleWrite为：
- SerializedShuffleHandle -> UnsafeShuffleWriter
- BypassMergeSortShuffleHandle -> BypassMergeSortShuffleWriter
- BaseShuffleHandle -> SortShuffleWriter
Write
BypassMergeSortShuffleWriter
此ShuffleWriter缺点比较大，仅供不需要聚合排序，分区数小于等于200时使用，其内部原理简单，但是需要注意其中一行代码：
public void write(Iterator> records) throws IOException {
//...
partitionWriters = new DiskBlockObjectWriter[numPartitions];
//...
} DiskBlockObjectWriter为缓冲块block，此方式为每个partition都开辟了一块32K大小的block缓冲区，当partition数量为10000个时，单一个MapTask就需要约320M内存，内存消耗过大，这也就是分区数小于200个使用的原因。
最后会将数据封装成一个MapStatus发送给Driver。
SortShuffleWriter
override def write(records: Iterator[Product2[K, V]]): Unit = {
// 一：
sorter = if (dep.mapSideCombine) {
require(dep.aggregator.isDefined, "Map-side combine without Aggregator specified!")
new ExternalSorter[K, V, C](
context, dep.aggregator, Some(dep.partitioner), dep.keyOrdering, dep.serializer)
} else {
// In this case we pass neither an aggregator nor an ordering to the sorter, because we don't
// care whether the keys get sorted in each partition; that will be done on the reduce side
// if the operation being run is sortByKey.
new ExternalSorter[K, V, V](
context, aggregator = None, Some(dep.partitioner), ordering = None, dep.serializer)
}
// 将 Map 任务的输出记录插入到缓存中
sorter.insertAll(records)
// Don't bother including the time to open the merged output file in the shuffle write time,
// because it just opens a single file, so is typically too fast to measure accurately
// (see SPARK-3570).
// 数据 shuffle 数据文件
val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
val tmp = Utils.tempFileWith(output)
try { // 将 map 端缓存的数据写入到磁盘中, 并生成 Block 文件对应的索引文件.
val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
// 记录各个分区数据的长度
val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
// 生成 Block 文件对应的索引文件. 此索引文件用于记录各个分区在 Block文件中的偏移量, 以便于
// Reduce 任务拉取时使用
shuffleBlockResolver.writeIndexFileAndCommit(dep.shuffleId, mapId, partitionLengths, tmp)
mapStatus = MapStatus(blockManager.shuffleServerId, partitionLengths)
} finally {
if (tmp.exists() && !tmp.delete()) {
logError(s"Error while deleting temp file ${tmp.getAbsolutePath}")
}
}
}- 根据不同的场景选取不同的方式。如果map端开启了聚合，则
new ExternalSorter[K, V, C](context, dep.aggregator, Some(dep.partitioner), dep.keyOrdering, dep.serializer)
传入了一个聚合器和一个排序器。而如果没有聚合，则
new ExternalSorter[K, V, V](context, aggregator = None, Some(dep.partitioner), ordering = None, dep.serializer)
所以只有在map端需要聚合的算子才会排序，像sortByKey这种算子map不会排序，而是放在reduce端排序。
- 接下来一步 sorter.insertAll(records)
records就是write方法的形参，也就是我们传过来的数据。那这个insertAll是干嘛的？
insertAll方法是我们第一步new出来的ExternalSorter类中的方法，
def insertAll(records: Iterator[Product2[K, V]]): Unit = {
// TODO: stop combining if we find that the reduction factor isn't high
val shouldCombine = aggregator.isDefined
if (shouldCombine) {
// Combine values in-memory first using our AppendOnlyMap
val mergeValue = aggregator.get.mergeValue
val createCombiner = aggregator.get.createCombiner
var kv: Product2[K, V] = null
val update = (hadValue: Boolean, oldValue: C) => {
if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
}
//循环放数据
while (records.hasNext) {
//记录自上次spill以来内存中的数据条数
addElementsRead()
//数据
kv = records.next()
//PartitionedAppendOnlyMap中update数据
map.changeValue((getPartition(kv._1), kv._1), update)
//检查数据是否存放的下，如果存放不下则会扩展内存或者spill
maybeSpillCollection(usingMap = true)
}
} else {
// Stick values into our buffer
//PartitionedPairBuffer
while (records.hasNext) {
addElementsRead()
val kv = records.next()
buffer.insert(getPartition(kv._1), kv._1, kv._2.asInstanceOf[C])
maybeSpillCollection(usingMap = false)
}
}
}为了提高聚合和排序的性能，Spark为ShuffleWriter的聚合排序过程设计了两种数据结构：
- PartitionedAppendOnlyMap：如果设置了Aggregator，我们可以将对象放入AppendOnlyMap中以将它们组合在一起
- PartitionedPairBuffer：如果没有设置Aggregator，则放入buffer中。
下面我们先来看看这两个数据结构
PartitionedAppendOnlyMap：是AppendOnlyMap的间接子类，在ShuffleWriter端使用，其内部如下：
操作数据仅有一个insert方法。调用的是其父类AppendOnlyMap的update方法。
AppendOnlyMap
此类是一个类似于HashMap的数据结构，使用Hash表进行寻址，但不同于HashMap的是，AppendOnlyMap只有update操作（暂且认为添加也算update），因为Shuffle的数据不需要删除，只需要聚合，所以这样设计一个轻量级的HashMap合情合理。另外一点不同的是，AppendOnlyMap解决Hash碰撞的方法为具有2的幂的哈希表大小的开放地址法的二次探测法，而HashMap为链表法。另外需要注意，AppendOnlyMap最多可存放375809638（0.7 * 2 ^ 29）个元素，0.7为加载因子，也就是当利用率为70%时扩容为原来的2倍。
至于怎么实现添加操作，如果读者感兴趣可以搜此类的update方法详细了解。
PartitionedPairBuffer：本质上类似于一个Array数组，仅供排序使用，可按partitionID或partitionID+Key排序。最多支持1073741823 (2 ^ 30 - 1)个元素。
回过头来看上面的insertAll方法，如果有Aggregator，则使用PartitionedAppendOnlyMap，否则使用PartitionedPairBuffer。在最后，都有一个maybeSpillCollection操作，此方法是判断每条数据插入内存缓冲区之后，是否需要spill文件。
private def maybeSpillCollection(usingMap: Boolean): Unit = {
var estimatedSize = 0L
if (usingMap) {
//估计当前集合大小
estimatedSize = map.estimateSize()
if (maybeSpill(map, estimatedSize)) {//判断是否可以spill
//如果spill了，就新开辟一块新内存
map = new PartitionedAppendOnlyMap[K, C]
}
} else {
estimatedSize = buffer.estimateSize()
if (maybeSpill(buffer, estimatedSize)) {
buffer = new PartitionedPairBuffer[K, C]
}
}
//记录最大使用内存量
if (estimatedSize > _peakMemoryUsedBytes) {
_peakMemoryUsedBytes = estimatedSize
}
}首先估计了当前缓冲区（集合）的大小，这一步是比较困难的。因为虽然我们知道AppendOnlyMap中持有的数据的长度和大小，但是数组里面存放的是Key和Value的引用，并不是实际对象的大小，而且Value会不断更新，实际大小不断变化。因此想要准确的获取到其大小相当困难，如果简单的每次插入一条数据都扫描数组中的record并对其大小相加，那么时间复杂度太高会极大影响效率。此estimateSize方法是Spark设计的一个增量式的高效估算算法，在每个record插入或更新时根据历史统计值和当前变化量直接估算当前AppendOnlyMap的大小，算法时间复杂度为O(1)，开销很小。在record插入和聚合过程中会定期对当前AppendOnlyMap中的record进行抽样，然后精确计算这些record的总大小，总个数，更新个数及平均值等，并作为历史统计值。进行抽样是因为record个数过多，难以对每个record进行精确计算。之后，每当有record插入或更新时，会根据历史统计值和历史平均的变化值，估算AppendOnlyMap的总大小。
拿到估算大小之后，就会尝试spill，maybeSpill方法如下：
protected def maybeSpill(collection: C, currentMemory: Long): Boolean = {
var shouldSpill = false
//elementsRead：自上次spill后，读取的元素个数。myMemoryThreshold：默认5M
if (elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold) {
//申请内存，扩容2倍。
//扩容2倍需要开辟的新内存。比如已占用6M，则需要新申请 12-5 = 7M
val amountToRequest = 2 * currentMemory - myMemoryThreshold
//实际申请到的新内存
val granted = acquireMemory(amountToRequest)
//现在实际拥有的内存。在内存足够用也就是可以全部申请来的时候，myMemoryThreshold=2*currentMemory
myMemoryThreshold += granted
//如果拥有的内存小于已使用的内存，则会spill
shouldSpill = currentMemory >= myMemoryThreshold
}
//如果数据量过大，已经超出了long的最大值，则仍然会spill
shouldSpill = shouldSpill || _elementsRead > numElementsForceSpillThreshold
if (shouldSpill) {// 执行spill
_spillCount += 1
logSpillage(currentMemory)
spill(collection)
_elementsRead = 0
_memoryBytesSpilled += currentMemory
releaseMemory()
}
shouldSpill
}elementsRead为缓冲区元素个数，如果是32的倍数，并且使用内存已经大于默认的spill阈值（默认5M），则扩容两倍。如果扩容之后的容量还不足以放下当前使用容量，则溢出。举个例子：假若当前使用了8M内存，也就是currentMemory=8M，而myMemoryThreshold=5M，所以会扩容到16M，需要新申请16-5=11M，但是我只申请到了1M，那么此时myMemoryThreshold=6M 小于 currentMemory，就需要溢出。另外如果元素个数超过了Long类型的最大值，则也会spill。
在insertAll方法之后，会将数据封装成一个MapStatus发送给Driver。
UnsafeShuffleWriter
此方式是Spark2.0为提高效率新加入的方式，核心思想是在堆外内存中操作序列化的record对象（二进制数据），降低内存消耗和GC开销。正好弥补了分区数超过200时BypassMergeSortShuffleWriter的不足。
Stop
至此stop，Shuffle的Wirte阶段结束。