Spark Shuffle源码分析系列之SortShuffleWriter

最新推荐文章于 2022-02-13 01:08:12 发布

小白数据猿

最新推荐文章于 2022-02-13 01:08:12 发布

阅读量619

点赞数 6

分类专栏： Spark 文章标签： spark 大数据 java

本文链接：https://blog.csdn.net/lidongmeng0213/article/details/109488654

版权

Spark 专栏收录该内容

42 篇文章 10 订阅

订阅专栏

上一节我们分析了BypassMergeSortShuffleWriter，它是Hash风格的ShuffleWriter，主要适用于不需要map-side聚合排序，并且分区数目小于spark.shuffle.sort.bypassMergeThreshold<默认200>；本节我们来介绍SortBased-ShuffleWriter，它可以处理map端的聚合排序操作，是默认的ShuffleWriter。

ShuffleMapTask整体流程

Spark中负责具体的Shuffle Map端执行任务的逻辑在ShuffleMapTask，当任务提交后，Executor会将任务交给处理任务的线程池，最终调用的是Task中的runTask方法，ShuffleMapTask的具体实现步骤如下:

首先反序列化出Task，得到改ShuffleTask依赖的rdd和dependency；
获取shuffleManager，默认是SortShuffleManager，负责获取相应的handle，writer，reader等；
根据dependency的shuffleHandle以及当前任务处理的依赖rdd的具体分区id得到相应的写入者writer；
处理依赖的RDD的分区id的相应的迭代数据，将计算的中间结果写入磁盘文件，生成索引文件和数据文件；
返回MapStatus，回传给MapOutputTracker进行中间结果的记录，供ShuffleReader获取Map数据的位置和大小信息。

// org.apache.spark.scheduler.ShuffleMapTask
override def runTask(context: TaskContext): MapStatus = {
  val threadMXBean = ManagementFactory.getThreadMXBean
  // 记录反序列化开始时间
  val deserializeStartTime = System.currentTimeMillis()
  val deserializeStartCpuTime = if (threadMXBean.isCurrentThreadCpuTimeSupported) {
    threadMXBean.getCurrentThreadCpuTime
  } else 0L
  // 获取反序列化器
  val ser = SparkEnv.get.closureSerializer.newInstance()
  // 对Task数据进行反序列化，得到RDD和ShuffleDependency
  val (rdd, dep) = ser.deserialize[(RDD[_], ShuffleDependency[_, _, _])](
    ByteBuffer.wrap(taskBinary.value), Thread.currentThread.getContextClassLoader)
  // 得到反序列化时间
  _executorDeserializeTime = System.currentTimeMillis() - deserializeStartTime
  _executorDeserializeCpuTime = if (threadMXBean.isCurrentThreadCpuTimeSupported) {
    threadMXBean.getCurrentThreadCpuTime - deserializeStartCpuTime
  } else 0L

  // Shuffle写出器
  var writer: ShuffleWriter[Any, Any] = null
  try {
    // 将计算的中间结果写入磁盘文件
    val manager = SparkEnv.get.shuffleManager
    // 获取对指定分区的数据进行磁盘写操作的SortShuffleWriter
    writer = manager.getWriter[Any, Any](dep.shuffleHandle, partitionId, context)
    // 调用RDD的iterator方法进行迭代计算，将计算的中间结果写入磁盘文件
    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
  }
}

SortShuffleWriter

SortShuffleWriter主要负责Shuffle数据写入内存+磁盘的逻辑，SortShuffleWriter使用ExternalSorter作为排序器，具体的执行流程如下:

org.apache.spark.shuffle.sort.SortShuffleWriter
org.apache.spark.util.collection.ExternalSorter
org.apache.spark.util.collection.{PartitionedAppendOnlyMap, SizeTrackingAppendOnlyMap,
                                  AppendOnlyMap, SizeTracker, WritablePartitionedPairCollection, Spillable}

(records: Iterator[Product2[K, V]])   --改Task处理的数据迭代器
-- SortShuffleWriter.write
	-- ExternalSorter.sorter.insertAll
		----- k-v record write
				--- SizeTrackingAppendOnlyMap.map.changeValue  --有聚合时候
					==> AppendOnlyMap.changeValue SizeTrackingAppendOnlyMap.afterUpdate
				--- PartitionedPairBuffer.buffer.insert    --不需要聚合时候，只缓冲
					==> SizeTracker.afterUpdate
		-- ExternalSorter.maybeSpillCollection  --- 申请内存，不能申请到会溢写磁盘
			==> SizeTracker.estimateSize   -- 估计内存占用大小
      ==> Spillable.maybeSpill
         ==> MemoryConsumer.acquireMemory  <获取执行内存> 
             ==> 可能使个consumer进行forceSpill操作
 -- ExternalSorter.writePartitionedFile    -- 合并磁盘+内存数据
 -- shuffleBlockResolver.writeIndexFileAndCommit  --生成索引和数据文件

成员变量

SortShuffleWriter具有以下成员变量:

shuffleBlockResolver用来生成索引和数据文件；
BaseShuffleHandle用来记录其中依赖的ShuffleDependency，可以得到相应的分区器，以及是否mapSideCombine，排序器，聚合器等；
mapId该任务处理的MapTask的唯一标识；
context是任务的一些上下文信息；
ExternalSorter实质性具体的内存+磁盘写入的类；
MapStatus是该任务生成的结果。

private[spark] class SortShuffleWriter[K, V, C](
    shuffleBlockResolver: IndexShuffleBlockResolver,
    handle: BaseShuffleHandle[K, V, C],
    mapId: Int,
    context: TaskContext) {
      private val dep = handle.dependency
      private val blockManager = SparkEnv.get.blockManager
      private var sorter: ExternalSorter[K, V, _] = null
      
      // map任务的状态，即MapStatus。
      private var mapStatus: MapStatus = null
 }

写入函数

SortShuffleWriter实现了ShuffleWriter的抽象方法write，负责具体数据的处理，步骤如下:

