Spark-ShuffleWriter-SortShuffleWriter

一、上下文

《Spark-ShuffleWriter》中对ShuffleWriter的获取、分类和写入做了简单的分析，下面我们对其中的SortShuffleWriter做更详细的学习

二、构建ExternalSorter

ExternalSorter是一个通用的排序器

可以使用 spark.shuffle.conspress 控制块压缩

将（K，V）型的数据进行排序并合并，生成类型为（K，C）的键组合器对。

使用Partitioner首先将Key进行分区，然后使用自定义Comparator对每个分区中的Key进行排序。输出单个分区文件，每个分区具有不同的字节范围，方便用于下游Shuffle拉取。

如果map端没有聚合，那么类型C = 类型 V

构建ExternalSorter所需要的参数如下：

aggregator ：可选的聚合器，带有用于合并数据的组合函数
partitioner ： 可选分区器；如果给定，则按分区ID排序，然后按Key排序
ordering ： 可选排序; 每个分区内的键；全排序
serializer ： 溢写数据到磁盘所用的序列化器

请注意：

        如果map端没有聚合，会将ordering置为None，以避免额外的排序

        如果map端有聚合，会有ordering，且效率比不设置更高

使用方法：

        1、实例化ExternalSorter

        2、调用insertAll()将一组数据写进去

        3、使用 iterator() 以遍历排序/聚合的数据或者调用writePartitionedFile()来创建一个包含排序/聚合输出的文件，用于Spark的排序洗牌

该类的内部工作方式：

        反复填充内存中数据的缓冲区，如果想按Key组合，则使用PartitionedAppendOnlyMap，如果我们不想按Key结合，则使用PartitionedPairBuffer。在这些缓冲区中，我们按分区ID对元素进行排序，然后也可能按Key排序。为了避免使用每个Key多次调用分区器，我们将分区ID存储在每条记录旁边。

        当每个缓冲区达到设置的内存限制时，将其溢出到一个文件中。如果我们想进行聚合，这个文件首先按分区ID排序，然后可能按Key或Key的哈希码排序。对于每个文件，我们跟踪内存中每个分区中有多少对象，所以不必为每个元素写出分区ID。

        当用户请求迭代器或文件输出时，溢出的文件将与任何剩余的内存数据一起合并，使用上面定义的相同排序顺序（除非禁用排序和聚合）。如果我们需要按键聚合，我们要么使用排序参数的总排序，要么使用相同的哈希码读取Key，并将它们相互比较相等性以合并值。

注意：在shuffle过程中map端指的是写入结果的一端，reduce端指的是拉取结果的一端

1、map端有聚合

new ExternalSorter[K, V, C](
        context, dep.aggregator, Some(dep.partitioner), dep.keyOrdering, dep.serializer)

context：有值

aggregator：有值

partitioner：有值

ordering：有值

serializer：有值

2、map端无聚合

new ExternalSorter[K, V, V](
        context, aggregator = None, Some(dep.partitioner), ordering = None, dep.serializer)

context：有值

aggregator：无值

partitioner：有值

ordering：无值

