Spark Shuffle源码分析系列之Shuffle介绍&演进过程介绍了Shuffle的演进过程，本文将介绍shuffle过程中使用到的基础类shuffleManager，ShuffleWriter ，ShuffleReader ， ShuffleHandle， ShuffleBlockResolver的实现，为后续进行Map-Reduce实现分析打下基础。
ShuffleManager
概述
ShuffleManager是Spark系统中可插拔的Shuffle系统接口，ShuffleManager会在Driver或Executor的SparkEnv被创建时一并创建，可以通过spark.shuffle.manage配置指定具体的实现类。目前指定唯一实现类org.apache.spark.shuffle.sort.SortShuffleManager，Driver会将Shuffle操作注册到该组件上，Executor可以询问该组件以读取或写入Shuffle过程中的数据，我们先来看下SparkEnv中创建ShuffleManager的代码:
// org.apache.spark.SparkEnv
// sort和tungsten-sort两种ShuffleManager的实现了都是SortShuffleManager
val shortShuffleMgrNames = Map(
"sort" -> classOf[org.apache.spark.shuffle.sort.SortShuffleManager].getName,
"tungsten-sort" -> classOf[org.apache.spark.shuffle.sort.SortShuffleManager].getName)
// 由参数spark.shuffle.manager来指定ShuffleManager，默认是SortShuffleManager
val shuffleMgrName = conf.get("spark.shuffle.manager", "sort")
val shuffleMgrClass = shortShuffleMgrNames.getOrElse(shuffleMgrName.toLowerCase, shuffleMgrName)
// 初始化ShuffleManager，利用反射机制创建
val shuffleManager = instantiateClass[ShuffleManager](shuffleMgrClass)
接口
ShuffleManager中定义了注册、注销Shuffle，获取ShuffleWriter，ShuffleReader，以及文件解析的接口，供ShuffleTask使用。ShuffleMapTask通过getWriter方法获取相应的ShuffleWriter进行Map任务的写入，写入结果生成数据和索引文件，通过ShuffleBlockResolver来进行解析，ShuffleReduceTask通过getReader方法获取相应的ShuffleReader进行Reduce任务的执行，从其中读取数据，进行处理。
private[spark] trait ShuffleManager {
// 注册一个Shuffle过程
def registerShuffle[K, V, C](shuffleId: Int, numMaps: Int, dependency: ShuffleDependency[K, V, C]): ShuffleHandle
// 获取输出数据用的ShuffleWriter；会被Executor的Map任务调用
def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext): ShuffleWriter[K, V]
// 获取输出数据用的ShuffleReader；会被Executor的Map任务调用
def getReader[K, C](handle: ShuffleHandle, startPartition: Int, endPartition: Int, context: TaskContext): ShuffleReader[K, C]
// 取消对指定的Shuffle过程额注册
def unregisterShuffle(shuffleId: Int): Boolean
// 用于获取Shuffle Block数据的ShuffleBlockResolver
def shuffleBlockResolver: ShuffleBlockResolver
// shutdown
def stop(): Unit
}
ShuffleHandle
ShuffleHandle是不透明的Shuffle句柄，包含了关于Shuffle的一些信息，主要给ShuffleManager使用。本质上来说，它是一个标志位，除了包含一些用于shuffle的一些属性之外，没有其他额外的方法，使用case class来实现。
abstract class ShuffleHandle(val shuffleId: Int) extends Serializable {}
private[spark] class BaseShuffleHandle[K, V, C](shuffleId: Int, val numMaps: Int, val dependency: ShuffleDependency[K, V, C])
extends ShuffleHandle(shuffleId)
private[spark] class SerializedShuffleHandle[K, V](shuffleId: Int, numMaps: Int, dependency: ShuffleDependency[K, V, V])
extends BaseShuffleHandle(shuffleId, numMaps, dependency) {}
// 用于确定何时选择绕开合并和排序的Shuffle路径
private[spark] class BypassMergeSortShuffleHandle[K, V](shuffleId: Int, numMaps: Int, dependency: ShuffleDependency[K, V, V])
extends BaseShuffleHandle(shuffleId, numMaps, dependency) {}
ShuffleWriter
// 定义了将map任务的中间结果输出到磁盘上的功能规范，包括将数据写入磁盘和关闭ShuffleWriter
private[spark] abstract class ShuffleWriter[K, V] {
// 用于将map任务的结果写到磁盘
@throws[IOException]
def write(records: Iterator[Product2[K, V]]): Unit
// 用于关闭ShuffleWriter
def stop(success: Boolean): Option[MapStatus]
}
ShuffleWriter是Spark提供的ShuffleMapTask写入数据的主要类，它和shuffleHandle是一一对应的，会根据上面选择的shuffleHandle的具体实现来选择相应的writer，有三种writer，分别是UnsafeShuffleWriter，BypassMergeSortShuffleWriter和SortShuffleWriter。ShuffleManager通过getWriter方法获取合适的ShuffleWriter，然后通过write方法写入数据到存储系统中。
ShuffleReader
private[spark] trait ShuffleReader[K, C] {
// 读取数据
def read(): Iterator[Product2[K, C]]
}
ShuffleReader是Shuffle read任务从上游ShuffleMapTask的结果MapStatus获取文件信息，读取数据产生迭代器，是后续Task使用的源数据的生产者，目前唯一实现是BlockStoreShuffleReader，实现了read方法，后面讲到Shuffle Reduce任务读取数据时候会具体分析。
ShuffleBlockResolver
概述
特质ShuffleBlockResolver定义了对Shuffle Block进行解析的规范，IndexShuffleBlockResolver是目前的它的唯一实现类， 主要用于shuffle blocks从逻辑block到物理文件之间的映射关系。实现了获取Shuffle数据文件、获取Shuffle索引文件、删除指定的Shuffle数据文件和索引文件、生成Shuffle索引文件、获取Shuffle块的数据等。
数据文件
- 命名格式: ShuffleMapTask生成的文件是按照下面的格式进行定义的，所以知道了shuffleId，mapId以及reduceId就能通过BlockManager获取该文件。
case class ShuffleDataBlockId(shuffleId: Int, mapId: Int, reduceId: Int) extends BlockId {
override def name: String = "shuffle_" + shuffleId + "_" + mapId + "_" + reduceId + ".data"
}
- 获取数据文件
def getDataFile(shuffleId: Int, mapId: Int): File = {
blockManager.diskBlockManager.getFile(ShuffleDataBlockId(shuffleId, mapId, NOOP_REDUCE_ID))
}
索引文件
- 命名格式: ShuffleMapTask生成的文件是按照下面的格式进行定义的，所以知道了shuffleId，mapId以及reduceId就能通过BlockManager获取该文件。
case class ShuffleIndexBlockId(shuffleId: Int, mapId: Int, reduceId: Int) extends BlockId {
override def name: String = "shuffle_" + shuffleId + "_" + mapId + "_" + reduceId + ".index"
}
- 获取数据文件
def getDataFile(shuffleId: Int, mapId: Int): File = {
blockManager.diskBlockManager.getFile(ShuffleDataBlockId(shuffleId, mapId, NOOP_REDUCE_ID))
}
检查Index文件和数据文件是否合法
索引文件存储的是每个Reduce任务起始位置在数据文件中的偏移位置，索引值为Long型，占8个字节，索引文件头的8个字节为标记值，即Long型整数0，所以索引文件的其长度需要是【数据块数量 + 1】 * 8，紧邻的两个索引值的差为对应Reduce的数据的大小，检查时候我们需要保证两方面一是索引文件数据是【数据块数量 + 1】 * 8，另一方面要保证读取索引得到的数据文件的大小等于数据文件读取到的大小。
private def checkIndexAndDataFile(index: File, data: File, blocks: Int): Array[Long] = {
// 检查索引文件长度，其长度需要是【数据块数量 + 1】 * 8， 索引值为Long型，占8个字节，索引文件头的8个字节为标记值，即Long型整数0，不算索引值
if (index.length() != (blocks + 1) * 8) {
return null
}
// 创建数据块数量大小的数组
val lengths = new Array[Long](blocks)
// 得到索引文件的输入流
val in = try {
new DataInputStream(new NioBufferedFileInputStream(index))
} catch {
case e: IOException =>
return null
}
try {
// 读取第一个Long型整数
var offset = in.readLong()
// 第一个Long型整数标记值必须为0，如果不为0，则直接返回null
if (offset != 0L) {
return null
}
var i = 0
while (i < blocks) {
// 读取Long型偏移量
val off = in.readLong()
// 记录对应的数据块长度
lengths(i) = off - offset
// offset更新
offset = off
i += 1
}
} catch {
case e: IOException =>
return null
} finally {
in.close()
}
// 数据文件的长度，等于lengths数组所有元素之和，表示校验成功
if (data.length() == lengths.sum) {
// 返回lengths数组
lengths
} else {
null
}
}
写入
Spark Shuffle中间文件是由数据文件和索引文件共同组成的，在ShuffleWriter写入到数据文件后，会记录每个Reduce的偏移量，writeIndexFileAndCommit 是把偏移量写入到索引文件中，由于同一个ShuffleMapTask可能会有多个尝试，所以考虑并发问题，需要进行同步，需要经过以下步骤：
- 先根据偏移量情况写偏移量到索引临时文件中；
- 获取索引文件和数据文件，然后加锁，检查indexFile文件、dataFile文件及数据块的长度是否相同，这个操作是为了检查可能已经存在的索引文件与数据文件是否匹配，如果检查发现时匹配的，则说明有其他的TaskAttempt已经完成了该Map任务的写出，那么此时产生的临时索引文件就无用了。
- 如果匹配成功说明其他任务已经写入了，删除临时的索引文件和数据文件；
- 如果不匹配，将临时文件重命名为正式文件
def writeIndexFileAndCommit(shuffleId: Int, mapId: Int, lengths: Array[Long], dataTmp: File): Unit = {
// 获取指定Shuffle中指定map任务输出的索引文件
val indexFile = getIndexFile(shuffleId, mapId)
// 根据索引文件获取临时索引文件的路径
val indexTmp = Utils.tempFileWith(indexFile)
try {
// 构建临时文件的输出流
val out = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(indexTmp)))
Utils.tryWithSafeFinally { // 写入临时索引文件的第一个标记值，即Long型整数0
var offset = 0L
out.writeLong(offset)
// 遍历每个数据块的长度，并作为偏移量写入临时索引文件
for (length <- lengths) { // 在原有offset上加上数据块的长度
offset += length
// 写入临时文件
out.writeLong(offset)
}
} {
out.close()
}
// 获取指定Shuffle中指定map任务输出的数据文件
val dataFile = getDataFile(shuffleId, mapId)
synchronized { // 主要该步骤是加锁的
// 检查indexFile文件、dataFile文件及数据块的长度是否相同，注意，这里传入的不是临时索引文件indexTmp，
// 这个操作是为了检查可能已经存在的索引文件与数据文件是否匹配，如果检查发现时匹配的，则说明有其他的TaskAttempt已经完成了该Map任务的写出，
// 那么此时产生的临时索引文件就无用了。
// 由于当前代码块是使用synchronized同步的，因此不用担心并发问题，同一个时间点只会有一个TaskAttempt成功写出数据。
val existingLengths = checkIndexAndDataFile(indexFile, dataFile, lengths.length)
if (existingLengths != null) { // 如果匹配
// 将checkIndexAndDataFile()方法记录的数据文件的每个分区的长度组成的数组复制到lengths数组中
System.arraycopy(existingLengths, 0, lengths, 0, lengths.length)
if (dataTmp != null && dataTmp.exists()) {
// 将临时索引文件删除
dataTmp.delete()
}
// 将临时索引文件删除
indexTmp.delete()
} else { // 如果不匹配，说明当前还没有其他的TaskAttempt进行索引文件的写入，本次操作产生的临时索引文件可以用
if (indexFile.exists()) {
indexFile.delete()
}
if (dataFile.exists()) {
dataFile.delete()
}
// 临时的索引文件和数据文件作为正式的索引文件和数据文件, 将indexTmp重命名为indexFile
if (!indexTmp.renameTo(indexFile)) {
throw new IOException("fail to rename file " + indexTmp + " to " + indexFile)
}
// 将dataTmp重命名为dataFile
if (dataTmp != null && dataTmp.exists() && !dataTmp.renameTo(dataFile)) {
throw new IOException("fail to rename file " + dataTmp + " to " + dataFile)
}
}
}
} finally {
// 如果临时索引文件还存在，一定要将其删除
if (indexTmp.exists() && !indexTmp.delete()) {
logError(s"Failed to delete temporary index file at ${indexTmp.getAbsolutePath}")
}
}
}
数据读取
getBlockData用于获取某个MapId文件中对应ReduceId的数据，由于索引文件中存储了各个ReduceId对应的偏移量，而且每个偏移量占8字节，所以很明显可以从索引文件中获取到ReduceId的起始位置<reduceId * 8>，然后读取下一个reduceId的偏移量，两者相减即为数据长度，所以我们可以根据偏移量以及reduce需要的数据的长度，可以从数据文件中截取相应的数据，构建FileSegmentManagedBuffer返回给使用方使用。
override def getBlockData(blockId: ShuffleBlockId): ManagedBuffer = {
// 获取指定map任务输出的索引文件
val indexFile = getIndexFile(blockId.shuffleId, blockId.mapId)
// 读取索引文件的输入流
val in = new DataInputStream(new FileInputStream(indexFile))
try {
// 跳过与当前Reduce任务无关的字节，在索引文件中是按照Reduce任务ID的顺序记录每个Reduce对应的数据块的索引数据的。
ByteStreams.skipFully(in, blockId.reduceId * 8)
// 读取偏移量
val offset = in.readLong()
// 读取下一个偏移量
val nextOffset = in.readLong()
// 构造并返回FileSegmentManagedBuffer
new FileSegmentManagedBuffer(
transportConf,
getDataFile(blockId.shuffleId, blockId.mapId),
// 读取的起始偏移量为offset
offset,
// 读取长度为nextOffset - offset
nextOffset - offset)
} finally {
in.close()
}
}
SortShuffleManager
最后我们来看下ShuffleManager的唯一实现类SortShuffleManager，他实现了获取writer和reader的具体接口。
-
属性，
numMapsForShuffle记录了shuffle的唯一标识和Map数量的关系，在获取writer时候添加，在注销shuffle时候删除。// Shuffle的ID与为此Shuffle生成输出的map任务的数量之间的映射关系。 private[this] val numMapsForShuffle = new ConcurrentHashMap[Int, Int]() -
注册Shuffle，根据Shuffle的数据情况，选择合适的
ShuffleHandle- 当map个数比较少，小于
spark.shuffle.sort.bypassMergeThreshold而且不需要做map-side的聚合操作时候，可以使用BypassMergeSortShuffleHandle，这种类似于早期的HashShuffle，每个reduce一个文件，直接写入到文件中，最后合并为一个文件，不需要序列化和反序列化开销，加速程序运行，但是会同时打开多个文件； - 当shuffle可以使用序列化shuffle时候，使用
SerializedShuffleHandle，数据写入时候会进行序列化，缓存，直接在序列化的二进制数据上排序而不是在java 对象上，这样可以减少内存的消耗和GC的开销，另外将键值和指针结合到一起进行排序，可以更好的利用缓存； - 都不满足时候选用默认的
BaseShuffleHandle，通过SortShuffleWriter来写入数据。
override def registerShuffle[K, V, C](shuffleId: Int, numMaps: Int, dependency: ShuffleDependency[K, V, C]): ShuffleHandle = { // 判断使用哪一种ShuffleHandle if (SortShuffleWriter.shouldBypassMergeSort(SparkEnv.get.conf, dependency)) { // 需要绕开合并及排序，则创建BypassMergeSortShuffleHandle // 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)) { // 如果可以使用序列化的Shuffle，则创建SerializedShuffleHandle new SerializedShuffleHandle[K, V]( shuffleId, numMaps, dependency.asInstanceOf[ShuffleDependency[K, V, V]]) } else { // Otherwise, buffer map outputs in a deserialized form: // 其他情况，将创建BaseShuffleHandle new BaseShuffleHandle(shuffleId, numMaps, dependency) } } - 当map个数比较少，小于
-
获取writer，首先会将当前shuffleId先记录到shuffleManager中，根据
shuffleHandle选择合适的writer，writer负责ShuffleMapTask具体数据写入的过程。// 用于根据ShuffleHandle获取ShuffleWriter。 override def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext): ShuffleWriter[K, V] = { // 将指定的shuffleId和Shuffle对应的map任务数注册到numMapsForShuffle字典中 numMapsForShuffle.putIfAbsent(handle.shuffleId, handle.asInstanceOf[BaseShuffleHandle[_, _, _]].numMaps) val env = SparkEnv.get // 根据ShuffleHandle的具体类型，创建不同的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) } } -
获取reader，当Shuffle Reduce任务执行时候，要获取数据，这个是通过
BlockStoreShuffleReader来实现read方法获取迭代器数据的。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) } -
注销shuffle，注销后需要删除shuffleId的记录，并清理产生的数据文件和索引文件。
override def unregisterShuffle(shuffleId: Int): Boolean = { Option(numMapsForShuffle.remove(shuffleId)).foreach { numMaps => (0 until numMaps).foreach { mapId => // 删除map任务产生的数据文件和索引文件的 shuffleBlockResolver.removeDataByMap(shuffleId, mapId) } } true } -
是否可以序列化的判断，使用序列化的Shuffle，会创建SerializedShuffleHandle，这种模式下需要先对数据进行序列化，缓存，直接在序列化的二进制数据上排序而不是在java 对象上，这样可以减少内存的消耗和GC的开销，另外将键值和指针结合到一起进行排序，可以更好的利用缓存<缓存感知计算技术>，如果可以序列化，需要满足以下三个条件:
- 序列化器支持对已经序列化的对象重定位<例如
KryoSerializer>，重新排序序列化流输出中的序列化对象的字节等同于在序列化它们之前重新排序这些元素，这是为了直接对序列化的数据进行排序； - shuffle dependency没有指定aggregation；
- reduce端的分区数目小于等于16777216(排序过程中使用的是记录指针，其中分区ID占24位，则
MAXIMUM_PARTITION_ID = (1 << 24) - 1; //16777215，又因为ID是从0开始的，所以分区数目不能大于16777216)。
// Shuffle是否可以序列化 def canUseSerializedShuffle(dependency: ShuffleDependency[_, _, _]): Boolean = { // Shuffle ID和分区数量 val shufId = dependency.shuffleId val numPartitions = dependency.partitioner.numPartitions if (!dependency.serializer.supportsRelocationOfSerializedObjects) { // ShuffleDependency的序列化器无法对流中输出的序列化后的对象的字节进行排序，则返回false。 false } else if (dependency.aggregator.isDefined) { // ShuffleDependency指定了聚合器，说明存在聚合操作，则返回false。 false } else if (numPartitions > MAX_SHUFFLE_OUTPUT_PARTITIONS_FOR_SERIALIZED_MODE) { // Shuffle过程产生的分区数大于MAX_SHUFFLE_OUTPUT_PARTITIONS_FOR_SERIALIZED_MODE（16777216） false } else { // 三个条件都满足，返回true true } } - 序列化器支持对已经序列化的对象重定位<例如
574
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