SortShuffleWriter数据的写入主要是靠ExternalSorter来处理的，首先根据依赖关系，看是否需要进行mapSideCombine，如果需要则需要制定聚合器，key比较器来初始化ExternalSorter，否则这两系那个为None来初始化ExternalSorter即可；
利用ExternalSorter来进行真正的数据写入，优先内存写入，内存不足时候溢写磁盘处理；
由于ExternalSorter会生成多个磁盘文件，每个磁盘文件内部都是按照partitionId进行排序的，有可能分区内部还是按照key来排序的，所以需要进行聚合成一个文件，是通过ExternalSorter的writePartitionedFile来进行的，合并多个文件为一个，写在一个临时文件中，并且计算了各个分区的数据的大小；
利用shuffleBlockResolver生成索引文件和数据文件；
更新mapStatus，在任务结束时候，ShuffleMapTask清理ExternalSorter结束返回相应的mapStatus，完成一个ShuffleMapTask的运行任务。

override def write(records: Iterator[Product2[K, V]]): Unit = {
  // 创建ExternalSorter，dep.mapSideCombine决定了ExternalSorter
  // 选择内存数据的存储格式PartitionedAppendOnlyMap还是PartitionedPairBuffer。
  sorter = if (dep.mapSideCombine) {
    require(dep.aggregator.isDefined, "Map-side combine without Aggregator specified!")
    // 将ShuffleDependency的aggregator和keyOrdering传递给ExternalSorter的aggregator和ordering属性
    new ExternalSorter[K, V, C](
      context, dep.aggregator, Some(dep.partitioner), dep.keyOrdering, dep.serializer)
  } else {
    // 未传递ShuffleDependency的aggregator和keyOrdering给ExternalSorter的aggregator和ordering属性
    new ExternalSorter[K, V, V](
      context, aggregator = None, Some(dep.partitioner), ordering = None, dep.serializer)
  }

  // 将map任务的输出记录插入到缓存中
  sorter.insertAll(records)

  // 获取Shuffle数据文件
  val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
  val tmp = Utils.tempFileWith(output)
  try {
    // 将map端缓存的数据写入到磁盘中，并生成Block文件对应的索引文件
    val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
    // 将map端缓存的数据写入到磁盘中，返回值记录了各个分区的长度。
    val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
    // 生成Block文件对应的索引文件，用于记录各个分区在Block文件中对应的偏移量，以便于reduce任务拉取时使用。
    shuffleBlockResolver.writeIndexFileAndCommit(dep.shuffleId, mapId, partitionLengths, tmp)
    // 构造并返回MapStatus
    mapStatus = MapStatus(blockManager.shuffleServerId, partitionLengths)
  } finally {
    if (tmp.exists() && !tmp.delete()) {
      logError(s"Error while deleting temp file ${tmp.getAbsolutePath}")
    }
  }
}

所以可以看出来，具体的写入部分工作是交由ExternalSorterl来进行的，下面我们来详细分析下写入逻辑。

ExternalSorter

ExternalSorter对类型[K，V]的多个键值对进行排序并可能合并，以生成类型[K，C]的键组合对。执行流程是先使用分区器计算出key所在的分区，然后可以选择使用自定义排序对每个分区中的key进行排序，有可能还需要对相同key的数据进行聚合操作，最后为每个分区输出具有不同字节范围的单个分区文件，提供给后续ShuffleReader读取使用。

概述

成员变量

private[spark] class ExternalSorter[K, V, C](
    context: TaskContext,
    aggregator: Option[Aggregator[K, V, C]] = None,
    partitioner: Option[Partitioner] = None,
    ordering: Option[Ordering[K]] = None,
    serializer: Serializer = SparkEnv.get.serializer)
  extends Spillable[WritablePartitionedPairCollection[K, C]](context.taskMemoryManager())
  with Logging {
    ...
}

ExternalSorter的成员变量如下：

aggregator可选的聚合器，如果指定了，会进行相同key的数据合并操作；
partitioner可选的分区器，默认是hash分区，先按分区Id排序，再按key排序；
ordering 可选的排序器，它在每一个分区内按key进行排序，如果指定了，最后生成文件需要先按照分区排序，分区内按照key排序；
serializer用于溢出内存数据到磁盘的序列化器。

Partitioner分区器

ExternalSorter会根据分区器进行计算Reduce的分区个数，对于每个key都需要先计算分区Id，写入磁盘的数据也是要按照分区进行排序的。

// 分区数量。默认为1。
private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
// 是否有分区。当numPartitions大于1时为true。
private val shouldPartition = numPartitions > 1
// 使用分区器获取键的分区
private def getPartition(key: K): Int = {
  if (shouldPartition) partitioner.get.getPartition(key) else 0
}

分区器默认是基于Hash的排序器，分区的数目是按照以下顺序决定的：

如果依赖的RDD中有包含分区器的，而且分区数目大于0，选择它们中最大的那个RDD的分区器；
如果依赖的RDD中没有包含分区器的，需要使用默认的HashPartitioner，如果指定了spark.default.parallelism，则分区数目等于这个参数配置的默认分区数目；否则使用依赖RDD中最大分区数目。

// org.apache.spark.Partitioner
def defaultPartitioner(rdd: RDD[_], others: RDD[_]*): Partitioner = {
  val rdds = (Seq(rdd) ++ others)
  // 先判断rdds中是否有RDD的分区器存在且指定分区数量大于0，过滤出这些RDD
  val hasPartitioner = rdds.filter(_.partitioner.exists(_.numPartitions > 0))
  if (hasPartitioner.nonEmpty) { // 有RDD存在分区器
    // 使用分区数最大的RDD的分区器
    hasPartitioner.maxBy(_.partitions.length).partitioner.get
  } else { // 没有RDD存在分区器
    // 获取spark.default.parallelism参数配置的值
    if (rdd.context.conf.contains("spark.default.parallelism")) {
      // spark.default.parallelism有配置，返回对应分区总数的HashPartitioner分区器
      new HashPartitioner(rdd.context.defaultParallelism)
    } else {
      // 否则返回HashPartitioner，分区总数为所有RDD的最大分区数
      new HashPartitioner(rdds.map(_.partitions.length).max)
    }
  }
}

