Spark Shuffle源码分析系列之ShuffleReader 二

上一节我们分析了ShuffleReadTask是如何获取数据位置，得到各个块数据的，本节我们继续分析Spark对数据是如何存储、聚合、排序，然后供后续算子使用的。

概述

上一节我们讲到了ShuffleReduce任务获取数据的迭代器，那么获取来的数据存放到哪里？刚fetch来的FileSegment存放在缓冲区，经过解压缩和反序列化，如果shuffle算子规定了需要聚合操作，则要使用ExternalAppendOnlyMap进行数据聚合，优先在内存中使用SizeTrackerAppendOnlyMap进行聚合，如果内存空间不足时，ExternalAppendOnlyMap可以将内存中数据进行排序后spill到磁盘上，等到需要它们的时候再进行归并。使用“内存＋磁盘”的一个主要问题就是如何在两者之间取得平衡？在 Hadoop MapReduce 中，默认将reducer的70%的内存空间用于存放shuffle来的数据，等到这个空间利用率达到66%的时候就开始merge-combine-spill。在Spark中，也适用同样的策略，一旦ExternalAppendOnlyMap达到一个阈值就开始 spill；如果涉及到排序，则会使用ExternalSorter对聚合后的数据进行排序操作。

数据聚合

经过了数据的读取得到每个block的输入流，然后解压缩，反序列化操作，如果算子规定了聚合操作，则需要对各个map的结果按照key进行聚合操作，根据是否map端聚合操作，选用Aggregator中不同的方法，如下所示:

// 处理可能存在的聚合迭代
val aggregatedIter: Iterator[Product2[K, C]] = if (dep.aggregator.isDefined) {
  if (dep.mapSideCombine) {
    // 如果指定了聚合函数且允许在map端进行合并，在reduce端对数据进行聚合
    val combinedKeyValuesIterator = interruptibleIter.asInstanceOf[Iterator[(K, C)]]
    // 使用聚合器进行聚合
    dep.aggregator.get.combineCombinersByKey(combinedKeyValuesIterator, context)
  } else {
    // 如果指定了聚合函数，但不允许在map端进行合并，在reduce端对数据进行缓存
    val keyValuesIterator = interruptibleIter.asInstanceOf[Iterator[(K, Nothing)]]
    // 使用聚合器进行聚合
    dep.aggregator.get.combineValuesByKey(keyValuesIterator, context)
  }
} else {
  // 没有指定聚合函数，那么不作任何处理
  require(!dep.mapSideCombine, "Map-side combine without Aggregator specified!")
  interruptibleIter.asInstanceOf[Iterator[Product2[K, C]]]
}

我们来看下Aggregator的两个函数的实现，两者都使用ExternalAppendOnlyMap来进行数据的内存+磁盘的存储，不同的是:

  1. map端聚合了数据则得到的输入流迭代器数据类型为Product2[K, C]，则不需要创建createCombiner来初始化聚合的初始值；
  2. map端没有聚合，则得到的输入流迭代器数据类型为Product2[K, V]，由于初始的数据类型的V，所以需要用createCombiner来初始化聚合的初始值V => C。
// 调用此方法表示数据没有在Map端聚合，此时需要createCombiner处理初值
def combineValuesByKey(iter: Iterator[_ <: Product2[K, V]], context: TaskContext): Iterator[(K, C)] = {
  // 创建ExternalAppendOnlyMap，附带有createCombiner
  val combiners = new ExternalAppendOnlyMap[K, V, C](createCombiner, mergeValue, mergeCombiners)
  combiners.insertAll(iter)
  updateMetrics(context, combiners)
  combiners.iterator // 返回聚合后的迭代器
}

// 调用此方法表示处理的是已经通过Map端聚合的数据，不需要使用createCombiner再次进行初始化处理
def combineCombinersByKey(iter: Iterator[_ <: Product2[K, C]], context: TaskContext): Iterator[(K, C)] = {
  // 创建ExternalAppendOnlyMap，createCombiner传的是identity
  val combiners = new ExternalAppendOnlyMap[K, C, C](identity, mergeCombiners, mergeCombiners)
  combiners.insertAll(iter)
  updateMetrics(context, combiners)
  combiners.iterator // 返回聚合后的迭代器
}

ExternalAppendOnlyMap

