Spark核心Shuffle详解（二）ShuffleHandler

1、ShuffleWriter

ShuffleWriter 是一个抽象类，定义了 map 任务将中间结果输出到磁盘上的规范。ShuffleWriter 有三个实现类：SortShuffleWriter、BypassMergeSortShuffleWriter 和 UnsafeShuffleWriter。

1、BypassMergeSortShuffleWriter

BypassMergeSortShuffleWriter是Bypass机制的具体实现，其设计目标是在适当场景下消除排序开销。

1、实现原理

BypassMergeSortShuffleWriter的核心工作流程分为三个阶段：

分区缓冲区初始化：为每个Reduce分区（分区数N≤200）创建独立的磁盘写缓冲区和临时文件。
数据分发与写入：对于每条记录，计算其目标分区ID（通过key的hashcode对N取模），然后直接追加到对应分区的缓冲区。缓冲区满时，同步将数据溢写到对应分区的临时文件中。
文件合并：处理完所有数据后，将N个临时分区文件合并为一个数据文件，并创建索引文件。这一步是关键创新，既保持了Hash Shuffle避免排序的优点，又通过文件合并解决了小文件问题。

2、UnsafeShuffleWriter

UnsafeShuffleWriter是Spark Tungsten项目的重要组成部分，直接操作堆外内存，旨在消除JVM开销并提升内存使用效率。直接通过Java Unsafe API管理内存，避免JVM对象开销和GC压力。

3、SortShuffleWriter

SortShuffleWriter是SortShuffleManager的默认实现，采用了"内存聚合→排序→溢写→合并"。

1、内存管理与数据结构

SortShuffleWriter根据算子类型选择不同的内存数据结构：

Map数据结构（ PartitionedAppendOnlyMap）：用于聚合类算子（如reduceByKey）。该结构结合了Map和Array的特性，支持高效的数据更新和聚合操作。其核心优势包括：
- 使用线性探测法解决哈希冲突，减少指针开销
- 支持原地更新聚合值，避免创建新对象
- 在内存不足时能够扩展并溢出到磁盘
Array数据结构（ PartitionedPairBuffer）：用于非聚合类算子。该结构简单追加(key, value)对，内存效率更高。当数据量超过阈值时，触发排序和溢写操作。

scala 复制代码

//ExternalSorter
@volatile private var map = new PartitionedAppendOnlyMap[K, C]
@volatile private var buffer = new PartitionedPairBuffer[K, C]

2、溢出与合并机制

SortShuffleWriter采用多级溢出-合并策略平衡内存使用和磁盘I/O：

内存填充与排序 ：数据不断添加到内存数据结构，当达到阈值（spark.shuffle.spill.initialMemoryThreshold默认5MB）时，根据分区ID和Key进行排序。
批次溢写：排序后的数据以批次（默认10,000条）形式写入临时磁盘文件，减少磁盘I/O次数。
全局合并 ：最终将内存中剩余数据与所有临时磁盘文件进行归并排序，生成最终输出。这一过程使用高效的最小堆算法，只需一次扫描即可完成多路合并。

scala 复制代码

protected def maybeSpill(collection: C, currentMemory: Long): Boolean = {
    var shouldSpill = false
    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 = acquireMemory(amountToRequest)
      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
      shouldSpill = currentMemory >= myMemoryThreshold
    }
    shouldSpill = shouldSpill || _elementsRead > numElementsForceSpillThreshold
    // Actually spill
    if (shouldSpill) {
      _spillCount += 1
      logSpillage(currentMemory)
      spill(collection)
      _elementsRead = 0
      _memoryBytesSpilled += currentMemory
      releaseMemory()
    }
    shouldSpill
  }

scala 复制代码

//spark.shuffle.spill.batchSize
private[this] def spillMemoryIteratorToDisk(inMemoryIterator: Iterator[(K, C)])
      : DiskMapIterator = {
    val (blockId, file) = diskBlockManager.createTempLocalBlock()
    val writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize, writeMetrics)
    var objectsWritten = 0

    // List of batch sizes (bytes) in the order they are written to disk
    val batchSizes = new ArrayBuffer[Long]

    // Flush the disk writer's contents to disk, and update relevant variables
    def flush(): Unit = {
      val segment = writer.commitAndGet()
      batchSizes += segment.length
      _diskBytesSpilled += segment.length
      objectsWritten = 0
    }

    var success = false
    try {
      while (inMemoryIterator.hasNext) {
        val kv = inMemoryIterator.next()
        writer.write(kv._1, kv._2)
        objectsWritten += 1

        if (objectsWritten == serializerBatchSize) {
          flush()
        }
      }
      if (objectsWritten > 0) {
        flush()
        writer.close()
      } else {
        writer.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
        writer.closeAndDelete()
      }
    }

    new DiskMapIterator(file, blockId, batchSizes)
  }

scala 复制代码

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]]
    // Use the reverse order (compare(y,x)) because PriorityQueue dequeues the max
    val heap = new mutable.PriorityQueue[Iter]()(
      (x: Iter, y: Iter) => comparator.compare(y.head._1, x.head._1))
    heap.enqueue(bufferedIters: _*)  // Will contain only the iterators with hasNext = true
    new Iterator[Product2[K, C]] {
      override def hasNext: Boolean = heap.nonEmpty

      override def next(): Product2[K, C] = {
        if (!hasNext) {
          throw new NoSuchElementException
        }
        val firstBuf = heap.dequeue()
        val firstPair = firstBuf.next()
        if (firstBuf.hasNext) {
          heap.enqueue(firstBuf)
        }
        firstPair
      }
    }
  }