比较器

当分区内部需要按照key排序时候，会指定key的排序器，如果Ordering中已经制定了，则使用指定的，没有则使用基于hashCode的排序方式。

private val keyComparator: Comparator[K] = ordering.getOrElse(
  // 当没有指定比较器时，会使用默认的按照key的哈希值进行比较的比较器
  new Comparator[K] {
    override def compare(a: K, b: K): Int = {
      val h1 = if (a == null) 0 else a.hashCode()
      val h2 = if (b == null) 0 else b.hashCode()
      if (h1 < h2) -1 else if (h1 == h2) 0 else 1
    }
  }
)

// 获取比较器
private def comparator: Option[Comparator[K]] = {
  // 只有在ordering被定义，或开启了Map端聚合时才需要比较器
  if (ordering.isDefined || aggregator.isDefined) {
    Some(keyComparator)
  } else {
    None
  }
}

写入过程

ExternalSorter是通过insertAll进行写入数据的，具体步骤如下:

指定map端聚合时候，先指定聚合函数，然后遍历每条记录，使用AppendOnlyMap作为内存缓冲区，来不断的写入数据，每次写入数据完成后，都要进行是否要溢出判断，如果内存不足，则需要溢写磁盘；
不指定map端聚合时候，则使用PartitionedPairBuffer缓冲数据，每次写入数据，也要进行溢写磁盘判断。

def insertAll(records: Iterator[Product2[K, V]]): Unit = {
  val shouldCombine = aggregator.isDefined

  if (shouldCombine) { // 如果用户指定了聚合器，那么对数据进行聚合
    // 获取聚合器的mergeValue函数，此函数用于将新的Value合并到聚合的结果中
    val mergeValue = aggregator.get.mergeValue
    // 获取聚合器的createCombiner函数，此函数用于创建聚合的初始值。
    val createCombiner = aggregator.get.createCombiner
    var kv: Product2[K, V] = null
    // 定义偏函数，当有新的Value时，调用mergeValue函数将新的Value合并到之前聚合的结果中，
    // 否则说明刚刚开始聚合，此时调用createCombiner函数以Value作为聚合的初始值
    val update = (hadValue: Boolean, oldValue: C) => {
      if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
    }

    // 迭代输入的记录
    while (records.hasNext) {
      // 增加已经读取的元素数
      addElementsRead()
      kv = records.next()
      
      // 计算分区索引（ID），将分区索引与key、update偏函数作为参数对由分区索引与key组成的对偶进行聚合
      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()
      // 计算分区索引ID，调用PartitionedPairBuffer的insert()方法
      buffer.insert(getPartition(kv._1), kv._1, kv._2.asInstanceOf[C])
      
      // 进行可能的磁盘溢出
      maybeSpillCollection(usingMap = false)
    }
  }
}

内存写入

PartitionedAppendOnlyMap和PartitionedPairBuffer是内存中缓冲数据的数据结构，前面已经详细介绍过了，可以自行回顾下，主要是每次读入到数据后，进行内存缓存，每次缓存后都会考虑一下是否需要溢写到磁盘。

// 当设置了聚合器（Aggregator）时，Map端将中间结果溢出到磁盘前， 先利用此数据结构在内存中对中间结果进行聚合处理。
@volatile private var map = new PartitionedAppendOnlyMap[K, C]
// 当没有设置聚合器（Aggregator）时，Map端将中间结果溢出到磁盘前，先利用此数据结构将中间结果存储在内存中
@volatile private var buffer = new PartitionedPairBuffer[K, C]

内存申请&溢写磁盘

上面说到每次写入一条数据后都要进行判断是否要溢写磁盘，是通过maybeSpillCollection进行的，先估算占用的内存大小，然后判断是否内存不足，这个时候我们找到了内存交互的地方。

// 用于判断何时需要将内存中的数据写入磁盘
private def maybeSpillCollection(usingMap: Boolean): Unit = {
  var estimatedSize = 0L
  if (usingMap) { // 如果使用了PartitionedAppendOnlyMap
    // 对PartitionedAppendOnlyMap的大小进行估算；通过MemoryConsumer进行内存申请
    estimatedSize = map.estimateSize()
    if (maybeSpill(map, estimatedSize)) { // 将PartitionedAppendOnlyMap中的数据溢出到磁盘，申请的memory不够了，需要spill到磁盘并进行重建
      // 重新创建PartitionedAppendOnlyMap
      map = new PartitionedAppendOnlyMap[K, C]
    }
  } else { // 如果使用了PartitionedPairBuffer
    // 对PartitionedPairBuffer的大小进行估算
    estimatedSize = buffer.estimateSize()
    if (maybeSpill(buffer, estimatedSize)) { // 将PartitionedPairBuffer中的数据溢出到磁盘
      // 重新创建PartitionedPairBuffer
      buffer = new PartitionedPairBuffer[K, C]
    }
  }

  if (estimatedSize > _peakMemoryUsedBytes) { // 更新ExternalSorter已经使用的内存大小的峰值
    _peakMemoryUsedBytes = estimatedSize
  }
}

内存申请

既然使用到了内存缓冲，那么使用前就要先申请内存，不然就会有OOM风险，ExternalSorter继承了Spillable接口，Spillable继承了MemoryConsumer，前面我们说到了TaskMemoryManager接受MemoryConsumer的内存申请，统一由ExecutionMemoryPool来提供执行内存，当内存不足时候Spillable提供将内存中内容写入到磁盘的操作，所以作为MemoryConsumer可以向TaskMemoryManager申请内存，不足时候会进行溢写磁盘操作。让我们来看下是如何判断溢写的:

如果当前集合已经读取的元素数量是32的倍数，且集合当前的内存大小大于等于myMemoryThreshold，就会尝试获取内存；
尝试申请内存，申请量为2 * currentMemory - myMemoryThreshold；
更新申请到的内存大小；
判断是否要进行溢写操作，条件是写入的条数大于numElementsForceSpillThreshold或者没申请足够的内存；
如果要溢写，则调用子类的具体写入磁盘逻辑，然后释放内存，更新读入数据计数。