ExternalAppendOnlyMap是Shuffle Reduce任务对于需要聚合数据的shuffle算子的聚合操作使用的数据结构，包含了以下变量:

  1. createCombiner通过初始值产生聚合初始值的函数，如果map端没有聚合，则得到的数据迭代器是[K, V]，reduce端需要聚合成为[K, C]，所以需要初始化C操作；
  2. mergeValue也适用于map端没有聚合操作的情况，将[K, V]数据和[K, C]数据聚合得到结果；
  3. mergeCombiners适用于map端聚合的情况，聚合后生成[K, C]，不同task相同的key还要再次聚合即(C, C) => C；
class ExternalAppendOnlyMap[K, V, C](
    createCombiner: V => C,
    mergeValue: (C, V) => C,
    mergeCombiners: (C, C) => C,
    serializer: Serializer = SparkEnv.get.serializer,
    blockManager: BlockManager = SparkEnv.get.blockManager,
    context: TaskContext = TaskContext.get(),
    serializerManager: SerializerManager = SparkEnv.get.serializerManager)
  extends Spillable[SizeTracker](context.taskMemoryManager())
  with Serializable
  with Logging
  with Iterable[(K, C)] {
 		...
}

另外可以看到ExternalAppendOnlyMap继承了Spillable，是个MemoryConsumer，会申请执行内存，内存不足时候会溢写磁盘操作，另外实现了Iterable接口，所以需要提供iterator供外界迭代数据使用。

数据写入过程

insertAll是将从kv迭代器的数据写入到内存+磁盘中的入口函数，可以看出来和ShuffleMapTask相同，会遍历数据，先估算当前占用的数据大小，检查是否要溢写磁盘，如果需要溢写就溢写，不然聚合数据，然后对元素数目加1即可，流程比较简单，我们接下来分析内存和磁盘的数据写入过程。

def insertAll(entries: Iterator[Product2[K, V]]): Unit = {
  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() // 获取记录
    // 估算当前currentMap的大小
    val estimatedSize = currentMap.estimateSize()
    // 如果currentMap的大小大于之前记录的内存使用峰值，则更新内存使用峰值
    if (estimatedSize > _peakMemoryUsedBytes) {
      _peakMemoryUsedBytes = estimatedSize
    }
    // 检查是否需要溢写，该方法返回值表示是否发生了溢写
    if (maybeSpill(currentMap, estimatedSize)) {
      // 如果发生溢写，则重新构建一个新的SizeTrackingAppendOnlyMap赋值给currentMap
      currentMap = new SizeTrackingAppendOnlyMap[K, C]
    }
    // 进行数据聚合
    currentMap.changeValue(curEntry._1, update)
    // 更新已插入的键值对计数
    addElementsRead()
  }
}

内存存储

ExternalAppendOnlyMap内存使用SizeTrackingAppendOnlyMap这个带有大小估计的AppendOnlyMap来存储数据，主要内部使用数组模拟hashMap来存储数据，具体的分析可以参考AppendOnlyMap和PartitionedAppendOnlyMap，不在细说。

// 当前用于存储数据的Map，底层其实还是使用了AppendOnlyMap
@volatile private var currentMap = new SizeTrackingAppendOnlyMap[K, C]

主动溢写磁盘

每次添加一个元素都要使用maybeSpill来检查是否需要溢写磁盘，是在Spillable里面实现的，在ShuffleMapTask中分析过了，可以参考，主要工作是申请内存，如果能申请到，扩充map大小，不能申请到就溢写磁盘，ExternalAppendOnlyMap实现了spill逻辑，我们来看下溢写磁盘过程:

  1. 首先有两个关于写磁盘用到的参数spark.shuffle.spill.batchSize指定了一个批次写入的数据条数，默认是10000条，spark.shuffle.file.buffer是写入磁盘时候的缓冲区大小，默认是32k；还有一个变量spilledMaps记录溢写磁盘的文件迭代器；
  2. 首先对内存数据进行排序，排序大概过程是先对map中的数据移动到数组最左端，然后使用TimSort进行排序，返回排好序的迭代器；
  3. 通过spillMemoryIteratorToDisk写入迭代器中数据到磁盘中；
  4. 将磁盘文件对应的迭代器添加到spilledMaps中。