serializer：有值

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 with ShuffleChecksumSupport {

  //下游分区数
  private val numPartitions = partitioner.map(_.numPartitions).getOrElse(1)
  //如果下游分区数=1，那么就不需要分区了
  private val shouldPartition = numPartitions > 1
  private def getPartition(key: K): Int = {
    if (shouldPartition) partitioner.get.getPartition(key) else 0
  }

  //在每个节点（driver和executor）上运行的管理器，它提供了将块本地和远程放入和检索到各种存储（内存、磁盘和堆外）的接口。
  private val blockManager = SparkEnv.get.blockManager
  //创建并维护逻辑块和磁盘上物理位置之间的逻辑映射。一个块映射到一个文件，其名称由其BlockId给出。
  //块文件在spark.local.dir（或spark_local_DIRS，如果已设置）中列出的目录之间进行哈希运算。
  private val diskBlockManager = blockManager.diskBlockManager
  //为各种Spark组件配置序列化、压缩和加密的组件，包括自动选择用于混洗的[[Serializer]]。
  private val serializerManager = SparkEnv.get.serializerManager
  //序列化器。由于某些序列化库不是线程安全的，因此此类用于创建执行实际序列化并保证一次只能从一个线程调用的[[org.apache.sspark.serializer.SerializerInstance]]对象。
  private val serInstance = serializer.newInstance()

  //spark.shuffle.file.buffer 默认值 32K 
  //每个shuffle文件输出流的内存缓冲区大小，单位为KiB。这些缓冲区减少了在创建中间shuffle文件时进行的磁盘查找和系统调用的数量。
  private val fileBufferSize = conf.get(config.SHUFFLE_FILE_BUFFER_SIZE).toInt * 1024
  //spark.shuffle.spill.batchSize 默认值 10000
  //从序列化器读/写时对象批的大小。
  //对象是分批写的，每批都使用自己的序列化流。这减少了在反序列化流时构建的引用跟踪图的大小。
  //注意：将此值设置得太低可能会导致序列化时的过度复制，因为一些序列化器在每次对象数量加倍时都会通过growing + copying来增长内部数据结构。
  private val serializerBatchSize = conf.get(config.SHUFFLE_SPILL_BATCH_SIZE)

  //在溢写之前存储在内存中对象的数据结构。
  //根据我们map端是否有聚合，我们要么将对象放入AppendOnlyMap中进行组合，要么将它们存储在数组缓冲区中。
  //WritablePartitionedPairCollection的实现，它封装了一个map
  //  ((分区ID，K) ->  V)
  @volatile private var map = new PartitionedAppendOnlyMap[K, C]
  //仅附加键值对缓冲区，每个缓冲区都有相应的分区ID，以跟踪其估计的字节大小。
  //缓冲器最多可支持10 7374 1819个元素
  @volatile private var buffer = new PartitionedPairBuffer[K, C]

  //...................

  //溢写文件的信息
  private[this] case class SpilledFile(
    file: File,
    blockId: BlockId,
    serializerBatchSizes: Array[Long],
    elementsPerPartition: Array[Long])

  //溢写文件数组
  private val spills = new ArrayBuffer[SpilledFile]


}

三、将结果数据插入ExternalSorter

调用了ExternalSorter的insertAll(records)

  def insertAll(records: Iterator[Product2[K, V]]): Unit = {
    // map端是否聚合
    val shouldCombine = aggregator.isDefined

    if (shouldCombine) {
      // 使用 AppendOnlyMap 在内存中聚合 value
      // 聚合器有三个方法：
      // 1、createCombiner ：用于创建聚合的初始值
      // 2、mergeValue ：用于将新值合并到聚合结果中 
      // 3、mergeCombiners ：用于合并多个mergeValue函数的输出 （这个只能在reduce端做，因为map端只有一个key的部分数据）
      // 1和2可以在map端做也可以在reduce端做，如果业务符合在map端做可以压缩数据量，降低带宽消耗
      val mergeValue = aggregator.get.mergeValue
      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)
      }
      while (records.hasNext) {
        addElementsRead()
        kv = records.next()
        //最终形成 (分区id，Key)  ->  Values  的结构
        map.changeValue((getPartition(kv._1), kv._1), update)
        //如果需要，将当前的内存数据溢写到磁盘。
        maybeSpillCollection(usingMap = true)
      }
    } else {
      // 如果不聚合 将数据直接放入缓冲区 
      while (records.hasNext) {
        addElementsRead()
        val kv = records.next()
        //将(分区id，Key,Value) 元素插入 buffer
        //数组的第2n索引位置存：(partition, key.asInstanceOf[AnyRef])
        //数组的第2n+1索引位置存：value.asInstanceOf[AnyRef]
        //真是的Key 和 Value 都存在堆里
        buffer.insert(getPartition(kv._1), kv._1, kv._2.asInstanceOf[C])
        如果需要，将当前的内存数据溢写到磁盘
        maybeSpillCollection(usingMap = false)
      }
    }
  }

PartitionedPairBuffer

从上面可以看到，如果map端没有聚合就将数据放入PartitionedPairBuffer中，下面我们通过它的源码看下数据是如何放的