// 用于检测存放数据的集合是否需要进行溢写，若需要溢写，会调用spill()方法将集合中的数据溢出到磁盘。 返回值表示是否需要溢写。
protected def maybeSpill(collection: C, currentMemory: Long): Boolean = {
  var shouldSpill = false
  // 如果当前集合已经读取的元素数量是32的倍数，且集合当前的内存大小大于等于myMemoryThreshold
  if (elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold) {
    // 计算尝试获取的内存大小
    val amountToRequest = 2 * currentMemory - myMemoryThreshold
    // 从内存管理器里面获取内存 【为当前任务尝试获取期望大小的内存，得到实际获得的大小】，与Memory交互的地方终于找到了
    val granted = acquireMemory(amountToRequest)
    // 更新已经获得的内存大小
    myMemoryThreshold += granted
    // 判断是否需要溢出，如果没有获取到Memory，这个地方要进行spill操作
    shouldSpill = currentMemory >= myMemoryThreshold
  }
  // 如果_elementsRead大于numElementsForceSpillThreshold也代表需要溢出
  shouldSpill = shouldSpill || _elementsRead > numElementsForceSpillThreshold
  // Actually spill
  if (shouldSpill) { // 如果应该进行溢出
    // 溢写操作计数自增
    _spillCount += 1
    logSpillage(currentMemory)
    // 将集合中的数据溢出到磁盘，该方法是抽象方法，需要子类实现
    spill(collection)
    // 已读取的元素计数归零
    _elementsRead = 0
    // 更新已经溢出的内存大小
    _memoryBytesSpilled += currentMemory
    // 释放ExternalSorter占用的内存
    releaseMemory()
  }
  // 返回是否进行了溢出
  shouldSpill
}

磁盘溢写

溢写磁盘是在ExternalSorter中实现的，写入磁盘中要按照分区排序，指定了key排序器时候分区内按照key排序，首先将WritablePartitionedIterator迭代器中的数据溢写到磁盘，Iterator是遍历排序后的数组，PartitionedAppendOnlyMap的排序规则是<partitionId, key>, PartitionedPairBuffer则是partitionId，然后通过spillMemoryIteratorToDisk写入到磁盘中。

override protected[this] def spill(collection: WritablePartitionedPairCollection[K, C]): Unit = {
  // 获取WritablePartitionedIterator，传入的是比较器由comparator方法提供。
  // 如果定义了ordering或aggregator，那么比较器就是keyComparator，否则没有比较器。 这里会影响后面的排序过程：
  // - 如果指定了ordering或aggregator，此处传入keyComparator，后面的[[PartitionedAppendOnlyMap]]不仅会根据分区ID进行排序，还会根据keyComparator对键进行排序；
  // - 如果没有指定ordering或aggregator，此处传入None，后面的[[PartitionedPairBuffer]]则只会根据分区ID进行排序。
  val inMemoryIterator = collection.destructiveSortedWritablePartitionedIterator(comparator)
  // 将集合中的数据溢出到磁盘
  val spillFile = spillMemoryIteratorToDisk(inMemoryIterator)
  // 将溢出生成的文件添加到spills中
  spills += spillFile
}

我们来先看下溢写到磁盘中的数据结构，可以看出来它是比较规整的，包含文件名字，文件对应数据块，以及每个批次的大小，每个分区的文件大小等，这些都是为了后续读取时候方便的获取元信息。

private[this] case class SpilledFile(
  file: File, // 溢写的文件
  blockId: BlockId, // 对应的数据块的BlockId
  serializerBatchSizes: Array[Long], // 每个批次的数据大小，默认是10000条数据一个批次，最后一个批次可能小于10000条
  elementsPerPartition: Array[Long]) // 每个分区的元素数量

最后我们来看下是如何将排序好的迭代器数据写入到磁盘文件中的:

首先创建一个临时文件;
然后构造磁盘写入器DiskBlockObjectWriter，根据块id，文件名字，序列器，缓冲区大小来进行构建；
建立每个批次大小的记录数组，最后赋给SpilledFile属性，一般除了最后一个长度有可能不是10000外，其余都是10000；
建立每个partition长度的数组，写入时候给它记录最后赋给SpilledFile属性；
遍历按照分区或者分区+key排序好的数据迭代器，先获取相应的partitionId，对应的分区计数加1，写入量加1，等到写满10000条时候，主动flush数据到磁盘；
最后如果还有数据没有flush到磁盘，进行flush；
根据每次flush的条数，每个分区的大小以及文件名字构建SpilledFile返回。

private[this] def spillMemoryIteratorToDisk(inMemoryIterator: WritablePartitionedIterator): SpilledFile = {
  // 创建唯一的BlockId和文件
  val (blockId, file) = diskBlockManager.createTempShuffleBlock()

  // objectsWritten用于统计已经写入磁盘的键值对数量
  var objectsWritten: Long = 0
  // 用于对Shuffle中间结果写入到磁盘的度量与统计。
  val spillMetrics: ShuffleWriteMetrics = new ShuffleWriteMetrics
  // 获取DiskBlockObjectWriter
  val writer: DiskBlockObjectWriter =
    blockManager.getDiskWriter(blockId, file, serInstance, fileBufferSize, spillMetrics)

  // 创建存储批次大小的数组缓冲batchSizes
  val batchSizes = new ArrayBuffer[Long]
  // 创建存储每个分区有多少个元素的数组缓冲
  val elementsPerPartition = new Array[Long](numPartitions)

  def flush(): Unit = {
    // 将DiskBlockObjectWriter的输出流中的数据真正写入到磁盘
    val segment = writer.commitAndGet()
    // 将本批次写入的文件长度添加到batchSizes和_diskBytesSpilled中
    batchSizes += segment.length
    _diskBytesSpilled += segment.length
    // 将objectsWritten清零，以便下一批次的操作
    objectsWritten = 0
  }

  var success = false
  try {
    // 对WritablePartitionedIterator进行迭代
    while (inMemoryIterator.hasNext) {
      val partitionId = inMemoryIterator.nextPartition()  // 获取数据的分区ID
      inMemoryIterator.writeNext(writer)     // 将键值对写入磁盘
      // 将elementsPerPartition中统计的分区对应的元素数量加1
      elementsPerPartition(partitionId) += 1
      // 将objectsWritten加一
      objectsWritten += 1

      // 每写10000条记录进行一次刷盘
      if (objectsWritten == serializerBatchSize) {
        // 将DiskBlockObjectWriter的输出流中的数据真正写入到磁盘
        flush()
      }
    }
    // 遍历记录完毕，如果还有剩余未刷盘的数据，则进行刷盘
    if (objectsWritten > 0) {
      // 将DiskBlockObjectWriter的输出流中的数据真正写入到磁盘
      flush()
    } else {
      writer.revertPartialWritesAndClose()
    }
    // 标记写出成功
    success = true
  } finally {
    if (success) {
      // 关闭写出器
      writer.close()
    } else {
      writer.revertPartialWritesAndClose()
      if (file.exists()) {
        if (!file.delete()) {
          logWarning(s"Error deleting ${file}")
        }
      }
    }
  }

  // 创建并返回SpilledFile,batchSizes.toArray：记录了每个批次的数据大小
  // elementsPerPartition：记录了每个分区有多少条键值对记录
  SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition)
}