// 溢写过程中序列化器读写操作的批次大小
private val serializerBatchSize = sparkConf.getLong("spark.shuffle.spill.batchSize", 10000)
// 文件缓冲区大小
private val fileBufferSize = sparkConf.getSizeAsKb("spark.shuffle.file.buffer", "32k").toInt * 1024
// 溢写的Map
private val spilledMaps = new ArrayBuffer[DiskMapIterator]

// spill接口实现
override protected[this] def spill(collection: SizeTracker): Unit = {
  // 获取对currentMap中的数据使用keyComparator比较器进行排序后的键值对迭代器
  val inMemoryIterator = currentMap.destructiveSortedIterator(keyComparator)
  // 使用spillMemoryIteratorToDisk()方法进行溢写
  val diskMapIterator = spillMemoryIteratorToDisk(inMemoryIterator)
  // 将溢写后得到的迭代器保存到spilledMaps数组中
  spilledMaps += diskMapIterator
}

spillMemoryIteratorToDisk

spillMemoryIteratorToDisk是将内存中排好序的数据写入到磁盘中，步骤如下:

  1. 首先创建一个临时文件;
  2. 然后构造磁盘写入器DiskBlockObjectWriter，根据块id，文件名字，序列器，缓冲区大小来进行构建；
  3. 建立每个批次大小的记录数组，最后赋给DiskMapIterator属性，一般除了最后一个长度有可能不是10000外，其余都是10000；
  4. 遍历按照key排序好的数据迭代器，写入量加1，等到写满10000条时候，主动flush数据到磁盘；
  5. 最后如果还有数据没有flush到磁盘，进行flush；
  6. 根据每次flush的条数，文件名字以及块id构建DiskMapIterator返回。
private[this] def spillMemoryIteratorToDisk(inMemoryIterator: Iterator[(K, C)]): DiskMapIterator = {
  // 创建临时文件，命名为"temp_local_"前缀加上UUID字符串
  val (blockId, file) = diskBlockManager.createTempLocalBlock() 
  // 获取临时文件的DiskBlockObjectWriter
  val writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize, writeMetrics)
  // 记录每个批次溢写的键值对数量
  var objectsWritten = 0
  // 记录溢写操作中每个批次的数据的字节大小
  val batchSizes = new ArrayBuffer[Long]
 
  // 刷盘操作
  def flush(): Unit = { 
    val segment = writer.commitAndGet() // DiskBlockObjectWriter进行刷盘，返回FileSegment对象
    batchSizes += segment.length // 记录该批次刷盘后的数据大小
    _diskBytesSpilled += segment.length // 更新已溢写数据的总大小
    objectsWritten = 0  // 重置objectsWritten计数
  } 

  var success = false
  try { 
    while (inMemoryIterator.hasNext) { // 迭代键值对记录
      val kv = inMemoryIterator.next() 
      writer.write(kv._1, kv._2)  // 使用DiskBlockObjectWriter向临时文件写入键值对 
      objectsWritten += 1 // 维护objectsWritten计数

      // 如果写入的键值对数量达到了批次刷盘阈值serializerBatchSize，则进行刷盘
      // serializerBatchSize由spark.shuffle.spill.batchSize参数配置，默认为10000， 也即是10000条键值对刷一次盘
      if (objectsWritten == serializerBatchSize) {
        flush()
      }
    }

    // 所有键值对记录都迭代完了，检查是否还有未刷盘的数据
    if (objectsWritten > 0) {
      flush() // 进行刷盘，并关闭DiskBlockObjectWriter
      writer.close()
    } else { // objectsWritten为0
      // 否则说明最后一次没有任何写出，那么放弃最后一次写出并关闭DiskBlockObjectWriter
      writer.revertPartialWritesAndClose()
    }
    // 标记写出成功
    success = true
  } finally {
    if (!success) { // 写出不成功 
      // 放弃最后一次写出，关闭DiskBlockObjectWriter，并删除临时文件
      writer.revertPartialWritesAndClose()
      if (file.exists()) {
        if (!file.delete()) {
          logWarning(s"Error deleting ${file}")
        }
      }
    }
  }

  // 将溢写的文件、BlockId及记录了每个批次数据大小的数组封装为DiskMapIterator对象返回
  new DiskMapIterator(file, blockId, batchSizes)
}

DiskMapIterator

