前面我们分析了Spark中具体的Task的提交和运行过程，从本文开始我们开始进入Shuffle的世界，Shuffle对于分布式计算来说是至关重要的部分，它直接影响了分布式系统的性能，所以我将尽可能进行详细的分析。
我们首先来看Shuffle中的Write部分：
override def runTask(context: TaskContext): MapStatus = {
// Deserialize the RDD using the broadcast variable.
val deserializeStartTime = System.currentTimeMillis()
val ser = SparkEnv.get.closureSerializer.newInstance()
val (rdd, dep) = ser.deserialize[(RDD[_], ShuffleDependency[_, _, _])](
ByteBuffer.wrap(taskBinary.value), Thread.currentThread.getContextClassLoader)
_executorDeserializeTime = System.currentTimeMillis() - deserializeStartTime
metrics = Some(context.taskMetrics)
var writer: ShuffleWriter[Any, Any] = null
try {
val manager = SparkEnv.get.shuffleManager
writer = manager.getWriter[Any, Any](dep.shuffleHandle, partitionId, context)
writer.write(rdd.iterator(partition, context).asInstanceOf[Iterator[_ <: Product2[Any, Any]]])
writer.stop(success = true).get
} catch {
case e: Exception =>
try {
if (writer != null) {
writer.stop(success = false)
}
} catch {
case e: Exception =>
log.debug("Could not stop writer", e)
}
throw e
}
}
首先根据SparkEnv获得ShuffleManager，ShuffleManager是为Spark shuffle系统而抽象的可插拔的接口，它被创建在Driver和Executor上，具体是在SparkEnv实例化的时候进行配置的，源码如下：
val shortShuffleMgrNames = Map(
"hash" -> "org.apache.spark.shuffle.hash.HashShuffleManager",
"sort" -> "org.apache.spark.shuffle.sort.SortShuffleManager",
"tungsten-sort" -> "org.apache.spark.shuffle.sort.SortShuffleManager")
val shuffleMgrName = conf.get("spark.shuffle.manager", "sort")
val shuffleMgrClass = shortShuffleMgrNames.getOrElse(shuffleMgrName.toLowerCase, shuffleMgrName)
val shuffleManager = instantiateClass[ShuffleManager](shuffleMgrClass)
可以看到是由"spark.shuffle.manager"配置项来决定具体使用哪种实现方式，默认情况下使用的是sort的方式(本文参考的是Spark 1.6.3版本的源码)。Driver通过ShuffleManager来注册shuffles，Executors可以通过它来读写数据。
获得到ShuffleManager后，就根据它来获得ShuffleWriter(根据具体ShuffleManager的getWriter方法获得)，顾名思义就是用来写数据，而接下来的工作就是调用具体的ShuffleWriter的write方法来进行写数据的工作。
先来看一下getWriter方法，这里的第一个参数dep.shuffleHandle是ShuffleDependency的一个成员变量：
val shuffleHandle: ShuffleHandle = _rdd.context.env.shuffleManager.registerShuffle(
shuffleId, _rdd.partitions.size, this)
这里的registerShuffle方法用来向ShuffleManager注册一个shuffle并且获得一个用来传递任务的句柄，会根据不同ShuffleManager有不同的实现，HashShuffleManager返回的是BaseShuffleHandle，而SortShuffleManager又会根据不同的情况返回BypassMergeSortShuffleHandle、SerializedShuffleHandle或者BaseShuffleHandle。
HashShuffleManager：
override def registerShuffle[K, V, C](
shuffleId: Int,
numMaps: Int,
dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
new BaseShuffleHandle(shuffleId, numMaps, dependency)
}
SortShuffleManager：
override def registerShuffle[K, V, C](
shuffleId: Int,
numMaps: Int,
dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
if (SortShuffleWriter.shouldBypassMergeSort(SparkEnv.get.conf, dependency)) {
// 这里注释说的很清楚，根据spark.shuffle.sort.bypassMergeThreshold的值（默认是200）判断是否需要进行Map端的聚合操作
// 如果partitions的个数小于200就不进行该操作
// 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]])
// 这里是判断是否使用tungsten的方式
} 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)
}
}
使用一张图来总结一下上面的过程：
然后我们来看ShuffleManager的getWriter方法：
HashShuffleManager：
/** Get a writer for a given partition. Called on executors by map tasks. */
override def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext)
: ShuffleWriter[K, V] = {
new HashShuffleWriter(
shuffleBlockResolver, handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, context)
}
SortShuffleManager：
/** Get a writer for a given partition. Called on executors by map tasks. */
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 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)
}
}
从上述源码可以看出，getWriter方法内部实际上是根据传进来的ShuffleHandle的具体类型来判断使用哪种ShuffleWriter的，然后最终执行ShuffleWriter的write方法，下面我们就分为HashShuffleManager和SortShuffleManager两种类型来进行分析。
1、HashShuffleManager
从上面的源码中可以看到HashShuffleManager最终实例化的是HashShuffleWriter，实例化的时候有一行比较重要的代码：
private val shuffle = shuffleBlockResolver.forMapTask(dep.shuffleId, mapId, numOutputSplits, ser,
writeMetrics)
我们来看forMapTask这个方法：
def forMapTask(shuffleId: Int, mapId: Int, numReducers: Int, serializer: Serializer,
writeMetrics: ShuffleWriteMetrics): ShuffleWriterGroup = {
// 为每个ShuffleMapTask实例化了一个ShuffleWriterGroup
new ShuffleWriterGroup {
// 实例化ShuffleState并保存shuffleId和ShuffleState的对应关系
shuffleStates.putIfAbsent(shuffleId, new ShuffleState(numReducers))
// 根据shuffleId获得对应的ShuffleState
private val shuffleState = shuffleStates(shuffleId)
val openStartTime = System.nanoTime
val serializerInstance = serializer.newInstance()
// 获得该ShuffleWriterGroup的writers
val writers: Array[DiskBlockObjectWriter] = {
Array.tabulate[DiskBlockObjectWriter](numReducers) { bucketId =>
// 生成ShuffleBlockId，是一个case class，我们可以通过name方法看到其具体的组成：
// override def name: String = "shuffle_" + shuffleId + "_" + mapId + "_" + reduceId
val blockId = ShuffleBlockId(shuffleId, mapId, bucketId)
// 通过DiskBlockManager的getFile方法获得File
val blockFile = blockManager.diskBlockManager.getFile(blockId)
// 临时目录
val tmp = Utils.tempFileWith(blockFile)
// 使用BlockManager的getDiskWriter方法获得DiskBlockObjectWriter
// 注意这里的bufferSize默认情况下是32kb，可以通过spark.shuffle.file.buffer进行配置
// private val bufferSize = conf.getSizeAsKb("spark.shuffle.file.buffer", "32k").toInt * 1024
blockManager.getDiskWriter(blockId, tmp, serializerInstance, bufferSize, writeMetrics)
}
}
// Creating the file to write to and creating a disk writer both involve interacting with
// the disk, so should be included in the shuffle write time.
writeMetrics.incShuffleWriteTime(System.nanoTime - openStartTime)
override def releaseWriters(success: Boolean) {
shuffleState.completedMapTasks.add(mapId)
}
}
}
为每个ShuffleMapTask实例化一个ShuffleWriterGroup，其中包含了一组writers，每个writer对应一个reducer。
然后我们进入到HashShuffleWriter的write方法：
/** Write a bunch of records to this task's output */
override def write(records: Iterator[Product2[K, V]]): Unit = {
val iter = if (dep.aggregator.isDefined) {
// 判断是否进行map端的combine操作
if (dep.mapSideCombine) {
dep.aggregator.get.combineValuesByKey(records, context)
} else {
records
}
} else {
require(!dep.mapSideCombine, "Map-side combine without Aggregator specified!")
records
}
for (elem <- iter) {
val bucketId = dep.partitioner.getPartition(elem._1)
shuffle.writers(bucketId).write(elem._1, elem._2)
}
}
如果没有进行Map端的combine操作，根据key获得bucketId，实际上是进行取模运算：
def nonNegativeMod(x: Int, mod: Int): Int = {
val rawMod = x % mod
rawMod + (if (rawMod < 0) mod else 0)
}
下面就是根据bucketId获得ShuffleWriterGroup中对应的writer，然后执行其write方法，将key和value写入到对应的block file中：
/**
* Writes a key-value pair.
*/
def write(key: Any, value: Any) {
if (!initialized) {
open()
}
objOut.writeKey(key)
objOut.writeValue(value)
recordWritten()
}
写入完成后，我们回到ShuffleMapTask的runTask方法中，接下来执行的是：
writer.stop(success = true).get
即HashShuffleWriter的stop方法，该方法返回的是Option[MapStatus]，最主要的一句代码为：
Some(commitWritesAndBuildStatus())
进入到commitWritesAndBuildStatus方法：
private def commitWritesAndBuildStatus(): MapStatus = {
// Commit the writes. Get the size of each bucket block (total block size).
val sizes: Array[Long] = shuffle.writers.map { writer: DiskBlockObjectWriter =>
// 调用DiskBlockObjectWriter的commitAndClose方法
writer.commitAndClose()
// 获得每个bucket block的大小
writer.fileSegment().length
}
// 重命名所有的shuffle文件，每个executor只有一个ShuffleBlockResolver，所以使用了synchronized关键字
// rename all shuffle files to final paths
// Note: there is only one ShuffleBlockResolver in executor
shuffleBlockResolver.synchronized {
shuffle.writers.zipWithIndex.foreach { case (writer, i) =>
val output = blockManager.diskBlockManager.getFile(writer.blockId)
if (sizes(i) > 0) {
if (output.exists()) {
// Use length of existing file and delete our own temporary one
sizes(i) = output.length()
writer.file.delete()
} else {
// Commit by renaming our temporary file to something the fetcher expects
if (!writer.file.renameTo(output)) {
throw new IOException(s"fail to rename ${writer.file} to $output")
}
}
} else {
if (output.exists()) {
output.delete()
}
}
}
}
MapStatus(blockManager.shuffleServerId, sizes)
}
这里我们需要注意最终返回的是封装的MapStatus，它记录了产生的磁盘文件的位置，然后Executor中的MapOutputTrackerWorker将MapStatus信息发送给Driver中的MapOutputTrackerMaster，后面Shuffle Read的之后就会从Driver的MapOutputTrackerMaster获取MapStatus的信息，也就是获取对应的上一个ShuffleMapTask的计算结果的输出的文件位置信息。
Map端combine的情况
再来补充一下Map端combine的情况：
if (dep.mapSideCombine) {
dep.aggregator.get.combineValuesByKey(records, context)
} else {
records
}
进入到Aggregator的combineValuesByKey方法：
def combineValuesByKey(
iter: Iterator[_ <: Product2[K, V]],
context: TaskContext): Iterator[(K, C)] = {
// 首先实例化ExternalAppendOnlyMap
val combiners = new ExternalAppendOnlyMap[K, V, C](createCombiner, mergeValue, mergeCombiners)
// 执行ExternalAppendOnlyMap的insertAll方法
combiners.insertAll(iter)
updateMetrics(context, combiners)
combiners.iterator
}
首先实例化了ExternalAppendOnlyMap，然后执行ExternalAppendOnlyMap的insertAll方法：
def insertAll(entries: Iterator[Product2[K, V]]): Unit = {
if (currentMap == null) {
throw new IllegalStateException(
"Cannot insert new elements into a map after calling iterator")
}
// An update function for the map that we reuse across entries to avoid allocating
// a new closure each time
var curEntry: Product2[K, V] = null
val update: (Boolean, C) => C = (hadVal, oldVal) => {
if (hadVal) mergeValue(oldVal, curEntry._2) else createCombiner(curEntry._2)
}
while (entries.hasNext) {
curEntry = entries.next()
val estimatedSize = currentMap.estimateSize()
if (estimatedSize > _peakMemoryUsedBytes) {
_peakMemoryUsedBytes = estimatedSize
}
if (maybeSpill(currentMap, estimatedSize)) {
currentMap = new SizeTrackingAppendOnlyMap[K, C]
}
currentMap.changeValue(curEntry._1, update)
addElementsRead()
}
}
具体的实现方式就不再解释了，简单的说就是将key相同的value进行合并，如果某个key有对应的值就执行merge(也可以理解为更新)操作，如果没有对应的值就新建一个combiner，需要注意的是如果内存不够的话就会将数据spill到磁盘。
HashShuffle方式的Shuffle Write部分至此结束，使用一张图概括一下：
接下来看一下SortShuffle方式的具体流程。
2、SortShuffleManager
为了解决Hash Shuffle产生小文件过多的问题，产生了Sort Shuffle，解下来我们就一起看一下Sort Shuffle的Write部分。
上文中我们已经提到，SortShuffleManager中的getWriter会根据不同的ShuffleHandle产生相应的ShuffleWriter：
- SerializedShuffleHandle 对应 UnsafeShuffleWriter
- BypassMergeSortShuffleHandle 对应 BypassMergeSortShuffleWriter
- BaseShuffleHandle 对应 SortShuffleWriter
下面我们分别进行分析：
2.1、BaseShuffleHandle & SortShuffleWriter
首先来看一下SortShuffleWriter，直接来看它的write方法：
/** Write a bunch of records to this task's output */
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)
}
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).
val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
val tmp = Utils.tempFileWith(output)
try {
val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
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}")
}
}
}
我们看到，内部有一个非常重要的部分，即ExternalSorter，而关于ExternalSorter的使用，源码中的注释说的很清楚了，这里就不做翻译了：
/**
* Users interact with this class in the following way:
*
* 1. Instantiate an ExternalSorter.
*
* 2. Call insertAll() with a set of records.
*
* 3. Request an iterator() back to traverse sorted/aggregated records.
* - or -
* Invoke writePartitionedFile() to create a file containing sorted/aggregated outputs
* that can be used in Spark's sort shuffle.
*/
我们就根据这三个步骤进行说明：
2.1.1 第一步
首先就是实例化ExternalSorter，这里有一个判断，如果要进行map端的combine操作的话就需要指定Aggregator和Ordering，否则这两个参数为None。我们熟悉的reduceByKey就进行了Map端的combine操作，如下图所示：
2.1.2 第二步（这一步非常重要）
通过判断是否进行Map端combine操作而实例化出不同的ExternalSorter后，就会调用insertAll方法，将输入的记录写入到内存中，如果内存不足就spill到磁盘中，具体的实现我们来看insertAll方法：
def insertAll(records: Iterator[Product2[K, V]]): Unit = {
// TODO: stop combining if we find that the reduction factor isn't high
// 首先判断是否需要进行Map端的combine操作
val shouldCombine = aggregator.isDefined
if (shouldCombine) {
// 如果需要进行map端的combine操作，使用PartitionedAppendOnlyMap作为缓存
// 将record根据key对value按照获得的聚合函数进行聚合操作(combine)
// Combine values in-memory first using our AppendOnlyMap
// 获得聚合函数，例如我们使用reduceByKey时编写的函数
val mergeValue = aggregator.get.mergeValue
// 获取createCombiner函数
val createCombiner = aggregator.get.createCombiner
var kv: Product2[K, V] = null
// 定义update函数，主要的逻辑是：如果某个key已经存在记录(record)就使用上面获取
// 的聚合函数进行聚合操作，如果还不存在记录就使用createCombiner方法进行初始化操作
val update = (hadValue: Boolean, oldValue: C) => {
if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
}
// 循环遍历所有的records(记录)
while (records.hasNext) {
// 记录spill的频率，每当read一条record的时候都会记录一次
addElementsRead()
// 使用kv储存当前读的record
kv = records.next()
// 这里的map和下面else中的buffer都是用来缓存的数据结构
// 如果进行Map端的聚合操作，使用的就是PartitionedAppendOnlyMap[K, C]
// 如果不进行Map端的聚合操作，使用的是PartitionedPairBuffer[K, C]
// 调用上面定义的update函数将记录插入到map中
map.changeValue((getPartition(kv._1), kv._1), update)
// 判断是否要进行spill操作
maybeSpillCollection(usingMap = true)
}
} else {
// 如果不需要进行Map端的聚合操作，就直接将记录放到buffer(PartitionedPairBuffer)中
// Stick values into our buffer
while (records.hasNext) {
addElementsRead()
val kv = records.next()
buffer.insert(getPartition(kv._1), kv._1, kv._2.asInstanceOf[C])
maybeSpillCollection(usingMap = false)
}
}
}
具体的流程用注释的方式写在了上面的源码中，这里我们先来看一下PartitionedAppendOnlyMap和PartitionedPairBuffer分别是如何工作的：
PartitionedAppendOnlyMap：
首先来看PartitionedAppendOnlyMap的changeValue实现，实际上，PartitionedAppendOnlyMap是继承自SizeTrackingAppendOnlyMap，而SizeTrackingAppendOnlyMap又继承自AppendOnlyMap，这里调用的changeValue方法实际上是SizeTrackingAppendOnlyMap的changeValue方法：
override def changeValue(key: K, updateFunc: (Boolean, V) => V): V = {
// 首先调用父类的changeValue方法
val newValue = super.changeValue(key, updateFunc)
// 然后调用SizeTracker接口的afterUpdate方法
super.afterUpdate()
// 返回newValue
newValue
}
父类(AppendOnlyMap)的changeValue方法：
def changeValue(key: K, updateFunc: (Boolean, V) => V): V = {
assert(!destroyed, destructionMessage)
val k = key.asInstanceOf[AnyRef]
// key为空时候的处理，增加长度
if (k.eq(null)) {
if (!haveNullValue) {
incrementSize()
}
nullValue = updateFunc(haveNullValue, nullValue)
haveNullValue = true
return nullValue
}
var pos = rehash(k.hashCode) & mask
var i = 1
while (true) {
// 这里的data是一个数组，用来同时存储key和value：key0, value0, key1, value1, key2, value2, etc.
// 即2 * pos上存储的是key的值，2 * pos + 1上存储的是value的值
val curKey = data(2 * pos)
// 如果key已经存在，就调用updateFunc方法更新value
if (k.eq(curKey) || k.equals(curKey)) {
val newValue = updateFunc(true, data(2 * pos + 1).asInstanceOf[V])
data(2 * pos + 1) = newValue.asInstanceOf[AnyRef]
return newValue
} else if (curKey.eq(null)) {
// 如果key不存在就将该key和对应的value添加到data这个数组中
val newValue = updateFunc(false, null.asInstanceOf[V])
data(2 * pos) = k
data(2 * pos + 1) = newValue.asInstanceOf[AnyRef]
incrementSize()
return newValue
} else {
// 否则继续计算位置(pos)
val delta = i
pos = (pos + delta) & mask
i += 1
}
}
null.asInstanceOf[V] // Never reached but needed to keep compiler happy
}
然后是SizeTracker接口的afterUpdate方法
protected def afterUpdate(): Unit = {
numUpdates += 1
if (nextSampleNum == numUpdates) {
takeSample()
}
}
更新数据的更新次数，如果更新的次数达到nextSampleNum，就执行采样操作，主要用来评估内存的使用情况。
PartitionedPairBuffer：
再来看PartitionedPairBuffer的insert方法，也就是不进行Map端combine操作的情况：
def insert(partition: Int, key: K, value: V): Unit = {
// 如果当前的大小达到了capacity的值就需要扩大该数组
if (curSize == capacity) {
growArray()
}
// 存储key，这里存储的是(partition Id, key)的格式
data(2 * curSize) = (partition, key.asInstanceOf[AnyRef])
// 存储value
data(2 * curSize + 1) = value.asInstanceOf[AnyRef]
curSize += 1
// 参考上面PartitionedAppendOnlyMap的部分
afterUpdate()
}
直接将数据存储到buffer中。
执行完上面的更新数据操作后，就要判断是否要将数据spill到磁盘，即maybeSpillCollection方法：
private def maybeSpillCollection(usingMap: Boolean): Unit = {
var estimatedSize = 0L
// 这里需要判断使用的是map(PartitionedAppendOnlyMap)还是buffer(PartitionedPairBuffer)
// 如果true就是map，false就是buffer
if (usingMap) {
// 估计当前map的内存占用大小
estimatedSize = map.estimateSize()
// 如果超过内存的限制，就将缓存中的数据spill到磁盘
if (maybeSpill(map, estimatedSize)) {
// spill到磁盘后，重置缓存
map = new PartitionedAppendOnlyMap[K, C]
}
} else {
// 不进行Map端聚合操作的情况
estimatedSize = buffer.estimateSize()
if (maybeSpill(buffer, estimatedSize)) {
buffer = new PartitionedPairBuffer[K, C]
}
}
if (estimatedSize > _peakMemoryUsedBytes) {
_peakMemoryUsedBytes = estimatedSize
}
}
上面代码的主要作用就是估计当前缓存(map或者buffer)使用内存的大小，如果超过了内存使用的限制，就要将缓存中的数据spill到磁盘中，同时重置当前的缓存。
下面就来看一下maybeSpill方法：
// 如果成功spill到磁盘就返回true，否则返回false
protected def maybeSpill(collection: C, currentMemory: Long): Boolean = {
var shouldSpill = false
// 在进行真正的spill操作之前向TaskMemoryManager申请再多分配一些内存
if (elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold) {
// Claim up to double our current memory from the shuffle memory pool
val amountToRequest = 2 * currentMemory - myMemoryThreshold
val granted =
taskMemoryManager.acquireExecutionMemory(amountToRequest, MemoryMode.ON_HEAP, null)
myMemoryThreshold += granted
// If we were granted too little memory to grow further (either tryToAcquire returned 0,
// or we already had more memory than myMemoryThreshold), spill the current collection
// 如果内存仍然不够用，就认定为需要spill到磁盘
shouldSpill = currentMemory >= myMemoryThreshold
}
// 如果内存中元素的个数超过了强制spill的上限也会认定为需要进行spill操作
shouldSpill = shouldSpill || _elementsRead > numElementsForceSpillThreshold
// Actually spill
// 接下来就真正将数据spill到磁盘
if (shouldSpill) {
_spillCount += 1
logSpillage(currentMemory)
// spill操作
spill(collection)
_elementsRead = 0
_memoryBytesSpilled += currentMemory
releaseMemory()
}
shouldSpill
}
在进行真正的spill操作之前会向TaskMemoryManager申请再多分配一些内存，如果还不能够满足，或者不能分配更多的内存，或者内存中元素的个数超过了强制spill的上限，最终就会执行spill操作，接下来进入spill方法：
// 这里的collection就是指的map或者buffer
override protected[this] def spill(collection: WritablePartitionedPairCollection[K, C]): Unit = {
// Because these files may be read during shuffle, their compression 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 context.
// 获取临时的BlockId(TempShuffleBlockId)及对应的File
val (blockId, file) = diskBlockManager.createTempShuffleBlock()
// These variables are reset after each flush
var objectsWritten: Long = 0
var spillMetrics: ShuffleWriteMetrics = null
var writer: DiskBlockObjectWriter = null
def openWriter(): Unit = {
assert (writer == null && spillMetrics == null)
spillMetrics = new ShuffleWriteMetrics
writer = blockManager.getDiskWriter(blockId, file, serInstance, fileBufferSize, spillMetrics)
}
// 获得DiskWriter(DiskBlockObjectWriter)
openWriter()
// List of batch sizes (bytes) in the order they are written to disk
// 用来储存每个batch对应的size
val batchSizes = new ArrayBuffer[Long]
// How many elements we have in each partition
// 用来储存每个partition有多少元素
val elementsPerPartition = new Array[Long](numPartitions)
// Flush the disk writer's contents to disk, and update relevant variables.
// The writer is closed at the end of this process, and cannot be reused.
def flush(): Unit = {
val w = writer
writer = null
w.commitAndClose()
_diskBytesSpilled += spillMetrics.shuffleBytesWritten
batchSizes.append(spillMetrics.shuffleBytesWritten)
spillMetrics = null
objectsWritten = 0
}
var success = false
try {
// 排序部分的操作，返回迭代器
val it = collection.destructiveSortedWritablePartitionedIterator(comparator)
// 循环的到的迭代器，执行write操作
while (it.hasNext) {
val partitionId = it.nextPartition()
require(partitionId >= 0 && partitionId < numPartitions,
s"partition Id: ${partitionId} should be in the range [0, ${numPartitions})")
it.writeNext(writer)
elementsPerPartition(partitionId) += 1
objectsWritten += 1
// 如果写的对象达到serializerBatchSize的个数时就进行flush操作
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，并保存在数据结构ArrayBuffer[SpilledFile]中
spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
}
先来简单的看一下排序部分的逻辑：
def destructiveSortedWritablePartitionedIterator(keyComparator: Option[Comparator[K]])
: WritablePartitionedIterator = {
// 这里的partitionedDestructiveSortedIterator会根据是map或者buffer有不同的实现
val it = partitionedDestructiveSortedIterator(keyComparator)
// 最后返回的是WritablePartitionedIterator，上面进行写操作的时候就是调用该迭代器中的writeNext方法
new WritablePartitionedIterator {
private[this] var cur = if (it.hasNext) it.next() else null
def writeNext(writer: DiskBlockObjectWriter): Unit = {
writer.write(cur._1._2, cur._2)
cur = if (it.hasNext) it.next() else null
}
def hasNext(): Boolean = cur != null
def nextPartition(): Int = cur._1._1
}
}
如果collection是map，具体的实现为(PartitionedAppendOnlyMap)：
def partitionedDestructiveSortedIterator(keyComparator: Option[Comparator[K]])
: Iterator[((Int, K), V)] = {
val comparator = keyComparator.map(partitionKeyComparator).getOrElse(partitionComparator)
destructiveSortedIterator(comparator)
}
如果collection是buffer，具体的实现为(PartitionedPairBuffer)：
override def partitionedDestructiveSortedIterator(keyComparator: Option[Comparator[K]])
: Iterator[((Int, K), V)] = {
val comparator = keyComparator.map(partitionKeyComparator).getOrElse(partitionComparator)
new Sorter(new KVArraySortDataFormat[(Int, K), AnyRef]).sort(data, 0, curSize, comparator)
iterator
}
这里需要注意，首先都是获取比较器，比较的数据格式为((partition Id, key), value)，而两者的区别在于前者才是真正的Destructive级别的，具体的实现在destructiveSortedIterator方法中，而不管采用哪种方式，其底层都是通过timSort算法实现的，具体的排序逻辑就不详细说明了，有兴趣的朋友可以深入研究下去。
接下来就进入到第三步。
2.1.3 第三步
先贴出该步骤的代码：
...
// 这里注释说的很清楚了，只打开了一个文件
// 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).
val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
val tmp = Utils.tempFileWith(output)
try {
// 构造blockId
val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
// 写数据
val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
// 写index文件
shuffleBlockResolver.writeIndexFileAndCommit(dep.shuffleId, mapId, partitionLengths, tmp)
// 进行Shuffle Read的时候需要参考该信息
mapStatus = MapStatus(blockManager.shuffleServerId, partitionLengths)
} finally {
if (tmp.exists() && !tmp.delete()) {
logError(s"Error while deleting temp file ${tmp.getAbsolutePath}")
}
}
writePartitionedFile：
首先来看writePartitionedFile方法：
def writePartitionedFile(
blockId: BlockId,
outputFile: File): Array[Long] = {
// Track location of each range in the output file
val lengths = new Array[Long](numPartitions)
// 首先判断spills中是否有数据，即判断是否有数据被spill到了磁盘中
if (spills.isEmpty) {
// 数据只在内存中的情况
// Case where we only have in-memory data
val collection = if (aggregator.isDefined) map else buffer
// 获得迭代器
val it = collection.destructiveSortedWritablePartitionedIterator(comparator)
// 进行迭代并将数据写到磁盘
while (it.hasNext) {
val writer = blockManager.getDiskWriter(blockId, outputFile, serInstance, fileBufferSize,
context.taskMetrics.shuffleWriteMetrics.get)
val partitionId = it.nextPartition()
while (it.hasNext && it.nextPartition() == partitionId) {
it.writeNext(writer)
}
writer.commitAndClose()
val segment = writer.fileSegment()
// 最后返回的是每个partition写入的数据的长度
lengths(partitionId) = segment.length
}
} else {
// 如果有数据被spill到了磁盘中，我们就需要进行merge-sort操作
// We must perform merge-sort; get an iterator by partition and write everything directly.
for ((id, elements) <- this.partitionedIterator) {
if (elements.hasNext) {
val writer = blockManager.getDiskWriter(blockId, outputFile, serInstance, fileBufferSize,
context.taskMetrics.shuffleWriteMetrics.get)
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.internalMetricsToAccumulators(
InternalAccumulator.PEAK_EXECUTION_MEMORY).add(peakMemoryUsedBytes)
lengths
}
下面我们就看一下this.partitionedIterator即内存和磁盘中的数据是如何合到一起的：
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) {
// Special case: if we have only in-memory data, we don't need to merge streams, and perhaps
// we don't even need to sort by anything other than partition ID
if (!ordering.isDefined) {
// The user hasn't requested sorted keys, so only sort by partition ID, not key
groupByPartition(collection.partitionedDestructiveSortedIterator(None))
} else {
// We do need to sort by both partition ID and key
groupByPartition(collection.partitionedDestructiveSortedIterator(Some(keyComparator)))
}
} else {
// Merge spilled and in-memory data
merge(spills, collection.partitionedDestructiveSortedIterator(comparator))
}
}
我们只考虑spills不为空的情况，即执行merge方法：
private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
: Iterator[(Int, Iterator[Product2[K, C]])] = {
// 根据每个SpilledFile实例化一个SpillReader，这些SpillReader组成一个Seq
val readers = spills.map(new SpillReader(_))
// 获得内存BufferedIterator
val inMemBuffered = inMemory.buffered
// 根据partition的个数进行迭代
(0 until numPartitions).iterator.map { p =>
// 实例化IteratorForPartition，即当前partition下的Iterator
val inMemIterator = new IteratorForPartition(p, inMemBuffered)
// 这里就是合并操作
val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
if (aggregator.isDefined) {
// Perform partial aggregation across partitions
// 如果需要map端的combine操作则需要根据key进行聚合操作
(p, mergeWithAggregation(
iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
} else if (ordering.isDefined) {
// No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
// sort the elements without trying to merge them
// 排序合并，例如sortByKey
(p, mergeSort(iterators, ordering.get))
} else {
(p, iterators.iterator.flatten)
}
}
}
具体的mergeWithAggregation和mergeSort就不一一说明了，下面再来看一下writeIndexFileAndCommit
writeIndexFileAndCommit：
再来看writeIndexFileAndCommit方法：
def writeIndexFileAndCommit(
shuffleId: Int,
mapId: Int,
lengths: Array[Long],
dataTmp: File): Unit = {
val indexFile = getIndexFile(shuffleId, mapId)
val indexTmp = Utils.tempFileWith(indexFile)
try {
val out = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(indexTmp)))
Utils.tryWithSafeFinally {
// We take in lengths of each block, need to convert it to offsets.
var offset = 0L
out.writeLong(offset)
for (length <- lengths) {
offset += length
out.writeLong(offset)
}
} {
out.close()
}
val dataFile = getDataFile(shuffleId, mapId)
// There is only one IndexShuffleBlockResolver per executor, this synchronization make sure
// the following check and rename are atomic.
synchronized {
val existingLengths = checkIndexAndDataFile(indexFile, dataFile, lengths.length)
if (existingLengths != null) {
// Another attempt for the same task has already written our map outputs successfully,
// so just use the existing partition lengths and delete our temporary map outputs.
System.arraycopy(existingLengths, 0, lengths, 0, lengths.length)
if (dataTmp != null && dataTmp.exists()) {
dataTmp.delete()
}
indexTmp.delete()
} else {
// This is the first successful attempt in writing the map outputs for this task,
// so override any existing index and data files with the ones we wrote.
if (indexFile.exists()) {
indexFile.delete()
}
if (dataFile.exists()) {
dataFile.delete()
}
if (!indexTmp.renameTo(indexFile)) {
throw new IOException("fail to rename file " + indexTmp + " to " + indexFile)
}
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}")
}
}
}
具体的实现就不详细说明了，主要就是根据上一步的到的partition长度的数组将偏移量写入到index文件中。
最后就是实例化MapStatus，shuffle read的时候根据MapStatus获取数据。至此BaseShuffleHandle & SortShuffleWriter的部分就结束了。
BypassMergeSortShuffleHandle & BypassMergeSortShuffleWriter
再来看一下BypassMergeSortShuffleWriter的write方法：
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];
// 针对每一个reducer建立一个临时文件
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).open();
}
// Creating the file to write to and creating a disk writer both involve interacting with
// the disk, and can take a long time in aggregate when we open many files, so should be
// included in the shuffle write time.
writeMetrics.incShuffleWriteTime(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();
}
// 将所有的临时文件内容按照partition Id合并到一个文件
File output = shuffleBlockResolver.getDataFile(shuffleId, mapId);
File tmp = Utils.tempFileWith(output);
try {
// 将记录和partition长度信息分别写入到data文件和index文件中
partitionLengths = writePartitionedFile(tmp);
shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, tmp);
} finally {
if (tmp.exists() && !tmp.delete()) {
logger.error("Error while deleting temp file {}", tmp.getAbsolutePath());
}
}
mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths);
}
该writer在进行写记录之前会根据reducer的个数(例如R个)生成R个临时文件，然后将记录写入对应的临时文件中，最后将这些文件进行合并操作并写入到一个文件中，由于直接将记录写入了临时文件，并没有缓存在内存中，所以如果reducer的个数过多的话，就会为每个reducer打开一个临时文件，如果reducer的数量过多的话就会影响性能，所以使用该种方式需要满足一下条件(下面是源码中的注释)：
- no Ordering is specified.
- no Aggregator is specified.
- the number of partitions is less than spark.shuffle.sort.bypassMergeThreshold.
其中spark.shuffle.sort.bypassMergeThreshold的个数默认为200个。
SerializedShuffleHandle & UnsafeShuffleWriter
最后再来看一下UnsafeShuffleWriter，也就是通常所说的Tungsten。
UnsafeShuffleWriter的write方法：
public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOException {
// Keep track of success so we know if we encountered an exception
// We do this rather than a standard try/catch/re-throw to handle
// generic throwables.
boolean success = false;
try {
while (records.hasNext()) {
// 循环便利所有记录，对其作用insertRecordIntoSorter方法
insertRecordIntoSorter(records.next());
}
// 将数据输出到磁盘上
closeAndWriteOutput();
success = true;
} finally {
if (sorter != null) {
try {
sorter.cleanupResources();
} catch (Exception e) {
// Only throw this error if we won't be masking another
// error.
if (success) {
throw e;
} else {
logger.error("In addition to a failure during writing, we failed during " +
"cleanup.", e);
}
}
}
}
}
首先来看insertRecordIntoSorter：
void insertRecordIntoSorter(Product2<K, V> record) throws IOException {
assert(sorter != null);
final K key = record._1();
final int partitionId = partitioner.getPartition(key);
serBuffer.reset();
serOutputStream.writeKey(key, OBJECT_CLASS_TAG);
serOutputStream.writeValue(record._2(), OBJECT_CLASS_TAG);
serOutputStream.flush();
final int serializedRecordSize = serBuffer.size();
assert (serializedRecordSize > 0);
sorter.insertRecord(
serBuffer.getBuf(), Platform.BYTE_ARRAY_OFFSET, serializedRecordSize, partitionId);
}
可以看出实际上调用的是ShuffleExternalSorter的insertRecord方法，限于篇幅，具体的底层实现暂不说明，以后有时间会单独分析一下Tungsten的部分，接下来看一下closeAndWriteOutput方法：
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 {
try {
// 将spilled的文件合并并写入到临时文件
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());
}
}
}
// 将partition的长度信息写入index文件
shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, tmp);
} finally {
if (tmp.exists() && !tmp.delete()) {
logger.error("Error while deleting temp file {}", tmp.getAbsolutePath());
}
}
// mapStatus
mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths);
}
用图来总结一下上面描述的三种方式如下所示：
BaseShuffleHandle & SortShuffleWriter
SerializedShuffleHandle & UnsafeShuffleWriter
BypassMergeSortShuffleHandle & BypassMergeSortShuffleWriter
最后需要说明的是如果采用的是SortShuffleManager，最后每个task产生的文件的个数为2 * M(M代表Mapper端ShuffleMapTask的个数)，相对于Hash的方式来说文件的个数明显减少。
至此Shuffle Write的部分就分析完了，下一遍文章会继续分析Shuffle Read的部分。
本文参照的是Spark 1.6.3版本的源码，同时给出Spark 2.1.0版本的连接：
本文为原创，欢迎转载，转载请注明出处、作者，谢谢！