这样子我们就完成了数据的写入过程，写入的数据可能一部分在磁盘中，以SpilledFile的形式存储着，一部分在内存中，存在于PartitionedAppendOnlyMap或者PartitionedPairBuffer中，但是ShuffleMapTask需要的结果是在一个文件中，所以我们要对内存和磁盘中的文件进行合并操作。

文件合并

当insertAll插入完所有数据后，这时候磁盘和内存中都有可能有Shuffle的数据，最后我们要进行合并，根据是否排序，是否mapSideCombine来进行相应的聚合操作。

磁盘文件迭代器-SpillReader

上面我们分析了写入磁盘文件的过程，磁盘文件除了文件名字外，有两个重要的变量，如下所示：

serializerBatchSizes: Array[Long], // 每个批次的数据大小，默认是10000条数据一个批次，最后一个批次可能小于10000条
elementsPerPartition: Array[Long]) // 每个分区的元素数量

可以看出来SpilledFile存储了各个分区的长度，而且是按照分区进行排序的，所以我们可以得到每个分区开始位置在文件中的偏移量，从而我们可以获取每个分区的数据迭代器， SpillReader就是来做这个工作的，我们首先看下它的成员变量，比较简单我们看下源码及注释:

 private[this] class SpillReader(spill: SpilledFile) {
  // scanLeft函数会将当前索引位置之前的元素进行算子计算，如：
  //    val s = 1 to 10    => Range(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
  //    val ss = s.scanLeft(0)(_ + _)  => Vector(0, 1, 3, 6, 10, 15, 21, 28, 36, 45, 55)
  // 通过这种方式，可以得到每个批次的数据的偏移量
  val batchOffsets = spill.serializerBatchSizes.scanLeft(0L)(_ + _)

  var partitionId = 0 // 分区ID
  var indexInPartition = 0L // 当前分区读取的键值对条数
  var batchId = 0 // 批次ID
  var indexInBatch = 0 // 当前批次已读取的键值对条数
  var lastPartitionId = 0 // 记录最后读取的分区的ID
}

读取下一个分区

SpillReader是从SpilledFile读取数据迭代器的接口，外界使用时候通过readNextPartition获取分区的数据迭代器，构造的读取下一个分区数据的迭代器如下所示:

// 读取下一个分区的数据，返回的是所有该分区的键值对的迭代器
def readNextPartition(): Iterator[Product2[K, C]] = new Iterator[Product2[K, C]] {
  val myPartition = nextPartitionToRead   // 当前读取的分区的ID
  nextPartitionToRead += 1  // 更新nextPartitionToRead，自增1

  override def hasNext: Boolean = {
    if (nextItem == null) {
      // 读取当前读取的分区的下一条记录
      nextItem = readNextItem()
      if (nextItem == null) { // 读取为空，说明当前分区没有下一条记录了
        // 返回false
        return false
      }
    }
    assert(lastPartitionId >= myPartition)
    // 由于readNextItem()读取时可能会跳到下一个分区，如果没有跳到下一个分区，说明当前分区还有数据
    lastPartitionId == myPartition
  }

  override def next(): Product2[K, C] = {
    if (!hasNext) {
      throw new NoSuchElementException
    }
    // 将nextItem返回即可
    val item = nextItem
    nextItem = null
    item
  }
}

读取下一个k-v对

可以看出来主要是通过readNextItem来读取数据的，我们来看下这个函数如何进行的:

从序列化流中<读取的一个批次，前面讲到是以10000为一个批次的>读取key, value；
如果读取的条数达到10000阈值，就需要读取下一个批次；
判断是否要跳到下一个分区，即当前分区是否读取完了<分区号小于最大分区数，且分区内索引达到了该分区内最大元素的数量>；
如果所有分区数据都读取完了，就结束。

// 获取下一个键值对
private def readNextItem(): (K, C) = {
  // 检查是否已经完成读取，或者反序列化流是否为null
  if (finished || deserializeStream == null) {
    return null
  }
  // 读取键和值
  val k = deserializeStream.readKey().asInstanceOf[K]
  val c = deserializeStream.readValue().asInstanceOf[C]
  // 记录当前读取的分区ID
  lastPartitionId = partitionId 
  // 批次索引自增
  indexInBatch += 1

  // serializerBatchSize是每个批次刷盘的阈值，默认为10000, 即溢写键值对时，每10000条键值对就进行一次刷盘，
  // 因此如果当前批次读取的条数达到了10000阈值，则需要读取下一个批次了。
  if (indexInBatch == serializerBatchSize) {
    // 重置当前批次读取的条数记录为0
    indexInBatch = 0
    // 创建下一个批次的数据的反序列化流
    deserializeStream = nextBatchStream()
  }
  // Update the partition location of the element we're reading
  // 当前分区读取的键值对条数自增
  indexInPartition += 1
  // 判断是否要跳向下一个分区
  skipToNextPartition()
  // 如果读取的分区是最后一个分区，注意，partitionId是在skipToNextPartition()中自增的，
  // 当partitionId与numPartitions相等，说明所有分区都读完了，numPartitions - 1是最后一个分区的ID
  if (partitionId == numPartitions) {
    // 标记为true
    finished = true
    // 关闭反序列化流
    if (deserializeStream != null) {
      deserializeStream.close()
    }
  }
  // 返回本次读取的键值对
  (k, c)
}