DiskMapIterator是从一个文件中，读取数据的迭代器，有以下几个变量:

  1. batchOffsets，根据每个批次大小数组计算得到的偏移量，用于每个批次的读取数据段；
  2. batchIndex记录当前读取到批次索引；
  3. fileStream是文件的输入流；
  4. deserializeStream是对文件输入流进行反序列化
// 对批次偏移量进行scanLeft操作
// 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)
// 通过这种方式，可以得到每个批次的数据的偏移量
private val batchOffsets = batchSizes.scanLeft(0L)(_ + _)  
private var batchIndex = 0  // 记录批次的索引
// 文件输入流
private var fileStream: FileInputStream = null

// 获取下一个批次数据的反序列化流
private var deserializeStream = nextBatchStream()
private var nextItem: (K, C) = null
private var objectsRead = 0
nextBatchStream

nextBatchStream用于获取下一个批次的数据，先找到该批次的在文件中的开始结束位置，然后构建BufferedInputStream文件读取流，并使用SerializerManager进行解压缩，最后反序列化得到反序列化流，供next使用读取下一个元素.

private def nextBatchStream(): DeserializationStream = {
  // 批次计数不能超过总批次数量
  if (batchIndex < batchOffsets.length - 1) {
    // 将读取上一个批次数据的反序列化流关闭
    if (deserializeStream != null) {
      deserializeStream.close()
      fileStream.close()
      deserializeStream = null
      fileStream = null
    }

    // 获取当前批次数据的起始偏移量
    val start = batchOffsets(batchIndex)
    // 获取文件输入流
    fileStream = new FileInputStream(file)
    fileStream.getChannel.position(start)  // 使用FileChannel定位到起始偏移量的位置
    batchIndex += 1  // 批次索引自增1
    // 获取当前批次数据的终止偏移量
    val end = batchOffsets(batchIndex)

    // 根据对应的偏移量范围构建当前批次数据的缓冲输入流
    val bufferedStream = new BufferedInputStream(ByteStreams.limit(fileStream, end - start))
    // 使用SerializerManager进行解压缩、加解密处理
    val wrappedStream = serializerManager.wrapStream(blockId, bufferedStream)
    // 生成反序列化流
    ser.deserializeStream(wrappedStream)
  } else { // 没有剩余批次数据了，做清理操作
    // No more batches left
    cleanup()
    null
  }
}
next

next用来获取下一个元素，hasNext用来判断是否还有元素，readNextItem用来真正的获取元素，我们看下获取元素的过程:

  1. 从反序列化流中读取下一个键值，并记录读取记录加1；
  2. 如果读取记录等于10000时候，读取下一个批次数据；
  3. 返回键值对。
private def readNextItem(): (K, C) = {
  try {
    // 从反序列化流中读取下一个键和值
    val k = deserializeStream.readKey().asInstanceOf[K]
    val c = deserializeStream.readValue().asInstanceOf[C]
    // 构造为二元元组
    val item = (k, c)
    // 当前批次读取记录的计数自增1
    objectsRead += 1

    // 如果当前批次读取的数据条数达到了serializerBatchSize，
    // 则需要构建下一个批次的输入流了。serializerBatchSize默认为10000，溢写时也是按照这个值进行批次划分的。
    if (objectsRead == serializerBatchSize) {
      // 要开始新的批次读取了，重置当前批次读取记录的计数为0
      objectsRead = 0
      // 获取下一个批次的输入流
      deserializeStream = nextBatchStream()
    }
    item
  } catch {
    case e: EOFException => // 读取到文件末尾了
    // 进行清理工作
    cleanup()
    null
  }
}

// 是否还有下一个键值对
override def hasNext: Boolean = {
  if (nextItem == null) { // 当前批次没有下一个键值对了
    // 当前批次的输入流也为null ，说明没有剩余了键值对了
    if (deserializeStream == null) {
      return false
    }
    // 如果当前批次的输入流不为null，尝试读取当前批次的下一个键值对
    nextItem = readNextItem()
  }
  // 根据读取的当前批次的下一个键值对是否为null判断是否还有剩余的键值对
  nextItem != null
}

// 获取下一个键值对
override def next(): (K, C) = {
  // 如果nextItem不为null，则返回nextItem， 否则使用readNextItem()方法来获取
  val item = if (nextItem == null) readNextItem() else nextItem
  // 如果获取不到下一个键值对，抛出异常
  if (item == null) {
    throw new NoSuchElementException
  }
  nextItem = null // 将nextItem置为null
  item   // 返回获取的键值对
}

