SortShuffle源码分析上
在上一篇博客中,我们从源码的角度分析了HashShuffle的两种机制的区别,对应之前的理论分析,现在我们分析SortShuffle的两种机制的区别。
同理我们首先看SortShuffleWriter的write()方法:
override def write(records: Iterator[Product2[K, V]]): Unit = {
// 创建一个非常重要的组件ExternalSorter,它会对数据进行聚合、分区和排序等操作。
// 如果需要进行本地聚合,就将aggregator传入,假如不需要聚合,那么就传入none
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)
}
// 这个函数特别重要,对partition数据进行排序、聚合或一些到磁盘等操作
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).
// 获取partition数据最终要聚合的那个文件目录
val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
// 生成临时文件名
val tmp = Utils.tempFileWith(output)
// 生成一个ShuffleBlockId
val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
// 调用ExternalSorter的writePartitionedFile方法
// 将已经处理的每个partition的数据全部写入磁盘文件中,返回offset数组,数组下标就是partitionID,值是offset
val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
// 将写入磁盘文件中的partition的偏移量写入索引文件中,供reduce端进行数据拉取。
shuffleBlockResolver.writeIndexFileAndCommit(dep.shuffleId, mapId, partitionLengths, tmp)
// 返回MapStatus,包含了写入每个partition的在输出文件中的offset的信息
mapStatus = MapStatus(blockManager.shuffleServerId, partitionLengths)
}
上面创建了一个非常重要的组件,ExternalSorter,它会对数据进行聚合、分区和排序等操作,根据是否需要本地聚合及排序来创建两种不同的sorter(这是不是bypass和非bypass的区别?),然后使用insertAll()将records也即partition数据进行聚合、排序和溢写到磁盘操作,并生成多个临时文件。接着创建Task最终输出文件名,使用writePartitionedFile将之前溢写到磁盘的临时文件中的数据进行merge,并写入Task输出文件中,也即一个Task就只会创建一个输出文件,在写入的过程中记录每个partition的写入文件的offset,并封装在MapStatus返回(其实这里注意一点,这些spill创建的临时文件,在SortShuffleWriter调用stop的时候会被删除,这点在之前源码中有体现,在runTask()这个方法中)。
insertAll方法
这里面一个很重要的方法 insertAll(),我们下面分析一下这个方法:
def insertAll(records: Iterator[Product2[K, V]]): Unit = {
// 是否需要聚合
val shouldCombine = aggregator.isDefined
// 如果需要本地聚合
if (shouldCombine) {
// 合并输入类型,比如 (C, V) => C,将C类型和V进行合并得到C。
val mergeValue = aggregator.get.mergeValue
// 聚合器,它会在每个分区上都会执行,只要遇到没有处理过的key,就会执行该方法
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)
}
// 遍历partition数据
while (records.hasNext) {
// 用于记录写入数据
addElementsRead()
// 取出一个数据
kv = records.next()
// 将聚合后的值存入本地聚合缓存 PartitionedAppendOnlyMap中
// 这里的map是很重要的一个组件,后面会反复使用,一会儿单独分析
map.changeValue((getPartition(kv._1), kv._1), update)
// 判断是否需要将数据溢写到磁盘
maybeSpillCollection(usingMap = true)
}
} else {
// 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)
}
}
}
insertAll()将partition数据取出放入map或buffer中,其中map是一个很重要的组件,它自定义了比较器Comparator,用于对key进行排序。buffer相比于map没有比较器,它只是存储partition的数据。下面我们先看一下map的类型。
// map的定义,注意是可变类型
private var map = new PartitionedAppendOnlyMap[K, C]
private[spark] class PartitionedAppendOnlyMap[K, V]
extends SizeTrackingAppendOnlyMap[(Int, K), V] with WritablePartitionedPairCollection[K, V] {
// 定义比较器并返回
def partitionedDestructiveSortedIterator(keyComparator: Option[Comparator[K]])
: Iterator[((Int, K), V)] = {
// partitionKeyComparator: 这个比较器,获取partitionId的差值,或者 key的差值
// partitionComparator : 比较两个(partitionId, key)的大小
// 其实这里就有点类似自定义的一个二次排序,先按照partitionID进行排序,相等则按照key进行排序。
// 说白了就是,不同partition之间按照partitionID排序,相同partition内部数据按照key排序
val comparator = keyComparator.map(partitionKeyComparator).getOrElse(partitionComparator)
// 使用自定义的comparator,对数据进行排序,内部使用到的是归并排序。
destructiveSortedIterator(comparator)
}
def insert(partition: Int, key: K, value: V): Unit = {
update((partition, key), value)
}
}
def destructiveSortedIterator(keyComparator: Comparator[K]): Iterator[(K, V)] = {
destroyed = true
// Pack KV pairs into the front of the underlying array
var keyIndex, newIndex = 0
// 对data中的数据进行整理,都放在前面
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))
// 使用前面生成的keyComparator进行排序,内部使用的是TimSort,一种归并排序
new Sorter(new KVArraySortDataFormat[K, AnyRef]).sort(data, 0, newIndex, keyComparator)
// 返回数据迭代器,里面的数据是经过归并排序的数据,为什么是归并排序,
// 可能是因为多个executor并行运行的关系,多个task的数据进行归并排序
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
}
}
}
}
map的类型是PartitionedAppendOnlyMap,它里面自定义了一个比较器,这个比较器的实现原理也很简单,不同分区之间按照partition ID的大小进行排序,每个partition内部的数据则按照key的大小排序。接着我们分析里面一个很重要的方法maybeSpillCollection(),它是将数据溢写到磁盘的关键。
maybeSpillCollection是否溢写到磁盘
这个方法就是判断数据是否溢写到磁盘
private def maybeSpillCollection(usingMap: Boolean): Unit = {
var estimatedSize = 0L
// 假如在本地聚合
if (usingMap) {
// 当前数据大小
estimatedSize = map.estimateSize()
// maybeSpill,判断数据是否需要溢写到磁盘
if (maybeSpill(map, estimatedSize)) {
// 重新申请缓存
map = new PartitionedAppendOnlyMap[K, C]
}
} else {
estimatedSize = buffer.estimateSize()
if (maybeSpill(buffer, estimatedSize)) {
buffer = new PartitionedPairBuffer[K, C]
}
}
// 更新当前使用的内存大小
if (estimatedSize > _peakMemoryUsedBytes) {
_peakMemoryUsedBytes = estimatedSize
}
}
这个函数里面依然封装了一个函数maybeSpill(),用于判断当前数据量是否足够大,是否需要溢写到磁盘,下面分析一下这个方法:
maybeSpill()判断是否需要溢写到磁盘
protected def maybeSpill(collection: C, currentMemory: Long): Boolean = {
var shouldSpill = false
// 当前字节是否是32整数倍,并且必须大于5M
if (elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold) {
// 首先尝试对当前缓存进行扩容
val amountToRequest = 2 * currentMemory - myMemoryThreshold
// 去TaskMemoryManager上去尝试获取内存
val granted =
taskMemoryManager.acquireExecutionMemory(amountToRequest, MemoryMode.ON_HEAP, null)
// 将获取到的内存更新到阈值中
myMemoryThreshold += granted
// 假如当前数据大小依然大于扩容后的阈值,那么shouldSpill为true,否则为false
shouldSpill = currentMemory >= myMemoryThreshold
}
// 如果计数器的值大于Long.MaxValue,那么shouldSpill也为true,否则为false
shouldSpill = shouldSpill || _elementsRead > numElementsForceSpillThreshold
// 根据上述两个阈值判断结果,决定如果需要溢写到磁盘
if (shouldSpill) {
// 溢写计数器,记录产生多少文件
_spillCount += 1
// 打印溢写日志信息
logSpillage(currentMemory)
// 将数据集合(key-value)溢写到磁盘
spill(collection)
// 计数器清零
_elementsRead = 0
// 释放内存
_memoryBytesSpilled += currentMemory
releaseMemory()
}
// 返回是否溢写到磁盘,继而重新申请map
shouldSpill
}
maybeSpill的主要功能就是判断什么情况下需要溢写到磁盘,假设大于扩容后的门限,或者达到Long.MAX_VALUE,那么就会发生溢写操作,否则就不进行溢写,这就对应之前shuffle操作中的达到一定的阈值就写入磁盘。这里就是那个阈值,它通过两种方法判断。
由于SortShuffle的中方法嵌套调用比较多,这里拆分上下节进行分析;总结一下,insertAll()对整个Task的数据进行处理,包括对数据进行排序、聚合和溢写磁盘等操作,并合并溢写到磁盘的临时文件,以及对应的索引文件等。其中判断什么时候进行溢写操作,主要是通过两个阈值,一个是数据大小,比如超过5M(默认大小,但是当Task内存足够,那么会相应扩容),另一个就是根据数量来进行溢写,默认大小是Long.MaxValue。下篇我们会分析如何进行溢写操作spill(collection)这个方法。