// 分区号小于最大分区数，且分区内索引达到了该分区内最大元素的数量
private def skipToNextPartition() { 
  while (partitionId < numPartitions &&
         indexInPartition == spill.elementsPerPartition(partitionId)) {
    // 分区号自增1
    partitionId += 1
    // 重置分区内索引记录为0
    indexInPartition = 0L
  }
}

读取下一个批次数据

上面讲到SpilledFile里面会有一个变量记录各个批次的数据条数，在构造从文件中读取的数据流时候，是会进行缓存，读取一个批次大小的数据，首先获取这个批次的开始结束偏移量，然后构造BufferedInputStream流，再经过序列化器对经过压缩和加密的流封装为反序列化流，供nextItem读取数据使用。

// 反序列化流
var deserializeStream = nextBatchStream()  // Also sets fileStream

//  为下一个批次的数据构建反序列化流
def nextBatchStream(): DeserializationStream = {
  // batchOffsets数组的长度为批次数量 + 1， 可以使用batchOffsets数组的长度判断批次ID是否在范围内。
  if (batchId < batchOffsets.length - 1) {
    // 如果反序列化流不为null，说明之前读取了一个批次的数据。为了让两个批次的数据互不干扰，先关闭该反序列化流和文件输入流
    if (deserializeStream != null) {
      // 先关闭反序列化流和文件输入流
      deserializeStream.close()
      fileStream.close()
      deserializeStream = null
      fileStream = null
    }

    // batchOffsets记录的是偏移量前后两个偏移量之差就是该段的长度，一般是10000;获取起始偏移量
    val start = batchOffsets(batchId)
    // 以溢写文件构建输入流
    fileStream = new FileInputStream(spill.file)
    // 使用FileChannel定位到起始偏移量
    fileStream.getChannel.position(start)
    // 将batchId自增1
    batchId += 1

    // 获取终止偏移量
    val end = batchOffsets(batchId)

    // 将文件输入流包装为缓冲输入流
    val bufferedStream = new BufferedInputStream(ByteStreams.limit(fileStream, end - start))

    // 将缓冲输入流进行压缩和加密
    val wrappedStream = serializerManager.wrapStream(spill.blockId, bufferedStream)
    // 将经过压缩和加密的流封装为反序列化流
    serInstance.deserializeStream(wrappedStream)
  } else {
    // 没有更多的批次需要处理，执行清理操作
    cleanup()
    null
  }
}

内存数据迭代器

针对还未写入磁盘的数据，按照分区进行分组，返回分区的iterator，data是按照partition+key排序的结果迭代器，IteratorForPartition如其名字，就是各个分区的迭代器，由于内存中是按照分区排序好的，所以只需要考虑元素的分区id是否等于当前遍历的分区即可，如下所示:

// 用于对destructiveIterator方法返回的迭代器按照分区ID进行分区 
private def groupByPartition(data: Iterator[((Int, K), C)]): Iterator[(Int, Iterator[Product2[K, C]])] = {
  val buffered = data.buffered
  // 给每个分区生成了一个IteratorForPartition
  (0 until numPartitions).iterator.map(p => (p, new IteratorForPartition(p, buffered)))
}

private[this] class IteratorForPartition(partitionId: Int, data: BufferedIterator[((Int, K), C)])
extends Iterator[Product2[K, C]] {
  // 用于判断对于指定分区ID是否有下一个元素：1. data本身需要有下一个元素。  data的下一个元素对应的分区ID要与指定的分区ID一样。
  override def hasNext: Boolean = data.hasNext && data.head._1._1 == partitionId

  // 获取下一个元素
  override def next(): Product2[K, C] = {
    if (!hasNext) {
      throw new NoSuchElementException
    }
    // 直接从data中获取
    val elem = data.next()
    // 返回键值对
    (elem._1._2, elem._2)
  }
}

磁盘+内存数据合并

writePartitionedFile是文件合并的入口，是由SortShuffleWriter调用，生成最终的临时数据文件，同时返回每个分区的数据大小数组。

如果只有内存中有数据，则获取内存基于分区或者分区+key的排序迭代器，不断写入临时数据文件即可；
如果内存和磁盘中都有文件，需要将其进行合并，对于每个分区返回分区的迭代器，然后各个分区依次写入到文件中。

// 前ExternalSorter中的添加数据全部写出到磁盘的文件，这个方法主要被SortShuffleWriter调用。
def writePartitionedFile(blockId: BlockId, utputFile: File): Array[Long] = {
  // 创建对每个分区的长度进行跟踪的数组lengths
  val lengths = new Array[Long](numPartitions)
  // 获取DiskBlockObjectWriter
  val writer = blockManager.getDiskWriter(
    blockId, outputFile, serInstance, fileBufferSize, context.taskMetrics().shuffleWriteMetrics)

  if (spills.isEmpty) { // 没有缓存溢出到磁盘的文件，即所有的数据依然都在内存中
    val collection = if (aggregator.isDefined) map else buffer
    val it = collection.destructiveSortedWritablePartitionedIterator(comparator)
    while (it.hasNext) { // 将底层data数组中的数据按照分区ID分别写入到磁盘中
      val partitionId = it.nextPartition()
      while (it.hasNext && it.nextPartition() == partitionId) {
        it.writeNext(writer)
      }
      val segment = writer.commitAndGet()
      // 将分区的数据长度更新到lengths数组中
      lengths(partitionId) = segment.length
    }
  } else { // 如果spills中缓存了溢出到磁盘的文件，即有些数据在内存中，有些数据已经溢出到了磁盘上
    // 使用partitionedIterator()方法对磁盘和内存中的数据进行合并后，将各个元素写到磁盘
    for ((id, elements) <- this.partitionedIterator) {
      if (elements.hasNext) {
        for (elem <- elements) {
          writer.write(elem._1, elem._2)
        }
        val segment = writer.commitAndGet()
        // 将各个分区的数据长度更新到lengths数组中
        lengths(id) = segment.length
      }
    }
  }
  ...

  // 返回lengths数组
  lengths
}