被动溢写磁盘

ExternalAppendOnlyMap间接继承MemoryConsumer，是执行内存的消费方，可以申请内存，当内存不足时候，一些高优的任务可以强制从其他低优的内存消费者中抢占内存，前面分析ShuffleMapTask时候，由于map任务涉及大量的计算和复杂逻辑，不允许被强占内存，ShuffleReduce任务直接从落盘后的文件读取，可以被强占内存，先将内存中的数据溢写磁盘，然后将使用的内存归还给统一内存管理，从而为其他需要执行内存的任务提供足够的内存使用，加快处理速度，Spillable中定义了spill方法，trigger是需要抢占内存的内存消费者，size是需要申请的内存大小，这个方法是由org.apache.spark.memory.TaskMemoryManager.acquireExecutionMemory触发，具体的溢写逻辑是由forceSpill操作是由子类实现的。

override def spill(size: Long, trigger: MemoryConsumer): Long = {
  // 由于是因为trigger申请内存时导致的强制溢写，因此溢写的内存消费者不能是trigger
  // 触发溢写的MemoryConsumer不是自己，且Tungsten内存是ON_HEAP模式的
  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
  }
}

ExternalAppendOnlyMap实现了forceSpill方法，可以看出来会先检查readingIterator是否存在，如果存在溢写，溢写后归还内存，清空map。

override protected[this] def forceSpill(): Boolean = {
  assert(readingIterator != null)
  // 使用SpillableIterator进行溢写
  val isSpilled = readingIterator.spill()
  if (isSpilled) {
    currentMap = null
  }
  isSpilled
}

readingIterator是SpillableIterator实例，是ShuffleRead任务内存中数据读取的迭代器，我们来先看下它的初始化:

def destructiveIterator(inMemoryIterator: Iterator[(K, C)]): Iterator[(K, C)] = {
  // 返回SpillableIterator迭代器
  readingIterator = new SpillableIterator(inMemoryIterator)
  readingIterator
}

destructiveIterator是对当前内存中包装为SpillableIterator，供内存数据读取或者与磁盘文件进行合并操作，我们后面专门分析。

下游算子访问数据迭代器

ExternalAppendOnlyMap实现了Iterable接口，所以它实现了iterator方法，如果数据都在内存中，没有溢写到磁盘，则对内存中的数组进行排序<该部分我们之前分析AppendOnlyMap时候分析过>，返回CompletionIterator迭代器，当迭代完成后会调用freeCurrentMap法将currentMap置为null并释放内存; 如果磁盘中有数据则需要对磁盘和内存中的数据进行聚合然后返回迭代器，是通过ExternalIterator来进行的。

override def iterator: Iterator[(K, C)] = {
  if (spilledMaps.isEmpty) {
    CompletionIterator[(K, C), Iterator[(K, C)]](
      destructiveIterator(currentMap.iterator), freeCurrentMap())
  } else {
    // 否则返回ExternalIterator迭代器
    new ExternalIterator()
  }
}

只有内存数据的迭代器

SpillableIterator是用于没有溢写磁盘时候提供给后续算子计算使用的数据源迭代器，上面讲到它有可能会因为其他消费者<其他线程>内存不足而且抢占，让其落写磁盘，所以需要加锁，有以下几个变量:

  1. SPILL_LOCK用于加锁操作；
  2. nextUpstream初始化于磁盘溢写后返回的迭代器；
  3. cur是当前元素，如果为null，则表示已经读完了当前的数据；
  4. hasSpilled是否已经溢写了。
 private[this] class SpillableIterator(var upstream: Iterator[(K, C)])
    extends Iterator[(K, C)] {
      private val SPILL_LOCK = new Object()
      private var nextUpstream: Iterator[(K, C)] = null
      private var cur: (K, C) = readNext()
      private var hasSpilled: Boolean = false
  }
溢写磁盘

溢写磁盘这个是通过forceSpill调用的，给别的消费者腾出执行内存空间，由于会涉及到多线程访问这个iterator，所以需要先加锁，然后判断是否已经溢写过了，如果溢写过了就不在溢写，否则调用spillMemoryIteratorToDisk写到磁盘中，这时候会更新nextUpstream，这个类型是DiskMapIterator，在下次读取数据时候就从磁盘读取，完美的将磁盘读取和内存读取联系在一起。