private[spark] class PartitionedPairBuffer[K, V](initialCapacity: Int = 64)
  extends WritablePartitionedPairCollection[K, V] with SizeTracker
{
  import PartitionedPairBuffer._

  require(initialCapacity <= MAXIMUM_CAPACITY,
    s"Can't make capacity bigger than ${MAXIMUM_CAPACITY} elements")
  require(initialCapacity >= 1, "Invalid initial capacity")

  //基本的可增长数组数据结构。我们使用AnyRef的单个数组来保存键和值，这样我们就可以使用KVArraySortDataFormat对它们进行高效排序。
  //初始容量 64
  private var capacity = initialCapacity
  private var curSize = 0
  //数组长度为 64 * 2 因为数组是用2个字节来存一个pkv
  private var data = new Array[AnyRef](2 * initialCapacity)

  /** 添加一个pkv元素到 buffer */
  def insert(partition: Int, key: K, value: V): Unit = {
    //如果已经等于容量，那么进行扩容
    if (curSize == capacity) {
      growArray()
    }
    data(2 * curSize) = (partition, key.asInstanceOf[AnyRef])
    data(2 * curSize + 1) = value.asInstanceOf[AnyRef]
    curSize += 1
    afterUpdate()
  }

  /** 进行两倍扩容 */
  private def growArray(): Unit = {
    if (capacity >= MAXIMUM_CAPACITY) {
      throw new IllegalStateException(s"Can't insert more than ${MAXIMUM_CAPACITY} elements")
    }
    val newCapacity =
      if (capacity * 2 > MAXIMUM_CAPACITY) { // Overflow
        MAXIMUM_CAPACITY
      } else {
        capacity * 2
      }
    //创建一个新数组
    val newArray = new Array[AnyRef](2 * newCapacity)
    //从指定的源数组开始，将数组复制到目标数组的指定位置。
    //JDK本地方法
    System.arraycopy(data, 0, newArray, 0, 2 * capacity)
    data = newArray
    capacity = newCapacity
    resetSamples()
  }

  /** 按照给定的顺序迭代数据 */
  override def partitionedDestructiveSortedIterator(keyComparator: Option[Comparator[K]])
    : Iterator[((Int, K), V)] = {
    val comparator = keyComparator.map(partitionKeyComparator).getOrElse(partitionComparator)
    new Sorter(new KVArraySortDataFormat[(Int, K), AnyRef]).sort(data, 0, curSize, comparator)
    iterator
  }

  private def iterator(): Iterator[((Int, K), V)] = new Iterator[((Int, K), V)] {
    var pos = 0

    override def hasNext: Boolean = pos < curSize

    override def next(): ((Int, K), V) = {
      if (!hasNext) {
        throw new NoSuchElementException
      }
      val pair = (data(2 * pos).asInstanceOf[(Int, K)], data(2 * pos + 1).asInstanceOf[V])
      pos += 1
      pair
    }
  }
}

//最大容量为Integer.MAX_VALUE - 15 
//某些JVM无法分配长度为Integer.MAX_VALUE的数组。这里保守一点设置的比极限值稍微小一点
private object PartitionedPairBuffer {
  val MAXIMUM_CAPACITY: Int = ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH / 2
}

 AppendOnlyMap

class AppendOnlyMap[K, V](initialCapacity: Int = 64)
  extends Iterable[(K, V)] with Serializable {

  private val LOAD_FACTOR = 0.7
  private var capacity = nextPowerOf2(initialCapacity)
  private var mask = capacity - 1
  private var curSize = 0
  private var growThreshold = (LOAD_FACTOR * capacity).toInt

  //将Key和Value保存在同一个数组中以实现内存局部性；
  //具体来说，元素的顺序是key0、value0、key1、value1、key2、value2 ......
  private var data = new Array[AnyRef](2 * capacity)

  //将key的值设置为updateFunc（hadValue，oldValue）函数的结果，返回新更新的值
  def changeValue(key: K, updateFunc: (Boolean, V) => V): V = {
    assert(!destroyed, destructionMessage)
    val k = key.asInstanceOf[AnyRef]
    if (k.eq(null)) {
      if (!haveNullValue) {
        incrementSize()
      }
      nullValue = updateFunc(haveNullValue, nullValue)
      haveNullValue = true
      return nullValue
    }
    //计算该数据应该要放的位置
    var pos = rehash(k.hashCode) & mask
    var i = 1
    while (true) {
      val curKey = data(2 * pos)
      if (curKey.eq(null)) {
        //如果该数组中没有放过这个key的值，那么就直接放
        val newValue = updateFunc(false, null.asInstanceOf[V])
        data(2 * pos) = k
        data(2 * pos + 1) = newValue.asInstanceOf[AnyRef]
        incrementSize()
        return newValue
      } else if (k.eq(curKey) || k.equals(curKey)) {
        //如果该数组放过这个key的值，且与这个key的值内容相等，那么走聚合的第二个函数，合并旧值
        val newValue = updateFunc(true, data(2 * pos + 1).asInstanceOf[V])
        data(2 * pos + 1) = newValue.asInstanceOf[AnyRef]
        return newValue
      } else {
        //hash碰撞处理
        //如果hash后得到了位置，且有值，但内容又不相等，需要向右移动再进行判断
        //这的确会浪费性能，但是相比让cpu多转会儿，减少磁盘IO的次数多用点内存节省的时间比这更多
        val delta = i
        pos = (pos + delta) & mask
        i += 1
      }
    }
    null.asInstanceOf[V]
  }

}