分区数据迭代器获取

我们来看下是如何生成各个分区数据的迭代器的，partitionedIterator函数返回(Int, Iterator[Product2[K, C]])的迭代器，key表示的是分区id，Iterator[Product2[K, C]表示的是该分区的按照排序规则或者聚合+排序规则得到的数据迭代器。

如果只有内存中有数据，可以通过上面讲过的内存数据迭代器来获取分区数据迭代器；
如果内存和磁盘都有数据，就要进行merge操作。

def partitionedIterator: Iterator[(Int, Iterator[Product2[K, C]])] = { // partitionId, <<k,C>,....>
  val usingMap = aggregator.isDefined
  // 获取当前使用的数据结构
  val collection: WritablePartitionedPairCollection[K, C] = if (usingMap) map else buffer
  // 判断spills是否为空，如果为空表示还没有溢写到磁盘文件
  if (spills.isEmpty) { // 没有溢出到磁盘的文件，即所有的数据依然都在内存中，直接返回内存数据的迭代器即可
    if (!ordering.isDefined) { // 对底层data数组中的数据只按照分区ID排序
      // 对分区进行分组， 对底层data数组中的数据按照分区ID排序
      groupByPartition(destructiveIterator(
        collection.partitionedDestructiveSortedIterator(None)))
    } else { // 对底层data数组中的数据按照分区ID和key排序
      // 对分区进行分组，对底层data数组中的额数据按照分区ID和key排序
      groupByPartition(destructiveIterator(
        collection.partitionedDestructiveSortedIterator(Some(keyComparator))))
    }
  } else { // 如果spills中缓存了溢出到磁盘的文件，即有些数据在内存中，有些数据已经溢出到了磁盘上
    // 将溢出的磁盘文件和data数组中的数据合并
    merge(spills, destructiveIterator(
            // 按照分区ID和key排序
            collection.partitionedDestructiveSortedIterator(comparator))
     )
  }
}

磁盘+内存合并

merge进行数据的合并操作，会构造磁盘数据缓冲器SpillReader以及内存缓冲器，然后对于每个分区进行遍历，将内存中的数据提取当前需要分区的迭代器与磁盘中各个文件改分区的迭代器进行合并，然后根据是否聚合来分别调用基于聚合的merge逻辑还是简单排序的merge逻辑。

private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)]): Iterator[(Int, Iterator[Product2[K, C]])] = {
  // 为每个SpilledFile创建SpillReader读取器
  val readers = spills.map(new SpillReader(_))
  // 创建缓冲迭代器，该迭代器扩展了一个功能方法head()，即可以查看迭代器中的下一个元素，但不会将它移出
  val inMemBuffered = inMemory.buffered
  // 遍历所有分区的ID，顺序遍历
  (0 until numPartitions).iterator.map { p =>
    // 为当前分区创建一个迭代器
    val inMemIterator = new IteratorForPartition(p, inMemBuffered)
    // 使用SpillReader顺序读取包含了每个分区的数据的迭代器，并与inMemIterator迭代器合并
    val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
    // 判断是否需要聚合
    if (aggregator.isDefined) { // 需要聚合
      // 使用mergeWithAggregation()方法进行聚合，返回元素类型为 (分区ID, 对应的聚合数据的迭代器) 的迭代器
      (p, mergeWithAggregation(
        iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
    } else if (ordering.isDefined) { // 需要排序
      // 使用mergeSort()方法进行归并排序，返回元素类型为 (分区ID, 对应的有序数据的迭代器) 的迭代器
      (p, mergeSort(iterators, ordering.get))
    } else {
      // 不需要聚合，也不需要排序，直接合并溢写文件中对应分区的数据
      (p, iterators.iterator.flatten)
    }
  }
}

不带聚合的合并

如果不需要聚合操作，由于各个文件都是按照分区排好序的，所以就变成了对多个排序好的文件的归并排序，Spark中使用优先级队列来进行排序，先将当前分区所有的元素首元素都加入到优先级队列中<小顶堆>，迭代器取数据的方式就变成了，先从队列中取出最小元素<堆第一个元素>，然后返回，如果该文件数据迭代器还有元素，就在加入到堆中，直到读取到所有文件中都没有该分区的数据表示该分区数据已经结束了。

// 使用比较器对迭代器内的键值对数据进行归并排序
private def mergeSort(iterators: Seq[Iterator[Product2[K, C]]], comparator: Comparator[K]): Iterator[Product2[K, C]] = {
   // 先过滤还有元素的迭代器,再将每个迭代器转换为缓冲迭代器
  val bufferedIters = iterators.filter(_.hasNext).map(_.buffered)
  type Iter = BufferedIterator[Product2[K, C]]
  
  // 构造一个优先队列，作为小顶堆结构 
  val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
    // 比较方法，使用比较器比较每个两个迭代器的下一个元素（即键值对）的键
    override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
  })
  // 将bufferedIters添加到heap队列，利用迭代器的第一个元素进行排序
  heap.enqueue(bufferedIters: _*)
  // 返回一个新的迭代器，其中的元素都存放在heap优先队列中
  new Iterator[Product2[K, C]] {
    // 是否还有下一个元素
    override def hasNext: Boolean = !heap.isEmpty

    // 获取下一个元素
    override def next(): Product2[K, C] = {
      if (!hasNext) {
        throw new NoSuchElementException
      }
      // heap出队一个迭代器
      val firstBuf = heap.dequeue()
      // 使用该迭代器获取一个键值对
      val firstPair = firstBuf.next()
      if (firstBuf.hasNext) { // 如果该迭代器还有元素
        // 将该迭代器再次放入heap堆中，这样下次还可以从该迭代器获取键值对
        heap.enqueue(firstBuf)
      }
      // 返回迭代到的键值对
      firstPair
    }
  }
}