def spill(): Boolean = SPILL_LOCK.synchronized {
  if (hasSpilled) { // 如果发生过溢写，则直接返回false
    false
  } else {
    // 将upstream的数据通过spillMemoryIteratorToDisk()方法溢写到磁盘
    nextUpstream = spillMemoryIteratorToDisk(upstream)
    // 标记发生了溢写
    hasSpilled = true
    true
  }
}
读取元素

hasNext来判断是否还有元素，next来获取下一个元素，主要是通过readNext来进行的，可以看出来会先看是否nextUpstream已经初始化了，如果初始化了代表已经溢写磁盘了，接下来要从磁盘读取数据了，否则还是从内存读取，内存和磁盘的统一迭代器为upstream。

def readNext(): (K, C) = SPILL_LOCK.synchronized {
  // 如果nextUpstream不为null，说明发生过溢写，此时upstream的数据已经被溢写到磁盘了，
  // 接下来的遍历需要使用nextUpstream迭代器（DiskMapIterator类型）中的数据
  if (nextUpstream != null) {
    upstream = nextUpstream
    nextUpstream = null
  }
  // 如果还存在下一个元素，则返回
  if (upstream.hasNext) {
    upstream.next()
  } else {
    null
  }
}

// 如果cur为null，说明没有更多的元素了
override def hasNext(): Boolean = cur != null

override def next(): (K, C) = {
  // 当前cur是要返回的元素，先记录
  val r = cur
  // cur更新为下一个要返回的元素
  cur = readNext()
  // 返回之前记录的cur
  r
}

内存+磁盘合并的迭代器

ExternalIterator会对磁盘和内存的数据进行合并，来进行聚合然后提供给后续算子使用的聚合数据源，是典型的的归并排序的过程，使用堆来实现，内部是StreamBuffer，它实现了Comparable接口，指定了排序规则，如下所示:

private val mergeHeap = new mutable.PriorityQueue[StreamBuffer]

private class StreamBuffer(val iterator: BufferedIterator[(K, C)], val pairs: ArrayBuffer[(K, C)])
extends Comparable[StreamBuffer] {
  // 如果内部的数组还存在键值对，说明该StreamBuffer不为空
  def isEmpty: Boolean = pairs.length == 0

  // 获取键的hash
  def minKeyHash: Int = { hashKey(pairs.head) }

  override def compareTo(other: StreamBuffer): Int = { // 堆中元素比较规则
    if (other.minKeyHash < minKeyHash) -1 else if (other.minKeyHash == minKeyHash) 0 else 1
  }
}

ExternalIterator是外部排序中的一个典型的归并过程，磁盘中的数据都是按照key的hash值排序好的，这在前面溢写磁盘时候讲到过，内存部分还没排序，需要对齐排序，然后跟磁盘的迭代器合并作为整体的输入流，由于内部key是按照hash进行排序的，所以有冲突的可能性，同一个hash值的key可能不同，所以我们需要读取每个输入流中最小hash的一组数据，使用StreamBuffer包装，放到堆中，这样子就做好了初始化操作，可以不断从堆中取出数据进行合并了。

//  获取currentMap中数据经过排序后的迭代器，使用CompletionIterator进行封装
private val sortedMap = CompletionIterator[(K, C), Iterator[(K, C)]](destructiveIterator(
  currentMap.destructiveSortedIterator(keyComparator)), freeCurrentMap())

// 将未溢写的键值对迭代器和溢写得到的键值对迭代器保存到Seq中，并为它们创建BufferedIterator
private val inputStreams = (Seq(sortedMap) ++ spilledMaps).map(it => it.buffered)

inputStreams.foreach { it =>
  // 将it迭代器中接下来键相同的键值对存放到kcPairs中
  val kcPairs = new ArrayBuffer[(K, C)]
  // 将it中接下来键相同的键值对都存入kcPairs
  readNextHashCode(it, kcPairs)
  if (kcPairs.length > 0) {
    // 将kcPairs存入mergeHeap队列
    mergeHeap.enqueue(new StreamBuffer(it, kcPairs))
  }
}
readNextHashCode