四、构建MapOutputWriter

需要用到三个参数：

        shuffleId ：map任务所属shuffle的唯一标识符

        mapTaskId ：map任务的ID。该ID在此Spark应用程序中是唯一的。

        numPartitions ：map任务将写入的分区数。其中一些分区可能是空的。

ShuffleExecutorComponents会调用其子类LocalDiskShuffleExecutorComponents来创建，最近会new出来一个LocalDiskShuffleMapOutputWriter（ShuffleMapOutputWriter的子类）

LocalDiskShuffleMapOutputWriter

将shuffle数据与索引文件一起复制到本地磁盘的功能

public class LocalDiskShuffleMapOutputWriter implements ShuffleMapOutputWriter {

  //创建并维护逻辑块和物理文件位置之间的Shuffle块映射。
  //来自同一map任务的Shuffle块的数据存储在单个合并数据文件中。后缀为: .data
  //"shuffle_" + shuffleId + "_" + mapId + "_" + reduceId + ".data"
  //数据文件中数据块的偏移量存储在单独的索引文件中。  后缀为： .index
  //"shuffle_" + shuffleId + "_" + mapId + "_" + reduceId + ".index"
  private final IndexShuffleBlockResolver blockResolver;
  
  public LocalDiskShuffleMapOutputWriter(
      int shuffleId,
      long mapId,
      int numPartitions,
      IndexShuffleBlockResolver blockResolver,
      SparkConf sparkConf) {
    this.shuffleId = shuffleId;
    this.mapId = mapId;
    this.blockResolver = blockResolver;
    //spark.shuffle.unsafe.file.output.buffer  默认值 32k
    //在 unsafe shuffle writer中写入每个分区后，此缓冲区大小的文件系统
    this.bufferSize =
      (int) (long) sparkConf.get(
        package$.MODULE$.SHUFFLE_UNSAFE_FILE_OUTPUT_BUFFER_SIZE()) * 1024;
    this.partitionLengths = new long[numPartitions];
    //根据shuffleId, mapId 获取 Shuffle 数据文件 创建或获取文件
    this.outputFile = blockResolver.getDataFile(shuffleId, mapId);
    this.outputTempFile = null;
  }

}

五、将ExternalSorter结果写入MapOutputWriter

  def writePartitionedMapOutput(
      shuffleId: Int,
      mapId: Long,
      mapOutputWriter: ShuffleMapOutputWriter): Unit = {
    var nextPartitionId = 0
    if (spills.isEmpty) {
      // 所有数据都在内存的场景
      //因此数据源 来自 map 或者 buffer
      val collection = if (aggregator.isDefined) map else buffer
      //按照分区id排序，获取一个新的迭代器
      val it = collection.destructiveSortedWritablePartitionedIterator(comparator)
      while (it.hasNext) {
        val partitionId = it.nextPartition()
        var partitionWriter: ShufflePartitionWriter = null
        var partitionPairsWriter: ShufflePartitionPairsWriter = null
        TryUtils.tryWithSafeFinally {
          //为每个分区创建一个可以打开输出流的写入程序，以持久化给定reduce分区id的目标字节。
          partitionWriter = mapOutputWriter.getPartitionWriter(partitionId)
          //格式化shuffle block
          //"shuffle_" + shuffleId + "_" + mapId + "_" + reduceId
          val blockId = ShuffleBlockId(shuffleId, mapId, partitionId)
          //KV写入器，它将字节推送到任意分区写入器，而不是通过块管理器写入本地磁盘。
          partitionPairsWriter = new ShufflePartitionPairsWriter(
            partitionWriter,
            serializerManager,
            serInstance,
            blockId,
            context.taskMetrics().shuffleWriteMetrics,
            if (partitionChecksums.nonEmpty) partitionChecksums(partitionId) else null)
          //数据结构是 (分区id，Key)  ->  value 
          //一个分区下肯定有多个Key ,会将分区一样的聚合后的 Key Value 写入一起
          while (it.hasNext && it.nextPartition() == partitionId) {
            //只将 KV 写入
            it.writeNext(partitionPairsWriter)
          }
        } {
          if (partitionPairsWriter != null) {
            partitionPairsWriter.close()
          }
        }
        nextPartitionId = partitionId + 1
      }
    } else {
      // 按分区获取迭代器，必须执行合并排序；并直接写入所有内容。
      //partitionedIterator 这是一个方法，返回一个迭代器
      for ((id, elements) <- this.partitionedIterator) {
        val blockId = ShuffleBlockId(shuffleId, mapId, id)
        var partitionWriter: ShufflePartitionWriter = null
        var partitionPairsWriter: ShufflePartitionPairsWriter = null
        TryUtils.tryWithSafeFinally {
          partitionWriter = mapOutputWriter.getPartitionWriter(id)
          partitionPairsWriter = new ShufflePartitionPairsWriter(
            partitionWriter,
            serializerManager,
            serInstance,
            blockId,
            context.taskMetrics().shuffleWriteMetrics,
            if (partitionChecksums.nonEmpty) partitionChecksums(id) else null)
          if (elements.hasNext) {
            for (elem <- elements) {
              partitionPairsWriter.write(elem._1, elem._2)
            }
          }
        } {
          if (partitionPairsWriter != null) {
            partitionPairsWriter.close()
          }
        }
        nextPartitionId = id + 1
      }
    }

    context.taskMetrics().incMemoryBytesSpilled(memoryBytesSpilled)
    context.taskMetrics().incDiskBytesSpilled(diskBytesSpilled)
    context.taskMetrics().incPeakExecutionMemory(peakMemoryUsedBytes)
  }


 //返回一个迭代器，遍历写入此对象的所有数据，按分区分组，并由请求的聚合器聚合。对于每个分区，我们在其内容上都有一个迭代器，这些内容应该按顺序访问（如果不读取上一个分区，就不能“跳过”到一个分区）。保证按分区ID的顺序为每个分区返回一个键值对。