2、ShuffleReader

1、BlockStoreShuffleReader

Shuffle Read阶段负责从各个Map Task的输出中获取属于当前Reduce分区的数据，并进行必要的处理（如聚合、排序），为后续计算阶段提供输入。ShuffleReader的设计直接影响Reduce Task的性能和稳定性。

1、实现原理

第1步：获取Map输出状态（SortShuffleManager##getReader）

Reduce Task首先通过 MapOutputTracker（主节点上的 MapOutputTrackerMaster）获取当前Stage中所有Map Task的输出状态信息（MapStatus）。
MapStatus 包含了每个Map Task输出的数据位置（Executor的Host和Port）以及每个Result Partition的大小。

第2步：创建拉取迭代器 - ShuffleBlockFetcherIterator (BlockStoreShuffleReader##read)

这是数据拉取的核心。它会：

划分请求： 根据 MapStatus，将需要拉取的Block分为两组：
- 本地请求： 如果Map Task和Reduce Task在同一个Executor上，则直接通过 BlockManager 从本地磁盘读取。
- 远程请求： 需要通过网络从其他Executor拉取。
并发拉取： 它会并发的发起多个远程请求（默认并行度为 spark.reducer.maxReqsInFlight，通常为1，但可能不是最优配置），以加速数据拉取过程。
流式解压与反序列化： 拉取到的数据是以序列化、压缩的字节流形式存在的。ShuffleBlockFetcherIterator 会边拉取边进行解压和反序列化，将其转换为一条条的 (Key, Value) 记录。
容错与重试： 如果某个Block拉取失败，它会进行重试。

第3步：聚合与排序 - 关键的"组合器"(BlockStoreShuffleReader##read)

反序列化后的记录流会被传递给一个"组合器"进行处理。根据Shuffle的依赖类型（Aggregator， Ordering），这里的行为有所不同：

mapSideCombine = false（默认/未在Map端聚合）：
- 无聚合、无排序： 直接返回记录流。例如 groupBy() 的某些情况。
- 有聚合： 使用 Aggregator。Spark会使用一个哈希表（ExternalAppendOnlyMap）来在内存中维护聚合状态。当内存不足时，它会将内存中的数据溢写（Spill） 到磁盘，最后再合并所有内存和磁盘文件。这就是"外部聚合"。
- 有排序： 使用 Ordering。Spark会使用一个堆排序器（ExternalSorter）。同样，当内存不足时，它会进行溢写和归并排序。
mapSideCombine = true（在Map端已聚合）：
- Reduce端接收到的已经是部分聚合的结果。此时，Reduce端的聚合器只需要对这些"部分聚合结果"进行最终合并，逻辑相同，但数据量更小，效率更高。

第4步：输出最终结果

经过聚合和排序后的最终结果，会以一个迭代器的形式返回给ShuffleReader的调用者（通常是Result Task的 runTask 方法）。Task会遍历这个迭代器，并将最终结果写入到自己的输出存储中（或直接返回给Driver）。

3、Shuffle状态跟踪管理

1、MapOutputTracker

MapOutputTracker 是 Spark Shuffle 的核心元数据管理组件，负责跟踪和管理所有 Map 任务的输出位置信息，为 Reduce 任务提供数据位置服务。

2、MapOutputTrackerMaster (org.apache.spark)

维护所有shuffle的map输出元数据，包括每个shuffle的map任务数量、每个map任务的输出位置（MapStatus）等。
响应Executor的MapOutputTrackerWorker的请求，提供map输出位置信息。
当shuffle完成后，清理对应的元数据。

scala 复制代码

// HashMap for storing shuffleStatuses in the driver.
// Statuses are dropped only by explicit de-registering.
// Exposed for testing
  
// Shuffleid -> ShuffleStatus
val shuffleStatuses = new ConcurrentHashMap[Int, ShuffleStatus]().asScala

// ShuffleStatus
// MapStatus 索引为map任务ID，值为MapStatus（包含每个map任务输出的位置和大小信息）
val mapStatuses = new Array[MapStatus](numPartitions)

3、MapOutputTrackerWorker

与Driver上的MapOutputTrackerMaster通信，获取shuffle的map输出位置信息。
缓存获取到的map输出位置信息，避免重复请求。
当缓存的信息过期时（比如Executor丢失），清理缓存。

scala 复制代码

val mapStatuses: Map[Int, Array[MapStatus]] =
    new ConcurrentHashMap[Int, Array[MapStatus]]().asScala

val mergeStatuses: Map[Int, Array[MergeStatus]] =
    new ConcurrentHashMap[Int, Array[MergeStatus]]().asScala

4、MapStatus

MapStatus是ShuffleMapTask执行完成后向Driver报告的核心元数据，包含了两类关键信息：

Task运行所在的BlockManager地址
该Task输出中每个Reduce分区的大小估算信息

压缩的 MapStatus CompressedMapStatus 默认

高度压缩的 MapStatus（用于大量 Reduce 分区）当分数数>2000使用。当size 小于阈值 reduce size 取平均值

Spark核心Shuffle详解（二）ShuffleHandler

1、ShuffleWriter

1、BypassMergeSortShuffleWriter

1、实现原理

2、UnsafeShuffleWriter

3、SortShuffleWriter

1、内存管理与数据结构

2、 溢出与合并机制

2、ShuffleReader

1、BlockStoreShuffleReader

1、实现原理

3、Shuffle状态跟踪管理

1、MapOutputTracker

2、MapOutputTrackerMaster (org.apache.spark)

3、MapOutputTrackerWorker

4、MapStatus

2、溢出与合并机制