readNextHashCode会从输入流里面读取最小hash值的所有数据，首先拿取迭代器首元素，然后不断的往后查看数据，如果数据对应的hashCode跟首元素hashCode相同，说明是冲突的，加入到数组中，返回具有相同hashcode的[K,C]的数组。

private def readNextHashCode(it: BufferedIterator[(K, C)], buf: ArrayBuffer[(K, C)]): Unit = {
  if (it.hasNext) {  // 如果迭代器还存在下一个键值对
    var kc = it.next()  // 获取下一个键值对
    buf += kc  // 添加到buf中
    val minHash = hashKey(kc)   // 将键值对的键进行hash

    // 由于it迭代器中的键值对已经按照键排序了，hash冲突的情况存在，所以相同hash值的可能有多个key。
    // 我们不断迭代it迭代器，判断it中下一个键值对的键是否与kc键值对的键相同,如果相同则将该键值对存入buf。
    // 该操作会将it迭代器中元素按照键顺序取出，并将相同的键存入buf数组。
    while (it.hasNext && it.head._1.hashCode() == minHash) {
      kc = it.next()  // 获取下一个键值对
      buf += kc  // 添加到buf数组
    }
  }
}
mergeIfKeyExists

mergeIfKeyExists会根据指定的键，从与改键的hashCode相同的Buffer中拿取相同key的数据，然后合并到当前key已经聚合好的值上，相当于C=>C过程，同时从数组中删除改key，由于map端已经聚合过数据，会保证每个文件中键为key的结果只有一条数据，找到后，从数组中删除，然后合并返回结果就行。

private def mergeIfKeyExists(key: K, baseCombiner: C, buffer: StreamBuffer): C = {
  var i = 0
  while (i < buffer.pairs.length) { // 遍历StreamBuffer中的数组
    val pair = buffer.pairs(i)   // 获取键值对
    if (pair._1 == key) {  // 如果键与传入的key相同
      removeFromBuffer(buffer.pairs, i) // 则将该键值对从数组中移除
      return mergeCombiners(baseCombiner, pair._2) // 对值进行聚合操作
    }
    i += 1
  }
  baseCombiner    // 如果没有进行聚合，则直接返回原来的值
}
removeFromBuffer

removeFromBuffer会从相同hashcode的buffer数组中移除指定位置的键值对。该操作会在移除键值对之后，将数组最后一个键值对移动到被移除的键值对所在的索引位置上，并调整数组的大小，最后返回删除位置的元素。

private def removeFromBuffer[T](buffer: ArrayBuffer[T], index: Int): T = {
  val elem = buffer(index)  // 获取指定位置的键值对
  // 将数组最后一个键值对移动到被移除的键值对所在的索引位置上
  buffer(index) = buffer(buffer.size - 1) 
  buffer.reduceToSize(buffer.size - 1)  // 调整数组的大小
  // 返回移除的键值对
  elem
}
next

next方法供外界进行访问聚合后的数据迭代器，执行步骤如下:

  1. 如果堆中没有元素了，说明已经没有数据了，直接返回；
  2. 如果有数据，从堆中取出hashCode最小的一个StreamBuffer，该数组中包含了多个相同hashCode的key，我们拿取数组中首元素，先来查找所有首元素对应的Key的所有数据，进行合并然后返回，由于溢写磁盘以及内存中。
  3. 对于一个特定的key，都只有一条数据，但是一个key对应的数据可能在多个文件都存在，所以我们要不断的查看堆的首元素，如果它的hashcode和当前key的hashcode相同，那么里面有可能会有当前key的值，使用mergeIfKeyExists进行合并，然后把这个StreamBuffer放到mergedBuffers;
  4. 对于mergedBuffers，如果没有这个hashCode的数据了，需要从文件中读取下一个hashCode的数据，然后判断该文件中还有元素不，有的话需要放到堆中进行下一轮选择。

这样子就完成了内存+磁盘元素的归并排序，给到下游的是已经聚合好的数据迭代器。