//目前，只需一次性合并所有溢出的文件
 def partitionedIterator: Iterator[(Int, Iterator[Product2[K, C]])] = {
    val usingMap = aggregator.isDefined
    val collection: WritablePartitionedPairCollection[K, C] = if (usingMap) map else buffer
    if (spills.isEmpty) {
      // 特殊情况：如果我们只有内存中的数据，我们不需要合并流，也许我们甚至不需要按分区ID以外的任何东西排序
      //判断是否定义排序
      if (ordering.isEmpty) {
        // 用户没有请求排序Key，因此只按分区ID排序，而不是按Key排序
        groupByPartition(destructiveIterator(collection.partitionedDestructiveSortedIterator(None)))
      } else {
        //我们确实需要按分区ID和键进行排序
        groupByPartition(destructiveIterator(
          collection.partitionedDestructiveSortedIterator(Some(keyComparator))))
      }
    } else {
      // 合并溢出和内存中的数据
      merge(spills.toSeq, destructiveIterator(
        collection.partitionedDestructiveSortedIterator(comparator)))
    }
  }

六、MapOutputWriter写入最终结果

最后的数据写入有LocalDiskShuffleMapOutputWriter完成

LocalDiskShuffleMapOutputWriter

  public MapOutputCommitMessage commitAllPartitions(long[] checksums) throws IOException {
    // 检查transferTo循环后的位置是否正确，如果不正确，则引发异常。
    if (outputFileChannel != null && outputFileChannel.position() != bytesWrittenToMergedFile) {
      throw new IOException(
          ".......");
    }
    cleanUp();
    File resolvedTmp = outputTempFile != null && outputTempFile.isFile() ? outputTempFile : null;
    //为这个map任务写 Shuffle index 文件
    log.debug("Writing shuffle index file for mapId {} with length {}", mapId,
        partitionLengths.length);
    //将数据和元数据文件作为原子操作提交
    //有两种元数据文件：
    // ----index file
    //    索引文件包含每个块的偏移量，以及输出文件末尾的最终偏移量。[[getBlockData]]将使用它来计算每个块的开始和结束位置。
    //-----checksum file 
    //    校验和文件包含每个块的校验和。它将用于诊断块损坏的原因。请注意，空的“校验和”表示校验和已禁用。
    blockResolver
      .writeMetadataFileAndCommit(shuffleId, mapId, partitionLengths, checksums, resolvedTmp);
    //表示为ShuffleMapTask写入map输出的结果
    //分区长度表示映射任务中写入的每个块的长度。这可用于下游读取器分配资源，例如内存缓冲区。
    //映射输出编写器可以选择附加任意元数据标签，以向shuffle输出跟踪器（一个目前尚未在shuffle存储API的未来迭代中构建的模块）注册。用于下游读取时使用
    return MapOutputCommitMessage.of(partitionLengths);
  }

七、向调度器报告摘要信息

ShuffleMapTask向调度器返回的结果。包括任务存储shuffle文件的块管理器地址，以及每个reducer的输出大小，以便传递给reduce任务。

当调度器发现这个ShuffleMapTask执行完成，就会执行下一个ShuffleMapTask或者ResultTask