spark源码分析-shuffle write

Spark会将job划分为多个Stage，每个job会由多个ShuffleMapStage和一个ResultStage组成，然后每个Stage会由多个Task组成，Task数量和每个Stage的Partition的数量相同。每个Task任务由单独的线程执行，不同Stage的Task之间需要进行数据流动，并且下游Stage的Task会依赖上游Stage的多个Task，所以该过程需要将数据写入磁盘，并将属于下游Stage相同Task的数据汇总到一起，以供该Task进行拉取，该过程被称为Shuffle。
shuffle过程分为shuffle上游Stage的Task的写过程和shuffle下游Stage的Task的读过程。首先看一下写过程：
写过程发生在ShuffleMapTask的runTask方法

override def runTask(context: TaskContext): MapStatus = {
    ...

    var writer: ShuffleWriter[Any, Any] = null
    try {
        // 获取Spark的ShuffleManager(SortShuffleManager)，通过ShuffleManager获取对应ShuffleWriter的具体实现子类，进行写shuffle数据的
      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
    }
  }

Spark的ShuffleWriter总共有三个实现类：
UnsafeShuffleWriter：主要用于对序列化的二进制数据进行排序，避免使用java对象，减少内存使用和GC开销，而且不需要序列化和反序列化，排序数据只需要使用分区id、key和指向数据的内存地址，可以避免cache-miss，优化排序性能。
BypassMergeSortShuffleWriter：该类会同时打开R个shuffle文件进行写入，如果shuffle数据量太大会生成太多的分区，从而造成同时打开太多文件，所以主要用于分区数小于配置参数spark.shuffle.sort.bypassMergeThreshold值的shuffle。但是如果shuffle使用了map端聚合操作的，无论分区数大小都不会使用该类。
SortShuffleWriter：该ShuffleWriter会优化BypassMergeSortShuffleWriter写出shuffle文件数过多的问题，而且该类可以使用map端聚合优化。
接下来主要使用SortShuffleWriter进行shuffle write流程的分析，后边会对其余两个ShuffleWriter的分析。
首先看一下write()方法，该方法会将数据写到ExternalSorter中，ExternalSorter会将先将数据写入到内存中然后对数据进行排序，当内存中数据到达一定大小会先将数据写到临时文件中，然后ExternalSorter会将所有写出的临时文件和最后未写出的数据进行merge sort排序，最后将排序完的数据写出到一个数据文件、一个索引文件。

 override def write(records: Iterator[Product2[K, V]]): Unit = {
    sorter = if (dep.mapSideCombine) {
    // shuffle过程有map端聚合优化，所有构造的ExternalSorter类传入聚合方法
    // ExternalSorter类的K为数据的key的类型，V为数据的类型，C为用于数据排序过程中使用的比较器类型
      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 {
    // 没有map端聚合优化，构造的ExternalSorter类的泛型的第二个、第三个类型相同
      new ExternalSorter[K, V, V](
        context, aggregator = None, Some(dep.partitioner), ordering = None, dep.serializer)
    }
    // 将数据写入排序器进行数据排序
    sorter.insertAll(records)

    // 生成shuffle结果的临时文件，最终数据文件名为"shuffle_" + shuffleId + "_" + mapId + "_" + reduceId + ".data",临时文件名为"shuffle_" + shuffleId + "_" + mapId + "_" + reduceId + ".data" + "." + uuid
    val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
    val tmp = Utils.tempFileWith(output)
    try {
      val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
      // 写排好序的shuffle数据到临时文件
      val partitionLengths =