带聚合的合并

如果需要mapSideCombine，则要根据是否规定了按照key排序来进行不同的处理，这是因为如果规定了key的排序规则，则相同key的数据就会在一起，而没有规定时候是默认按照key的hashcode进行排序的，同一个key的数据可能不连续，但是我们要选择出来这些相同key的进行聚合，所以需要记录相同hashcode的各个key，分别进行聚合操作。

private def mergeWithAggregation(iterators: Seq[Iterator[Product2[K, C]]], 
                                 mergeCombiners: (C, C) => C, 
                                 comparator: Comparator[K], totalOrder: Boolean)
: Iterator[Product2[K, C]] = {
  if (!totalOrder) { // 不是全排序,没有指定key排序器，默认使用hashCode进行排序，这样子会有多个key具有相同的大小顺序，造成同一个key的数据不连续
    new Iterator[Iterator[Product2[K, C]]] {
      // 先使用mergeSort()方法进行合并排序，该方法迭代的键值对会按照键进行排序
      val sorted = mergeSort(iterators, comparator).buffered
      
      val keys = new ArrayBuffer[K]  // 装载键的数组
      val combiners = new ArrayBuffer[C] // 聚合的值的数组

      // 需要看sorted中是否还有下一个
      override def hasNext: Boolean = sorted.hasNext

      override def next(): Iterator[Product2[K, C]] = {
        if (!hasNext) {
          throw new NoSuchElementException
        }
        // 清空键数组和聚合值数组
        keys.clear()
        combiners.clear()
        val firstPair = sorted.next() // 获得下一个键值对
        keys += firstPair._1  // 将键存入keys
        combiners += firstPair._2  // 将值存入combiners
        val key = firstPair._1   // 记录当前迭代到的键

        // 满足下面两个条件，就循环： 1. sorted还有下一个键值对；2. 比较下一个键值对的键与当前key记录的键是否相等。
        // 如果1和2都满足，说明下一个键值对的键与当前键值对的键是一样的
        while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
          // 获取下一个键值对pair
          val pair = sorted.next()
          var i = 0
          var foundKey = false
          // 从keys数组中找到对应的键
          while (i < keys.size && !foundKey) {
            if (keys(i) == pair._1) { // 找到相同的键了
              // 取出聚合数组中对应的旧的聚合值，将旧的聚合值与当前pair的值使用mergeCombiners函数进行聚合， 将新值存入combiners
              combiners(i) = mergeCombiners(combiners(i), pair._2)
              // 找到对应的键，说明该键之前已经迭代到了
              foundKey = true
            }
            i += 1
          }
          // 没有找到对应的键，则将键和值分别存入keys和combiners
          if (!foundKey) {
            keys += pair._1
            combiners += pair._2
          }
        }

        // 将键和聚合的值进行zip，获取新的键值对
        keys.iterator.zip(combiners.iterator)
      }
    }.flatMap(i => i)
  } else { // 全排序,key相同的数据是连续的
    new Iterator[Product2[K, C]] {
      // 先使用mergeSort()方法进行合并排序，该方法迭代的键值对会按照键进行排序
      val sorted = mergeSort(iterators, comparator).buffered

      // 需要看sorted中是否还有下一个
      override def hasNext: Boolean = sorted.hasNext

      override def next(): Product2[K, C] = {
        if (!hasNext) {
          throw new NoSuchElementException
        }
        val elem = sorted.next()  // 获取sorted中的下一个键值对
        val k = elem._1 // 取出键
        var c = elem._2 // 取出值
        // 满足下面两个条件，就循环：1. sorted还有下一个键值对； 2. 比较下一个键值对的键与当前k记录的键是否相等。
        // 如果1和2都满足，说明下一个键值对的键与当前键值对的键是一样的
        while (sorted.hasNext && sorted.head._1 == k) {
          val pair = sorted.next()  // 获取下一个键值对
          // 将下一个键值对的值和当前值使用mergeCombiners()方法进行聚合，聚合的值赋值给c
          c = mergeCombiners(c, pair._2)
        }
        // 返回键和聚合后的值
        (k, c)
      }
    }
  }
}

总结

最后我们来总结下，ShuffleMapTask的整体过程：

首先，数据会被一条一条地向内部的map/buffer结构中插入；
每插入一条数据都会检查内存占用情况，如果内存占用超过阈值，并且申请不到足够的执行内存，就会将目前内存中的数据溢写到磁盘；
对于溢写的过程：首先会将数据按照分区和key进行排序，相同分区的数据排在一起，然后根据提供的排序器按照key的顺序排；然后通过BlockManager获取DiskBlockWriter将数据写入磁盘形成一个文件，并记录溢写的文件信息；在整个写入过程中，会溢写多个文件；
然后会对多个文件和内存中数据进行合并，将这些溢写的小文件和最后内存中剩下的数据进行归并排序，然后写入一个大文件中，并且在写入的过程中记录每个分区数据在文件中的位移；
最后写到一个数据文件和一个索引文件中，上报MapStatus给Driver，索引文件即记录了每个reduce端分区在数据文件中的位移，这样reduce在拉取数据的时候才能很快定位到自己分区所需要的数据。

参考

小白数据猿

关注

6
点赞
踩
3

收藏

觉得还不错? 一键收藏
0
评论
Spark Shuffle源码分析系列之SortShuffleWriter

上一节我们分析了BypassMergeSortShuffleWriter，它是Hash风格的ShuffleWriter，主要适用于不需要map-side聚合排序，并且分区数目小于spark.shuffle.sort.bypassMergeThreshold<默认200>；本节我们来介绍SortBased-ShuffleWriter，它可以处理map端的聚合排序操作，是默认的ShuffleWriter。ShuffleMapTask整体流程Spark中负责具体的Shuffle Map端执行任.
复制链接

扫一扫