spark shuffle模块与hadoop的map-reduce模块功能类似，都是为了完成数据的洗牌，即要为每个分区的数据改变分区规则，重新分配分区。常见在诸如groupByKey，reduceByKey，repartition，join等操作都会触发shuffle过程。Spark的stage也是按照是否存在shuffle来划分的。当出现shuffle时，下游动作必须等到上游任务完成后才能执行。
一个shuffle的操作可以分为两个阶段，即map阶段和reduce阶段。在map阶段，executor将其所管理的分区数据根据指定的分区规则重新分区，并将中间结果输出到本地文件中，并将输出结果信息上报给driver节点。在reduce阶段，executor首先从driver节点获取到当前shuffle阶段所需要的中间文件在集群上的位置，然后向保存中间文件的节点请求获取中间结果文件，在进行相关处理（如聚合，排序）将数据返回给调用者，供调用者使用。
Shuffle结构
SortShuffleManager
在Spark 1.6.1中，SortShuffleManager是ShuffleManager的唯一子类，负责实现shuffle阶段的管理。主要包括getWriter，getReader，registerShuffle，unregisterShuffle等方法供map和reduce阶段使用。
在SortShuffleManager中，输入的记录会根据shuffle的分区规则将记录分配到对应的分区上，最终会将结果输出到一个map输出文件中，文件中的记录按照分区ID从小到大排序。每一个map输出文件对应一个分区索引文件，记录每个分区在文件中的偏移量。reduce任务会根据分区索引文件在map输出文件中找到需要分区的记录。当记录过多，不能在内存中装下时，会将内存中的记录临时刷新到磁盘上，这些临时文件最终会合成一个map输出文件。
registerShuffle
向SortShuffleManager注册一个shuffle，并获取对应的writer handler。
当前一共有三种类型的handler：BypassMergeSortShuffleHandle，SerializedShuffleHandle，BaseShuffleHandle，不同的handle在map阶段调用的writer不同，不同的writer在map时有不同的处理逻辑（如是否对每个分区中的记录进行聚合）
- BypassMergeSortShuffleHandle :不对shuffle后分区的记录进行聚合，对应writer BypassMergeSortShuffleWriter 只是简单的将输入记录输出到shuffle后对应的reduce分区文件中，最终将所有的分区文件合并成一个map输出文件。由于这个handle对应的writer会为每一个输入的分区建立一个临时文件，所以当输入分区很多时，会存在性能问题，所以使用BypassMergeSortShuffleHandle需要输入的分区数量小于 spark.shuffle.sort.bypassMergeThreshold 配置的大小（默认为200）。
- SerializedShuffleHandle：在不对输入记录进行反序列化的情况下，可以将记录按照shuffle后的reduce分区进行排序，并最终输出到一个map输出文件中。对应 UnsafeShuffleWriter。
- BaseShuffleHandle：如果需要对输入记录进行聚合，则使用这个方法将输入记录最终按照shuffle后的reduce分区id排序，每个分区中按照记录key的hash值进行排序。对应的writer为 SortShuffleWriter。
/**
* 向shuffleManager注册shuffle，并返回shuffle对应的handle
*/
override def registerShuffle[K, V, C](
shuffleId: Int,
numMaps: Int, // shuffle map的数量（通常就是rdd partition的数量）
dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
// 在没有指定在map阶段进行聚合且输入分区数小于bypassMergeThreshold（默认200）时，使用Bypass handle
if (SortShuffleWriter.shouldBypassMergeSort(SparkEnv.get.conf, dependency)) {
// 这里对在没指定map阶段聚合情况下，做了优化。如果partition数量大于bypassMergeThreshold，则
// 不能使用此方法。这个handle避免了在合并spill file时需要进行两次序列化/反序列化。
// 缺点是在同一时间需要打开多个文件，从而需要更多的文件缓存。
new BypassMergeSortShuffleHandle[K, V](
shuffleId, numMaps, dependency.asInstanceOf[ShuffleDependency[K, V, V]])
} else if (SortShuffleManager.canUseSerializedShuffle(dependency)) {
// 在输入的数据支持按分区id序列化重排 and 在map阶段不需要聚合
// 则使用此handle在map阶段直接对序列化数据进行操作
new SerializedShuffleHandle[K, V](
shuffleId, numMaps, dependency.asInstanceOf[ShuffleDependency[K, V, V]])
} else {
// 如果需要在map阶段进行聚合操作，使用非序列化数据进行map端聚合
new BaseShuffleHandle(shuffleId, numMaps, dependency)
}
}
getWriter
根据上文在向shuffleManager注册时返回给调用者的handle，来决定返回给调用者什么类型的writer来将输入记录输出到map输出文件中。
// 根据handle的不同返回不同的writer对象给当前的partition。被executor上的map任务调用
override def getWriter[K, V](
handle: ShuffleHandle,
mapId: Int, // mapId即是partitionId
context: TaskContext): ShuffleWriter[K, V] = {
numMapsForShuffle.putIfAbsent(
handle.shuffleId, handle.asInstanceOf[BaseShuffleHandle[_, _, _]].numMaps)
val env = SparkEnv.get
// 1.6.1版本支持的三种shuffleWriter
handle match {
// 直接使用序列化数据进行分区排序并输出到map文件
case unsafeShuffleHandle: SerializedShuffleHandle[K @unchecked, V @unchecked] =>
new UnsafeShuffleWriter(
env.blockManager,
shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver],
context.taskMemoryManager(),
unsafeShuffleHandle,
mapId,
context,
env.conf)
// 直接建立shuffle分区大小的临时输出文件，将输入记录直接写入对应的文件，并合并为一个map文件
case bypassMergeSortHandle: BypassMergeSortShuffleHandle[K @unchecked, V @unchecked] =>
new BypassMergeSortShuffleWriter(
env.blockManager,
shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver],
bypassMergeSortHandle,
mapId,
context,
env.conf)
// 当需要在map端对记录进行聚合操作时（相同的key记录进行合并），使用此方法
case other: BaseShuffleHandle[K @unchecked, V @unchecked, _] =>
new SortShuffleWriter(shuffleBlockResolver, other, mapId, context)
}
}
getReader
executor上的reduce任务调用此方法来获取map输出的数据。实际是返回一个BlockStoreShuffleReader实例。
/**
* 获取reduce partition，这个方法被executor上的reduce task所调用
*/
override def getReader[K, C](
handle: ShuffleHandle,
startPartition: Int,
endPartition: Int,
context: TaskContext): ShuffleReader[K, C] = {
new BlockStoreShuffleReader(
handle.asInstanceOf[BaseShuffleHandle[K, _, C]], startPartition, endPartition, context)
}
unregisterShuffle
// 删除shuffle数据，首先删除shuffleManager上的元数据（map数组）
// 再对每个partition删除磁盘上对应的data file和index file
override def unregisterShuffle(shuffleId: Int): Boolean = {
Option(numMapsForShuffle.remove(shuffleId)).foreach { numMaps =>
(0 until numMaps).foreach { mapId =>
shuffleBlockResolver.removeDataByMap(shuffleId, mapId)
}
}
true
}
ShuffleWriter
spark使用ShuffleWriter在map阶段将记录按照shuffle的分区输出到map文件中，供reduce阶段的任务来使用。ShuffleWriter在1.6.1版本中存在三个实现类：BypassMergeSortWriter（简单的建立shuffle分区数量个文件，将记录按shuffle后的reduce分区输出到对应的文件中，最后在合并成一个map文件），UnsafeShuffleWriter（直接将记录按照序列化的方式进行shuffle分区，并最终输出到map文件中），SortShuffleWriter（在map阶段对记录进行聚合，并最终输出到map文件中）。
ShuffleWriter抽象类有两个方法：
/** 将一系列的记录输出到map文件中 */
@throws[IOException]
def write(records: Iterator[Product2[K, V]]): Unit
/** 关闭当前writer，并返回map任务执行状态 */
def stop(success: Boolean): Option[MapStatus]
BypassMergeSortShuffleWriter
这个writer在写记录前，根据reduce分区数量，打开reduce分区数量个临时文件。根据reduce分区不同，将记录写入对应的临时文件中，然后将这些文件进行合并，输出到一个map文件中供reduce阶段使用。由于这个writer直接将记录写入对应的分区临时文件，所以记录不会在内存中缓存。如果shuffle后存在大量的reduce分区，则由于这个writer会为每一个reduce分区打开一个文件，所以在reduce分区众多时不会很高效。
这个writer只会在如下条件全部满足时才会被使用：
- 输入的rdd没有Ordering方法被指定
- 输入的rdd没有aggregator方法被指定
- 输入的rdd的分区数小于bypassMergeThreshold
write & writePartitionedFile
将输入的记录按照reduce分区进行排序并写入到输出的map文件中。
// 首先将记录写入对应reduce分区的临时文件中
public void write(Iterator<Product2<K, V>> records) throws IOException {
assert (partitionWriters == null);
if (!records.hasNext()) {
partitionLengths = new long[numPartitions];
shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, null);
mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths);
return;
}
final SerializerInstance serInstance = serializer.newInstance();
final long openStartTime = System.nanoTime();
partitionWriters = new DiskBlockObjectWriter[numPartitions];
// 对每一个reduce partition建立一个文件，所以如果partition多的话，同时会打开很多的文件
for (int i = 0; i < numPartitions; i++) {
final Tuple2<TempShuffleBlockId, File> tempShuffleBlockIdPlusFile =
blockManager.diskBlockManager().createTempShuffleBlock();
final File file = tempShuffleBlockIdPlusFile._2();
final BlockId blockId = tempShuffleBlockIdPlusFile._1();
partitionWriters[i] =
blockManager.getDiskWriter(blockId, file, serInstance, fileBufferSize, writeMetrics);
}
writeMetrics.incWriteTime(System.nanoTime() - openStartTime);
// 将记录根据写入的partition写入到不同的临时文件中
while (records.hasNext()) {
final Product2<K, V> record = records.next();
final K key = record._1();
partitionWriters[partitioner.getPartition(key)].write(key, record._2());
}
for (DiskBlockObjectWriter writer : partitionWriters) {
writer.commitAndClose();
}
// 将所有的reduce分区临时文件内容按照分区id合并到一个文件
File output = shuffleBlockResolver.getDataFile(shuffleId, mapId);
File tmp = Utils.tempFileWith(output);
partitionLengths = writePartitionedFile(tmp);
shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, tmp);
mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths);
}
private long[] writePartitionedFile(File outputFile) throws IOException {
final long[] lengths = new long[numPartitions];
if (partitionWriters == null) {
return lengths;
}
final FileOutputStream out = new FileOutputStream(outputFile, true);
final long writeStartTime = System.nanoTime();
boolean threwException = true;
try {
// 按照reduce分区id从0到n遍历
for (int i = 0; i < numPartitions; i++) {
// 如果存在当前reduce分区的临时文件，则将文件中的内容都追加到outputFile中
// 从而实现在outputFile中，记录按照分区id进行排序
final File file = partitionWriters[i].fileSegment().file();
if (file.exists()) {
final FileInputStream in = new FileInputStream(file);
boolean copyThrewException = true;
try {
lengths[i] = Utils.copyStream(in, out, false, transferToEnabled);
copyThrewException = false;
} finally {
Closeables.close(in, copyThrewException);
}
if (!file.delete()) {
logger.error("Unable to delete file for partition {}", i);
}
}
}
threwException = false;
} finally {
Closeables.close(out, threwException);
writeMetrics.incWriteTime(System.nanoTime() - writeStartTime);
}
partitionWriters = null;
return lengths;
}
UnsafeShuffleWriter
UnsafeShuffleWriter包含一个排序类 ShuffleExternalSorter。UnsafeShuffleWriter都是使用这个类来插入数据。
write方法实际会调用 ShuffleExternalSorter 的 insertRecode 方法来插入记录。
write & closeAndWriteOutput
将输入的记录按照shuffle的给定的分区方法进行分区，然后按照分区的id进行排序并输出到map文件中。
@Override
public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOException {
boolean success = false;
try {
// 将所有记录写入，最终会调用ShuffleExternalSorter中的insertRecord方法
while (records.hasNext()) {
insertRecordIntoSorter(records.next());
}
// 关闭writer，并将磁盘上的spill文件进行合并，输出到map文件中
closeAndWriteOutput();
success = true;
} finally {
// 省略代码
}
}
void closeAndWriteOutput() throws IOException {
assert(sorter != null);
updatePeakMemoryUsed();
serBuffer = null;
serOutputStream = null;
// 获取所有spilled的文件
final SpillInfo[] spills = sorter.closeAndGetSpills();
sorter = null;
final long[] partitionLengths;
// 获取最终输出的文件
final File output = shuffleBlockResolver.getDataFile(shuffleId, mapId);
// 创建最终输出文件的临时文件
final File tmp = Utils.tempFileWith(output);
try {
// 将spill文件合并到临时输出文件
partitionLengths = mergeSpills(spills, tmp);
} finally {
for (SpillInfo spill : spills) {
if (spill.file.exists() && ! spill.file.delete()) {
logger.error("Error while deleting spill file {}", spill.file.getPath());
}
}
}
// 将临时文件覆盖最终输出文件
shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, tmp);
mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths);
}
ShuffleExternalSorter
ShuffleExternalSorter用来将UnsafeShuffleWriter输入的记录保存并进行排序。输入的记录会首先被保存到内存的数据页中，当所有记录被插入或当前线程所使用的内存达到上限时，将内存页数据按照分区id进行排序后（使用ShuffleInMemorySorter方法），spill到磁盘中。
此sorter与SortShuffleWriter所使用的ExternalSorter不同，这个sorter不负责spill文件的合并，具体的合并操作由UnsafeShuffleWriter实现。UnsafeShuffleWriter支持直接对序列化的数据进行合并，从而避免了额外的序列化/反序列化操作。
这个类继承自MemoryConsumer，所以其所使用的内存来自execution memory，内存由TaskMemoryManager管理。当内存不足时，TaskMemoryManager会调用ShuffleExternalSorter的spill方法将记录刷新到磁盘上。
insertRecord
UnsafeShuffleWriter中的writer方法调用ShuffleExternalSorter中的insertRecord方法将数据写入。insertRecord实际调用ShuffleInMemorySorter中的insertRecord方法实际写入数据。
/**
* 将记录写入sorter中.
*/
public void insertRecord(Object recordBase, long recordOffset, int length, int partitionId)
throws IOException {
// 如果内存超过上限，则将内存中的记录spill到磁盘中
assert(inMemSorter != null);
if (inMemSorter.numRecords() > numElementsForSpillThreshold) {
spill();
}
growPointerArrayIfNecessary();
// Need 4 bytes to store the record length.
final int required = length + 4;
acquireNewPageIfNecessary(required);
assert(currentPage != null);
// 生成内存页地址
final Object base = currentPage.getBaseObject();
final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);
Platform.putInt(base, pageCursor, length);
pageCursor += 4;
Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);
pageCursor += length;
// 实际调用ShuffleInMemorySorter的insertRecord方法保存数据
inMemSorter.insertRecord(recordAddress, partitionId);
}
spill & writeSortedFile
当内存不足时，ShuffleExternalSorter会调用spill方法将内存中数据刷新到磁盘中。spill实际会调用writeSortedFile方法，首先将内存页中的数据按照reduce分区id进行排序（由ShuffleMemorySorter完成），然后输出到磁盘文件中。
/**
* 对内存中的记录进行排序并保存到文件中
*/
@Override
public long spill(long size, MemoryConsumer trigger) throws IOException {
// 省略代码
// 实际调用writeSorterFile进行spill操作
writeSortedFile(false);
final long spillSize = freeMemory();
inMemSorter.reset();
// 省略代码
// 返回实际释放的空间
return spillSize;
}
/**
* 对内存中的记录进行排序并将排好序的记录写入磁盘文件中。
*/
private void writeSortedFile(boolean isLastFile) throws IOException {
// 省略代码
// 获取内存中排好序的记录(根据shuffle后的partitionId进行排序)
// ShuffleInMemorySorter只在getSortedIterator方法被调用时，才会对内存中的数据进行排序
final ShuffleInMemorySorter.ShuffleSorterIterator sortedRecords =
inMemSorter.getSortedIterator();
// 获取写磁盘的writer
DiskBlockObjectWriter writer;
// 磁盘缓存
final byte[] writeBuffer = new byte[DISK_WRITE_BUFFER_SIZE];
// Because this output will be read during shuffle, its compression codec must be controlled by
// spark.shuffle.compress instead of spark.shuffle.spill.compress, so we need to use
// createTempShuffleBlock here; see SPARK-3426 for more details.
final Tuple2<TempShuffleBlockId, File> spilledFileInfo =
blockManager.diskBlockManager().createTempShuffleBlock();
final File file = spilledFileInfo._2();
final TempShuffleBlockId blockId = spilledFileInfo._1();
final SpillInfo spillInfo = new SpillInfo(numPartitions, file, blockId);
// 为了调用getDiskWriter而生成的序列化器实例，实际这个方法不会用到这个序列化器
final SerializerInstance ser = DummySerializerInstance.INSTANCE;
writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSizeBytes, writeMetricsToUse);
int currentPartition = -1;
// 将记录按照partition写入临时文件
while (sortedRecords.hasNext()) {
sortedRecords.loadNext();
final int partition = sortedRecords.packedRecordPointer.getPartitionId();
assert (partition >= currentPartition);
if (partition != currentPartition) {
// 如果已经遍历到下一个分区，则记录当前分区所包含的记录数
if (currentPartition != -1) {
writer.commitAndClose();
// 记录当前partition所包含的记录数
spillInfo.partitionLengths[currentPartition] = writer.fileSegment().length();
}
// 对每一个分区都会新建一个writer
currentPartition = partition;
writer =
blockManager.getDiskWriter(blockId, file, ser, fileBufferSizeBytes, writeMetricsToUse);
}
// 将内存中的记录实际写入磁盘
final long recordPointer = sortedRecords.packedRecordPointer.getRecordPointer();
final Object recordPage = taskMemoryManager.getPage(recordPointer);
final long recordOffsetInPage = taskMemoryManager.getOffsetInPage(recordPointer);
int dataRemaining = Platform.getInt(recordPage, recordOffsetInPage);
long recordReadPosition = recordOffsetInPage + 4; // skip over record length
while (dataRemaining > 0) {
final int toTransfer = Math.min(DISK_WRITE_BUFFER_SIZE, dataRemaining);
Platform.copyMemory(
recordPage, recordReadPosition, writeBuffer, Platform.BYTE_ARRAY_OFFSET, toTransfer);
writer.write(writeBuffer, 0, toTransfer);
recordReadPosition += toTransfer;
dataRemaining -= toTransfer;
}
writer.recordWritten();
}
if (writer != null) {
writer.commitAndClose();
// 记录最后一个分区所包含的记录数
if (currentPartition != -1) {
spillInfo.partitionLengths[currentPartition] = writer.fileSegment().length();
// 将当前spill的文件的每个分区对应的fileSegment信息保存到spills中
spills.add(spillInfo);
}
}
// 省略代码
}
}
ShuffleInMemorySorter
ShuffleInMemorySorter实际用来保存输入记录在内存页中的指针，并对输入的记录进行排序，支持基数排序和timsort排序两种排序方法，默认使用基数排序。
insertRecord
使用数组来保存输入记录以及其所在reduce分区id的指针（由PackedRecordPointer封装）。
public void insertRecord(long recordPointer, int partitionId) {
if (!hasSpaceForAnotherRecord()) {
throw new IllegalStateException("There is no space for new record");
}
array.set(pos, PackedRecordPointer.packPointer(recordPointer, partitionId));
pos++;
}
getSortedIterator
对内存中的数据进行排序，并返回iterator。
/**
* 返回排序后的iterator
*/
public ShuffleSorterIterator getSortedIterator() {
int offset = 0;
// 如果使用基数排序
if (useRadixSort) {
// 对数据排序
offset = RadixSort.sort(
array, pos,
PackedRecordPointer.PARTITION_ID_START_BYTE_INDEX,
PackedRecordPointer.PARTITION_ID_END_BYTE_INDEX, false, false);
} else { // 使用timsort排序
MemoryBlock unused = new MemoryBlock(
array.getBaseObject(),
array.getBaseOffset() + pos * 8L,
(array.size() - pos) * 8L);
LongArray buffer = new LongArray(unused);
Sorter<PackedRecordPointer, LongArray> sorter =
new Sorter<>(new ShuffleSortDataFormat(buffer));
// 对数据排序
sorter.sort(array, 0, pos, SORT_COMPARATOR);
}
return new ShuffleSorterIterator(pos, array, offset);
}
SortShuffleWriter
SortShuffleWriter与BypassMergeSortShuffleWriter和UnsafeShuffleWriter最大的不同在于他可以在map阶段对输入的记录按照key hashcode进行聚合以及对分配到每个reduce分区里的记录按照key hashcode进行排序。SortShuffleWriter具体的排序方法委托给 ExternalSorter 这个排序器来完成。
write
/** 将输入记录排序后输出到map文件中 */
override def write(records: Iterator[Product2[K, V]]): Unit = {
// 根据是否需要对map的输出进行排序，来构建externalSorter
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 {
// 如果在map阶段不需要聚合排序，则与bypass和unsafe shuffleWriter一样，只对输入记录按照reduce分区排序，并输出到map文件中
// 实际的聚合排序在reduce阶段进行
new ExternalSorter[K, V, V](
context, aggregator = None, Some(dep.partitioner), ordering = None, dep.serializer)
}
// 将记录保存到内存中，并有sorter负责刷新到磁盘上
sorter.insertAll(records)
// 获取输出文件位置
val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
// 构建临时文件
val tmp = Utils.tempFileWith(output)
// 生成ShuffleBlockId
val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
// 将数据写入临时文件，返回每个partition对应的大小
val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
// 生成索引文件并将临时文件重命名为最终文件
shuffleBlockResolver.writeIndexFileAndCommit(dep.shuffleId, mapId, partitionLengths, tmp)
// 更新map状态
mapStatus = MapStatus(blockManager.shuffleServerId, partitionLengths)
}
ExternalSorter
ExternalSorter实现对输入记录的排序，以及按照key hashcode对记录按照给定的聚合函数进行聚合，对reduce分区内的按照key hashcode进行排序，并最终将结果输出到一个map文件中。
ExternalSorter内部有两种缓存数据结构：PartitionedAppendOnlyMap和PartitionedPairBuffer。如果需要在map阶段进行聚合排序操作，则使用PartitionedAppendOnlyMap，否则则使用PartitionedPairBuffer。在记录输入过程中，会不断向缓存中写入数据，如果缓存达到内存上限，则要将缓存中的内容spill到磁盘中。当调用者（SortShuffleWriter）请求输出文件时，则将spill到磁盘上的所有文件和缓存中剩余的记录进行合并，并以一个map文件的形式输出。
ExternalSorter继承自Spillable类，而Spillable继承自MemoryConsumer，所以ExternalSorter所使用的内存也都来自execution memory，由TaskMemoryManager管理。
insertAll
将输入记录写入sorter内存中，如果内存不足则将内存中记录spill到磁盘上，并清空内存。
def insertAll(records: Iterator[Product2[K, V]]): Unit = {
val shouldCombine = aggregator.isDefined
// 如果需要在map端进行聚合，则使用PartitionedAppendOnlyMap作为缓存，将record根据key放入map中，将value按照给定的聚合函数进行聚合
if (shouldCombine) {
// 获取聚合函数
val mergeValue = aggregator.get.mergeValue
// 获取createCombiner函数
val createCombiner = aggregator.get.createCombiner
var kv: Product2[K, V] = null
// 如果已经存在相同key记录，则使用聚合函数进行合并，否则使用createCombiner进行初始化赋值
val update = (hadValue: Boolean, oldValue: C) => {
if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
}
// 遍历所有输入记录
while (records.hasNext) {
addElementsRead()
kv = records.next()
// 将记录插入map中
map.changeValue((getPartition(kv._1), kv._1), update)
// 判断是否需要进行spill操作
maybeSpillCollection(usingMap = true)
}
} else {
// 如果不需要合并，则直接将record放入PartitionedPairBuffer中
while (records.hasNext) {
addElementsRead()
val kv = records.next()
buffer.insert(getPartition(kv._1), kv._1, kv._2.asInstanceOf[C])
maybeSpillCollection(usingMap = false)
}
}
}
maybeSpillCollection
使用预估当前缓存（map/buffer）已使用大小方法，判断是否超过内存上限，如果超过则将缓存中数据spill到磁盘中，并重置缓存。
/**
* 预估当前map/buffer使用量是否超过设定的上限，如果超过则将map/buffer内容写入磁盘，清空map/buffer
*/
private def maybeSpillCollection(usingMap: Boolean): Unit = {
var estimatedSize = 0L
// 使用estimateSize方法预估当前map/buffer使用量
if (usingMap) {
estimatedSize = map.estimateSize()
// 如果超过内存上限，色调用maybeSpill方法将内存数据刷新到磁盘中，并生成新的map
if (maybeSpill(map, estimatedSize)) {
map = new PartitionedAppendOnlyMap[K, C]
}
} else {
estimatedSize = buffer.estimateSize()
// 如果超过内存上限，色调用maybeSpill方法将内存数据刷新到磁盘中，并生成新的buffer
if (maybeSpill(buffer, estimatedSize)) {
buffer = new PartitionedPairBuffer[K, C]
}
}
if (estimatedSize > _peakMemoryUsedBytes) {
_peakMemoryUsedBytes = estimatedSize
}
}
Spillable.maybeSpill
ExternalSorter实际调用父类Spillable的maybeSpill方法在预估到map/buffer中内存不足时，尝试进行spill操作。在真正进行spill之前，会先尝试向TaskMemoryManager申请内存，如果申请内存后的缓存大小大于预估的使用大小，则不需要进行spill，否则进行实际的spill操作
/**
* 将内存中的集合spill到磁盘中，在真正spill之前尝试获取更多的内存。
*/
protected def maybeSpill(collection: C, currentMemory: Long): Boolean = {
var shouldSpill = false
// 在将内存刷新到磁盘之前先尝试多分配一些内存
if (elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold) {
val amountToRequest = 2 * currentMemory - myMemoryThreshold
// 这里通过MemoryConsumer来申请execution memory
val granted = acquireMemory(amountToRequest)
myMemoryThreshold += granted
// 判断是否真的需要spill
shouldSpill = currentMemory >= myMemoryThreshold
}
// 如果申请的内存还是不够，或者内存中保存的元素数量达到需要强制刷新的上限
// 才真正刷磁盘
shouldSpill = shouldSpill || _elementsRead > numElementsForceSpillThreshold
// Actually spill
if (shouldSpill) {
_spillCount += 1
logSpillage(currentMemory)
spill(collection)
_elementsRead = 0
_memoryBytesSpilled += currentMemory
// 释放execution memory
releaseMemory()
}
shouldSpill
}
Spillable.spill
将内存中数据spill到磁盘中。这个方法只对堆内内存模式的execution memory进行spill操作。实际调用子类（ExternalSorter）的forceSpill来将内存数据刷新到磁盘中。返回释放的内存大小。
override def spill(size: Long, trigger: MemoryConsumer): Long = {
if (trigger != this && taskMemoryManager.getTungstenMemoryMode == MemoryMode.ON_HEAP) {
val isSpilled = forceSpill()
if (!isSpilled) {
0L
} else {
val freeMemory = myMemoryThreshold - initialMemoryThreshold
_memoryBytesSpilled += freeMemory
releaseMemory()
freeMemory
}
} else {
0L
}
}
spillMemoryIteratorToDisk
forceSpill会最终调用spillMemoryIteratorToDisk将内存中排好序/聚合好的数据刷新到磁盘上。map/buffer的iterator在返回前会对其中的元素进行排序。
/**
* 将内存数据输出到磁盘的临时文件
*/
private[this] def spillMemoryIteratorToDisk(inMemoryIterator: WritablePartitionedIterator)
: SpilledFile = {
// 创建临时文件
val (blockId, file) = diskBlockManager.createTempShuffleBlock()
// These variables are reset after each flush
var objectsWritten: Long = 0
var spillMetrics: ShuffleWriteMetrics = null
var writer: DiskBlockObjectWriter = null
// 获取临时文件的writer
def openWriter(): Unit = {
assert (writer == null && spillMetrics == null)
spillMetrics = new ShuffleWriteMetrics
writer = blockManager.getDiskWriter(blockId, file, serInstance, fileBufferSize, spillMetrics)
}
openWriter()
// 记录每次flush磁盘的数据大小
val batchSizes = new ArrayBuffer[Long]
// How many elements we have in each partition
// 记录文件中每个分区记录的个数
val elementsPerPartition = new Array[Long](numPartitions)
// 将缓存中的内容刷新到磁盘上，并更新对应变量
def flush(): Unit = {
val w = writer
writer = null
w.commitAndClose()
_diskBytesSpilled += spillMetrics.bytesWritten
// 记录一次刷新多大的数据到磁盘
batchSizes.append(spillMetrics.bytesWritten)
spillMetrics = null
objectsWritten = 0
}
var success = false
try {
// 遍历内存中的数据
while (inMemoryIterator.hasNext) {
val partitionId = inMemoryIterator.nextPartition()
require(partitionId >= 0 && partitionId < numPartitions,
s"partition Id: ${partitionId} should be in the range [0, ${numPartitions})")
// 将记录写入writer缓存
inMemoryIterator.writeNext(writer)
// 记录每个分区有多少个记录
elementsPerPartition(partitionId) += 1
objectsWritten += 1
// 当writer缓存中的记录达到刷新上限时，一次将serializerBatchSize的记录刷新到磁盘上
if (objectsWritten == serializerBatchSize) {
flush()
openWriter()
}
}
if (objectsWritten > 0) {
flush()
} else if (writer != null) {
val w = writer
writer = null
w.revertPartialWritesAndClose()
}
success = true
} finally {
if (!success) {
// This code path only happens if an exception was thrown above before we set success;
// close our stuff and let the exception be thrown further
if (writer != null) {
writer.revertPartialWritesAndClose()
}
if (file.exists()) {
if (!file.delete()) {
logWarning(s"Error deleting ${file}")
}
}
}
}
SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition)
}
writePartitionedFile
此方法由SortShuffleWriter调用，来获取按照shuffle分区规则重新分区后的map输出文件。将spill到磁盘的所有临时文件和内存中的数据进行合并，并写入一个map输出文件，并返回每个分区在map输出文件中的偏移量数组。实际合并缓存和磁盘文件的操作由partitionedIterator实现。
/**
* 将ExternalSorter上的数据(包括spill到磁盘的数据和缓存中的数据)都写到磁盘中的一个文件中
*/
def writePartitionedFile(
blockId: BlockId,
outputFile: File): Array[Long] = {
val writeMetrics = context.taskMetrics().shuffleWriteMetrics
// 用来保存每个reduce分区在最终的outputFile中的偏移量
val lengths = new Array[Long](numPartitions)
// 如果在向externalSorter添加记录时，没有发生将内存数据刷到磁盘上，则spill=null
if (spills.isEmpty) {
val collection = if (aggregator.isDefined) map else buffer
// 根据排序比较器comparator对缓存中的数据进行排序并返回iterator
val it = collection.destructiveSortedWritablePartitionedIterator(comparator)
// 将partition数据写到一个文件中，通过lengths记录每个partition的大小
while (it.hasNext) {
// 对每个分区都新建一个writer
val writer = blockManager.getDiskWriter(
blockId, outputFile, serInstance, fileBufferSize, writeMetrics)
val partitionId = it.nextPartition()
// 遍历当前分区的内容并写入writer
while (it.hasNext && it.nextPartition() == partitionId) {
it.writeNext(writer)
}
writer.commitAndClose()
val segment = writer.fileSegment()
// 记录分区偏移量
lengths(partitionId) = segment.length
}
} else {
// 如果记录存在在磁盘和内存中，则需要进行合并
// 调用partitionedIterator对内存和磁盘中的数据进行合并，返回（分区id，数据）的iterator
for ((id, elements) <- this.partitionedIterator) {
if (elements.hasNext) {
// 为每一个分区新建writer
val writer = blockManager.getDiskWriter(
blockId, outputFile, serInstance, fileBufferSize, writeMetrics)
for (elem <- elements) {
writer.write(elem._1, elem._2)
}
writer.commitAndClose()
val segment = writer.fileSegment()
// 记录分区偏移量
lengths(id) = segment.length
}
}
}
// 监控指标更新
context.taskMetrics().incMemoryBytesSpilled(memoryBytesSpilled)
context.taskMetrics().incDiskBytesSpilled(diskBytesSpilled)
context.taskMetrics().incPeakExecutionMemory(peakMemoryUsedBytes)
// 返回每个分区在outputFile中的偏移量的数组
lengths
}
partitionedIterator & merge
mergeWithAggregation&mergeSort用来对map结果进行聚合/排序。对已经排好序/聚合好的多个文件取出每个文件中的第一个元素，使用堆排进行排序。当调用者遍历iterator时，对聚合来说，是首先拿到堆顶元素，并找到堆中key与对应元素相同的元素，进行聚合，并返回。同时从对应文件取出下一个元素填充到堆中。对于排序来说，则直接返回堆顶元素，并从元素所属文件中取出下一个元素填充到堆中。
/**
* 将内存和磁盘中的数据进行合并，并按照partition分组，按照请求的聚合方式聚合数据并返回。
*
* 如果数据都在内存中，则直接对内存中的数据进行排序（先按照分区id排序，如果需要对分区中的数据进行排序（即
* ordering.isDefined）则再按照指定的排序规则进行排序），然后将数据按照分区id进行分组，并返回。
* 如果数据部分在内存中，部分在磁盘中，则调用merge函数将磁盘中的数据和内存中的数据合并，并返回。
*/
def partitionedIterator: Iterator[(Int, Iterator[Product2[K, C]])] = {
val usingMap = aggregator.isDefined
val collection: WritablePartitionedPairCollection[K, C] = if (usingMap) map else buffer
// 数据只在内存中
if (spills.isEmpty) {
// 是否要对每个分区中的数据进行排序
if (!ordering.isDefined) {
// 将数据排序后按照分区id分组并返回
groupByPartition(destructiveIterator(collection.partitionedDestructiveSortedIterator(None)))
} else {
// 使用指定的key比较器对数据排序后按照分区id分组并返回
groupByPartition(destructiveIterator(
collection.partitionedDestructiveSortedIterator(Some(keyComparator))))
}
} else {
// 调用merge函数将磁盘和内存数据合并
merge(spills, destructiveIterator(
collection.partitionedDestructiveSortedIterator(comparator)))
}
}
private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
: Iterator[(Int, Iterator[Product2[K, C]])] = {
// 获取每个spill文件的reader
val readers = spills.map(new SpillReader(_))
// 获取内存中的数据（传入merge方法的内存数据是已经排好序的）
val inMemBuffered = inMemory.buffered
// 遍历每个reduce分区
(0 until numPartitions).iterator.map { p =>
// 获取内存中对应分区的数据
val inMemIterator = new IteratorForPartition(p, inMemBuffered)
// 将内存中数据和磁盘中对应分区的数据合并（reader使用NIO进行遍历，每次只读固定大小的数据到内存中）
val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
if (aggregator.isDefined) {
// 如果定义了聚合器，则按照聚合方法进行合并
(p, mergeWithAggregation(
iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
} else if (ordering.isDefined) {
// 否则如果定义了排序方法，则按照排序规则进行合并
(p, mergeSort(iterators, ordering.get))
} else {
// 如果没有定义聚合，排序方法，则直接按照分区id排序分组后返回
(p, iterators.iterator.flatten)
}
}
}
ExternalSorter缓存
PartitionedAppendOnlyMap
PartitionedAppendOnlyMap继承自SizeTrackingAppendOnlyMap，而SizeTrackingAppendOnlyMap继承自AppendOnlyMap。可以将PartitionedAppendOnlyMap看做是带预估map大小功能的AppendOnlyMap。用来向map中追加和键值对，更新（更新操作提供了两个函数update和changeValue，其中update直接用新值覆盖旧值，changeValue则使用传入的updateFunc来更新值）已存在的键的值。底层使用数组来保存键值对，偶数位保存键，奇数位保存值。此map使用二次探测方式处理键的hash值碰撞。
插入更新操作update
首先判断是否需要进行扩容，具体扩容操作由growTable方法完成。然后判断插入的键是否存在，如果存在则更新值，不存在则插入值。
/** Set the value for a key */
def update(key: K, value: V): Unit = {
assert(!destroyed, destructionMessage)
val k = key.asInstanceOf[AnyRef]
if (k.eq(null)) {
if (!haveNullValue) {
// 增加大小并判断是否需要扩容
incrementSize()
}
nullValue = value
haveNullValue = true
return
}
var pos = rehash(key.hashCode) & mask
var i = 1
while (true) {
val curKey = data(2 * pos)
// 插入/更新值
if (curKey.eq(null)) {
data(2 * pos) = k
data(2 * pos + 1) = value.asInstanceOf[AnyRef]
incrementSize() // Since we added a new key
return
} else if (k.eq(curKey) || k.equals(curKey)) {
data(2 * pos + 1) = value.asInstanceOf[AnyRef]
return
} else {
val delta = i
pos = (pos + delta) & mask
i += 1
}
}
}
扩容growTable
// 扩容map
protected def growTable() {
// 将数组容量扩充2倍
val newCapacity = capacity * 2
require(newCapacity <= MAXIMUM_CAPACITY, s"Can't contain more than ${growThreshold} elements")
// 新建数组
val newData = new Array[AnyRef](2 * newCapacity)
val newMask = newCapacity - 1
// 将老数组中的数据rehash后放入新数组，如果发生hash碰撞，则二次探测
var oldPos = 0
while (oldPos < capacity) {
if (!data(2 * oldPos).eq(null)) {
val key = data(2 * oldPos)
val value = data(2 * oldPos + 1)
var newPos = rehash(key.hashCode) & newMask
var i = 1
var keepGoing = true
while (keepGoing) {
val curKey = newData(2 * newPos)
if (curKey.eq(null)) {
newData(2 * newPos) = key
newData(2 * newPos + 1) = value
keepGoing = false
} else {
val delta = i
newPos = (newPos + delta) & newMask
i += 1
}
}
}
oldPos += 1
}
// 使用新数组替换老数组
data = newData
capacity = newCapacity
mask = newMask
growThreshold = (LOAD_FACTOR * newCapacity).toInt
}
遍历map
当ExternalSorter要获取map中的数据时，会调用此方法。这个方法首先会对数组进行压缩，去除为null的项。然后使用timsort对数组进行排序并返回iterator。调用这个函数的缺点是由于是直接对map底层的数组进行操作，会破坏map结构。
/**
* 返回排好序的map，这个方法不会使用额外的内存，缺点就是会改变map原来的key-value顺序
*/
def destructiveSortedIterator(keyComparator: Comparator[K]): Iterator[(K, V)] = {
destroyed = true
var keyIndex, newIndex = 0
// 去掉数组中key=null的项（数组压缩）
while (keyIndex < capacity) {
if (data(2 * keyIndex) != null) {
data(2 * newIndex) = data(2 * keyIndex)
data(2 * newIndex + 1) = data(2 * keyIndex + 1)
newIndex += 1
}
keyIndex += 1
}
assert(curSize == newIndex + (if (haveNullValue) 1 else 0))
// 使用timSort对数组进行排序
new Sorter(new KVArraySortDataFormat[K, AnyRef]).sort(data, 0, newIndex, keyComparator)
// 返回排好序的数组的iterator
new Iterator[(K, V)] {
var i = 0
var nullValueReady = haveNullValue
def hasNext: Boolean = (i < newIndex || nullValueReady)
def next(): (K, V) = {
if (nullValueReady) {
nullValueReady = false
(null.asInstanceOf[K], nullValue)
} else {
val item = (data(2 * i).asInstanceOf[K], data(2 * i + 1).asInstanceOf[V])
i += 1
item
}
}
}
}
PartitionedPairBuffer
继承自WritablePartitionedPairCollection。一个只追加键值对的数据结构，并记录保存在其中数据的预估大小。底层使用数组来保存数据，偶数位保存键，奇数位保存值。
插入操作insert
/** Add an element into the buffer */
def insert(partition: Int, key: K, value: V): Unit = {
if (curSize == capacity) {
// 如果数组大小达到容量上限，则进行扩容
growArray()
}
// 偶数位保存键，奇数位保存值
data(2 * curSize) = (partition, key.asInstanceOf[AnyRef])
data(2 * curSize + 1) = value.asInstanceOf[AnyRef]
curSize += 1
// 在每次插入后，判断是否需要重新采样，预估当前保存的数据大小
afterUpdate()
}
扩容操作 growArray
扩容操作只是简单的将容量扩大2倍，并重置采样结果
/** Double the size of the array because we've reached capacity */
private def growArray(): Unit = {
if (capacity >= MAXIMUM_CAPACITY) {
throw new IllegalStateException(s"Can't insert more than ${MAXIMUM_CAPACITY} elements")
}
val newCapacity =
if (capacity * 2 < 0 || capacity * 2 > MAXIMUM_CAPACITY) { // Overflow
MAXIMUM_CAPACITY
} else {
capacity * 2
}
val newArray = new Array[AnyRef](2 * newCapacity)
System.arraycopy(data, 0, newArray, 0, 2 * capacity)
data = newArray
capacity = newCapacity
resetSamples()
}
遍历元素partitionedDestructiveSortedIterator
此方法是对WritablePartitionedPairCollection中方法的具体实现，保证每次返回给调用者的iterator都是按照keyComparator进行了排序的有序结果。内部使用timsort进行排序。
/** Iterate through the data in a given order. For this class this is not really destructive. */
override def partitionedDestructiveSortedIterator(keyComparator: Option[Comparator[K]])
: Iterator[((Int, K), V)] = {
// buffer只是根据partition id对内存中保存的记录做了排序
val comparator = keyComparator.map(partitionKeyComparator).getOrElse(partitionComparator)
new Sorter(new KVArraySortDataFormat[(Int, K), AnyRef]).sort(data, 0, curSize, comparator)
iterator
}
private def iterator(): Iterator[((Int, K), V)] = new Iterator[((Int, K), V)] {
var pos = 0
override def hasNext: Boolean = pos < curSize
override def next(): ((Int, K), V) = {
if (!hasNext) {
throw new NoSuchElementException
}
val pair = (data(2 * pos).asInstanceOf[(Int, K)], data(2 * pos + 1).asInstanceOf[V])
pos += 1
pair
}
}
SizeTracker
PartitionedAppendOnlyMap和PartitionedPairBuffer使用SizeTracker实现对自身保存数据大小的采样预估。主要包括采样（takeSample）和预估（estimateSize）两个方法。SizeTracker会维护一个计数器numUpdates来保存插入的记录次数。对每次插入的数据大小使用公式
$$ avgPerUpdateSize = (lastestSize - previousSize) / (lastestUpdateNum - previousUpdateNum)
$$
来计算，即平均每次插入的数据大小=当前缓存大小-上次采样时缓存大小/间隔的更新次数。
预估时使用公式：
$$ estimateSize = avgPerUpdateSize * (numUpdates - lastestUpdateNum) + lastestSize
$$
即当前缓存预估的大小=上次采样时的大小+采样后到现在预估增长的大小。
/**
* 重新采样，计算在两次采样之间每次更新平均扩大的bytes大小
*/
private def takeSample(): Unit = {
samples.enqueue(Sample(SizeEstimator.estimate(this), numUpdates))
// Only use the last two samples to extrapolate
if (samples.size > 2) {
samples.dequeue()
}
val bytesDelta = samples.toList.reverse match {
case latest :: previous :: tail =>
(latest.size - previous.size).toDouble / (latest.numUpdates - previous.numUpdates)
// If fewer than 2 samples, assume no change
case _ => 0
}
bytesPerUpdate = math.max(0, bytesDelta)
// 延长下次采样时间
nextSampleNum = math.ceil(numUpdates * SAMPLE_GROWTH_RATE).toLong
}
/**
* 估算当前大小。size = 最后一次采样时大小 + （当前updates-最后一次采样updates）* 平均采样每次更新大小
*/
def estimateSize(): Long = {
assert(samples.nonEmpty)
val extrapolatedDelta = bytesPerUpdate * (numUpdates - samples.last.numUpdates)
(samples.last.size + extrapolatedDelta).toLong
}
ShuffleReader （BlockStoreShuffleReader）
不同于shuffleWriter，ShuffleReader只有一个实现类即BlockStoreShuffleReader。当reduce任务需要获取map任务的输出时首先会向Driver请求得到保存有其需要的map文件的节点的地址信息。然后新建ShuffleBlockFetcherIterator实例去相应的节点获取map文件（对于非本地的节点则通过网络请求分批获取数据，对于本地节点则直接返回map文件在本地的物理存储位置）。在获取到iterator之后，会调用各种封装方法将原始的iterator封装为所需要的iterator（获取map数据是惰性的，这里只是对iterator进行各种包装变换，并没有真正去获取数据）。
如果rdd需要进行聚合，则会在reduce阶段对从map文件获取到的数据按key进行聚合（此处会真正调用ShuffleBlockFetcherIterator从远端和本地获取到所有的数据，并将数据写入缓存中，如果缓存不足则将数据按key排序后spill到磁盘上，返回聚合好的iterator。在真正遍历聚合好的iterator时，会使用堆排序方法首先将堆中key相同的结果进行聚合，然后将spill文件中后续数据放入堆中，并返回聚合好的数据）。
如果rdd需要按key排序，则在reduce阶段会对所有获取到的数据进行排序（此处会真正调用ShuffleBlockFetcherIterator从远端和本地获取到所有的数据，并将数据写入缓存中，如果缓存不足则将数据按partition&key排序后spill到磁盘上，返回排序好的iterator。在真正要遍历排好序的iterator时，会使用堆排序对每个spill文件的每个partition的第一个元素按key进行排序，依次返回堆顶的数据，并向堆中填充对应spill文件的partition中的下一个数据）。
也就是说，如果rdd需要进行聚合/排序，则在reduce阶段，需要首先从集群获取到所需要的全量数据，然后写入本地缓存中（缓存不够刷磁盘），然后才能进行聚合/排序，之后再返回给调用者处理好的iterator。即在需要聚合、排序时，都需要写磁盘（如果存在spill情况），都需要先将map的数据落盘，才能返回给调用者进行操作，在调用者对返回的iterator进行遍历时，才使用堆排对数据进行聚合/排序。如果rdd不需要进行排序/聚合，则直接将初始的iterator包装后返回给调用者，当调用者对返回的iterator遍历时，才真正从集群中获取数据，即边处理边fetch。
read
主要用来获取并返回需要的map数据经过处理后的iterator。首先通过新建ShuffleBlockFecherIterator实例得到可以从集群获取需要的map数据。如果当前rdd需要聚合，则调用聚合方法对iterator中的数据按key进行聚合，并返回iterator。如果当前rdd需要排序，则调用排序方法对iterator中的数据按key进行排序，并返回排好序的iterator。
// 读取当前reduce任务需要的键值对集合
override def read(): Iterator[Product2[K, C]] = {
// 获取数据块的遍历器
val blockFetcherItr = new ShuffleBlockFetcherIterator(
context,
blockManager.shuffleClient,
blockManager,
mapOutputTracker.getMapSizesByExecutorId(handle.shuffleId, startPartition, endPartition),
SparkEnv.get.conf.getSizeAsMb("spark.reducer.maxSizeInFlight", "48m") * 1024 * 1024,
SparkEnv.get.conf.getInt("spark.reducer.maxReqsInFlight", Int.MaxValue))
// 根据配置的压缩方法包装流
val wrappedStreams = blockFetcherItr.map { case (blockId, inputStream) =>
serializerManager.wrapForCompression(blockId, inputStream)
}
val serializerInstance = dep.serializer.newInstance()
// 对每一个流创建一个key/value遍历器
val recordIter = wrappedStreams.flatMap { wrappedStream =>
// Note: the asKeyValueIterator below wraps a key/value iterator inside of a
// NextIterator. The NextIterator makes sure that close() is called on the
// underlying InputStream when all records have been read.
serializerInstance.deserializeStream(wrappedStream).asKeyValueIterator
}
// 将流包装成带指标统计的流
val readMetrics = context.taskMetrics.createTempShuffleReadMetrics()
val metricIter = CompletionIterator[(Any, Any), Iterator[(Any, Any)]](
recordIter.map { record =>
readMetrics.incRecordsRead(1)
record
},
context.taskMetrics().mergeShuffleReadMetrics())
// 将流包装为可中断的流
val interruptibleIter = new InterruptibleIterator[(Any, Any)](context, metricIter)
// 如果定义了聚合函数，则进行聚合
val aggregatedIter: Iterator[Product2[K, C]] = if (dep.aggregator.isDefined) {
// 如果map阶段做了聚合，则使用agregator中的聚合方法对所有的map结果做一次聚合
if (dep.mapSideCombine) {
// We are reading values that are already combined
val combinedKeyValuesIterator = interruptibleIter.asInstanceOf[Iterator[(K, C)]]
dep.aggregator.get.combineCombinersByKey(combinedKeyValuesIterator, context)
} else {
// 如果map阶段没有聚合，则根据key对数据聚合
val keyValuesIterator = interruptibleIter.asInstanceOf[Iterator[(K, Nothing)]]
dep.aggregator.get.combineValuesByKey(keyValuesIterator, context)
}
} else {
// 如果不需要聚合，则直接返回iterator
require(!dep.mapSideCombine, "Map-side combine without Aggregator specified!")
interruptibleIter.asInstanceOf[Iterator[Product2[K, C]]]
}
// 如果定义了排序函数，则进行排序
dep.keyOrdering match {
case Some(keyOrd: Ordering[K]) =>
// 使用ExternalSorter对数据进行排序
val sorter =
new ExternalSorter[K, C, C](context, ordering = Some(keyOrd), serializer = dep.serializer)
// 将聚合好的iterator插入sorter缓存中，如果缓存达到上限则将缓存内容刷新到磁盘上
// 在刷新时，对缓存中的记录按照分区，key进行排序，并最终保存到磁盘spill文件中
sorter.insertAll(aggregatedIter)
context.taskMetrics().incMemoryBytesSpilled(sorter.memoryBytesSpilled)
context.taskMetrics().incDiskBytesSpilled(sorter.diskBytesSpilled)
context.taskMetrics().incPeakExecutionMemory(sorter.peakMemoryUsedBytes)
// 调用sorter的iterator方法，读取记录，其中如果记录在缓存和spill文件中都存在，
// 由于每个spill文件都是有序的，则使用堆排的方法，返回按分区&key排好序的iterator
CompletionIterator[Product2[K, C], Iterator[Product2[K, C]]](sorter.iterator, sorter.stop())
case None =>
aggregatedIter
}
}
ShuffleBlockFetcherIterator
BlockStoreShuffleReader实际使用ShuffleBlockFecherIterator来从集群中获取map任务输出的中间结果。map的中间结果可能存在本地，也可能存在集群的其他节点上。对于本地的数据块，直接从本地的BlockManager获取，对于远端的数据块，通过BlockTransferService从网络获取。
对于远端数据块的获取，限制了其缓存大小。每次只能最多获取指定内存大小的数据，并保存在内存中。当数据真正被遍历后，才能继续从远端获取数据。对于本地的数据块获取，由于数据本来就存在本地，这里只是返回了一个对数据块所在位置的引用，不需要缓存。
initialize
初始化ShuffleBlockFetcherIterator。初始化时，会向远端节点发送请求获取数据，并保存到本地缓存中，等待调用者遍历iterator时使用。
private[this] def initialize(): Unit = {
context.addTaskCompletionListener(_ => cleanup())
// 将本地块和远端块分开，对于远端块返回请求
val remoteRequests = splitLocalRemoteBlocks()
// 将请求随机放入队列中
fetchRequests ++= Utils.randomize(remoteRequests)
assert ((0 == reqsInFlight) == (0 == bytesInFlight),
"expected reqsInFlight = 0 but found reqsInFlight = " + reqsInFlight +
", expected bytesInFlight = 0 but found bytesInFlight = " + bytesInFlight)
// 发送请求
fetchUpToMaxBytes()
val numFetches = remoteRequests.size - fetchRequests.size
logInfo("Started " + numFetches + " remote fetches in" + Utils.getUsedTimeMs(startTime))
// 获取本地数据块
fetchLocalBlocks()
logDebug("Got local blocks in " + Utils.getUsedTimeMs(startTime))
}
fetchUpToMaxBytes & sendRequest
从远端获取数据的方法。fetchUpToMaxBytes会判断当前是否有fetchRequest队列中是否存在想要获取远端数据的请求，并检查当前缓存是否足够。如果都满足，则调用sendRequest方法实际从远端获取数据。
private def fetchUpToMaxBytes(): Unit = {
// Send fetch requests up to maxBytesInFlight
// 如果请求队列不空，且使用内存大小小于maxBytesInFlight，则发送请求
while (fetchRequests.nonEmpty &&
(bytesInFlight == 0 ||
(reqsInFlight + 1 <= maxReqsInFlight &&
bytesInFlight + fetchRequests.front.size <= maxBytesInFlight))) {
sendRequest(fetchRequests.dequeue())
}
}
/**
* 向远端请求数据块，如果成功将数据则会调用BlockFetchingListener，将buf中的数据包装为
* SuccessFetchResult
*
* @param req
*/
private[this] def sendRequest(req: FetchRequest) {
logDebug("Sending request for %d blocks (%s) from %s".format(
req.blocks.size, Utils.bytesToString(req.size), req.address.hostPort))
bytesInFlight += req.size
reqsInFlight += 1
// so we can look up the size of each blockID
val sizeMap = req.blocks.map { case (blockId, size) => (blockId.toString, size) }.toMap
val remainingBlocks = new HashSet[String]() ++= sizeMap.keys
val blockIds = req.blocks.map(_._1.toString)
val address = req.address
// 使用ShuffleClient从远端获取数据。
shuffleClient.fetchBlocks(address.host, address.port, address.executorId, blockIds.toArray,
new BlockFetchingListener {
override def onBlockFetchSuccess(blockId: String, buf: ManagedBuffer): Unit = {
// Only add the buffer to results queue if the iterator is not zombie,
// i.e. cleanup() has not been called yet.
ShuffleBlockFetcherIterator.this.synchronized {
if (!isZombie) {
// Increment the ref count because we need to pass this to a different thread.
// This needs to be released after use.
buf.retain()
remainingBlocks -= blockId
// 将数据放到results缓存中
results.put(new SuccessFetchResult(BlockId(blockId), address, sizeMap(blockId), buf,
remainingBlocks.isEmpty))
logDebug("remainingBlocks: " + remainingBlocks)
}
}
logTrace("Got remote block " + blockId + " after " + Utils.getUsedTimeMs(startTime))
}
override def onBlockFetchFailure(blockId: String, e: Throwable): Unit = {
logError(s"Failed to get block(s) from ${req.address.host}:${req.address.port}", e)
results.put(new FailureFetchResult(BlockId(blockId), address, e))
}
}
)
}
next
当调用者遍历ShuffleBlockFetcherIterator时，会调用next方法。这个方法会返回缓存的数据。如果获取的数据来自远端，则从远端再次获取数据并放入缓存。
/**
* 遍历block，如果当前遍历的block来自远端，则更新相关的值（bytesInFlight，reqsInFlight）
* 并调用fetchUpToMaxBytes来读取请求队列中的排队请求，继续从远端获取数据块。如果是本地
* 数据块，则不更新相关值。
*
* 最终返回blockId和对应的inputStream
*
*/
override def next(): (BlockId, InputStream) = {
numBlocksProcessed += 1
val startFetchWait = System.currentTimeMillis()
currentResult = results.take()
val result = currentResult
val stopFetchWait = System.currentTimeMillis()
shuffleMetrics.incFetchWaitTime(stopFetchWait - startFetchWait)
result match {
case SuccessFetchResult(_, address, size, buf, isNetworkReqDone) =>
if (address != blockManager.blockManagerId) {
shuffleMetrics.incRemoteBytesRead(buf.size)
shuffleMetrics.incRemoteBlocksFetched(1)
}
bytesInFlight -= size
if (isNetworkReqDone) {
reqsInFlight -= 1
logDebug("Number of requests in flight " + reqsInFlight)
}
case _ =>
}
// Send fetch requests up to maxBytesInFlight
fetchUpToMaxBytes()
result match {
case FailureFetchResult(blockId, address, e) =>
throwFetchFailedException(blockId, address, e)
case SuccessFetchResult(blockId, address, _, buf, _) =>
try {
(result.blockId, new BufferReleasingInputStream(buf.createInputStream(), this))
} catch {
case NonFatal(t) =>
throwFetchFailedException(blockId, address, t)
}
}
}