override def next(): (K, C) = {
  // 如果mergeHeap为空，说明没有键值对了，直接抛出异常
  if (mergeHeap.isEmpty) {
    throw new NoSuchElementException
  }
  val minBuffer = mergeHeap.dequeue() // 拿取hashCode最小的一个StreamBuffer
  val minPairs = minBuffer.pairs // 获取该StreamBuffer内部的数组<相同hashCode的数组，key不同>
  val minHash = minBuffer.minKeyHash   // 获取最小键的哈希值 
  val minPair = removeFromBuffer(minPairs, 0)  // 先得到数组minPairs的0号索引上的键值对
  
  val minKey = minPair._1 // 获取键和值，这一轮查询只得到minKey对应的数据聚合结果
  var minCombiner = minPair._2 // 聚合后的值

  // 创建合并数据用的数组
  val mergedBuffers = ArrayBuffer[StreamBuffer](minBuffer)
  
  // 上面只是一个输入流的排序好的数据，我们是要对所有输入流的相同key进行合并，
  // 所以我们要遍历所有输入流，如果hashCode跟当前hashCode相同，说明有可能有相同的key，需要拿出来进行merge操作
  while (mergeHeap.nonEmpty && mergeHeap.head.minKeyHash == minHash) {
    val newBuffer = mergeHeap.dequeue() // 出队StreamBuffer
    minCombiner = mergeIfKeyExists(minKey, minCombiner, newBuffer)  // 对该StreamBuffer中的键值对进行聚合
    mergedBuffers += newBuffer  // 将合并过的StreamBuffer记录到mergedBuffers中
  }

  // 遍历合并过的StreamBuffer，如果其中还剩有没处理的键值对，则还需要重新处理
  mergedBuffers.foreach { buffer =>
    if (buffer.isEmpty) {  // 说明这个输入流中没有跟minKey冲突的key，读取下一批键相同的键值对放入数组
      readNextHashCode(buffer.iterator, buffer.pairs)
    }
    if (!buffer.isEmpty) { // 说明这个输入流中有跟minKey冲突的key，将该StreamBuffer再次放入mergeHeap中
      mergeHeap.enqueue(buffer)
    }
  }

  // 返回聚合后的键值对
  (minKey, minCombiner)
}

数据排序

如果指定了数据排序，还需要对聚合后的数据进行排序，使用ExternalSorter来进行排序操作，由于不需要聚合操作了，所以ExternalSorter使用PartitionedPairBuffer来进行内存数据存储，内部使用数组进行存储，用两个位置分别存储<K, C>，内存不足时候，在内存中使用TimSort先进行排序，然后进行溢写磁盘操作，在数据写入完毕后，会对磁盘和内存中数据使用归并排序，完成整体的排序，这部分在ShuffleMapTask时候已经讲解过了，不做过多阐述，如下所示:

dep.keyOrdering match {
  // 如果指定了排序函数，则创建ExternalSorter，没有指定聚合函数和分区器，只指定了排序函数
  case Some(keyOrd: Ordering[K]) =>
  // 使用ExternalSorter进行排序
  val sorter = new ExternalSorter[K, C, C](context, ordering = Some(keyOrd), serializer = dep.serializer)
  // 对数据进行缓存
  sorter.insertAll(aggregatedIter)
  // 封装并返回CompletionIterator
  CompletionIterator[Product2[K, C], Iterator[Product2[K, C]]](sorter.iterator, sorter.stop())
  // 如果没有指定排序函数，那么返回aggregatedIter
  case None =>
  aggregatedIter
}

这样子经过数据的拉取、聚合、排序后得到了下游算子需要的原始数据，相当于初始化了下一个stage的计算需要的数据。

总结

最后我们来总结下ReduceTask的整体过程，Spark中Reduce任务可以看作是下一个map任务或者result任务的数据源头，这个数据源头经过shuffle算子Map端任务产生，用ShuffleRDD来进行抽象，首先通过MapOutTracker来获取相应的需要数据的数据位置，然后使用ShuffleBlockFetcherIterator来对数据进行切分，通过BlockManger或者ShuffleClient得到各个块数据的输入流，然后如果需要聚合，则通过ExternalAppendOnlyMap来进行内存+磁盘的聚合操作，最后如果需要排序，则要通过ExternalSorter的基于内存+磁盘的排序进行数据排序，给下游算子提供聚合好数据而且排好序的数据迭代器。

参考

  1. https://www.cnblogs.com/swordfall/p/9435949.html#auto_id_16
  2. https://blog.csdn.net/lidongmeng0213/article/details/109488654
  3. https://blog.csdn.net/lidongmeng0213/article/details/109330980
  4. https://blog.csdn.net/lidongmeng0213/article/details/109359609
版权声明：本文为lidongmeng0213原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接和本声明。