[SPARK][CORE] 面试问题之 SortShuffleWriter的实现详情

// sortShuffleWriter
override def write(records: Iterator[Product2[K, V]]): Unit = {
// [1] 首先创建基于JVM的外排器ExternalSorter， 如果是需要mapSide聚合的，封装进去aggregator和ordering
sorter= if (dep.mapSideCombine) {
    new ExternalSorter[K, V, C](
      context,dep.aggregator,Some(dep.partitioner),dep.keyOrdering,dep.serializer)
  } else {
    // In this case we pass neither an aggregator nor an ordering to the sorter, because we don't
    // care whether the keys get sorted in each partition; that will be done on the reduce side
    // if the operation being run is sortByKey.
    new ExternalSorter[K, V, V](
      context, aggregator = None,Some(dep.partitioner), ordering = None,dep.serializer)
  }
// [2] mapTask的records全部insert到外部排序器
sorter.insertAll(records)

  // Don't bother including the time to open the merged output file in the shuffle write time,
  // because it just opens a single file, so is typically too fast to measure accurately
  // (see SPARK-3570).
  // [3] 创建处理mapTask所有分区数据commit提交writer
  val mapOutputWriter = shuffleExecutorComponents.createMapOutputWriter(
dep.shuffleId, mapId,dep.partitioner.numPartitions)
// [4] 将写入ExternalSorter中的所有数据写出到一个map output writer中
sorter.writePartitionedMapOutput(dep.shuffleId, mapId, mapOutputWriter)
// [5] 提交所有分区长度，生成索引文件
partitionLengths= mapOutputWriter.commitAllPartitions(sorter.getChecksums).getPartitionLengths
mapStatus=MapStatus(blockManager.shuffleServerId,partitionLengths, mapId)
}

// ExternalSorter
def insertAll(records: Iterator[Product2[K, V]]): Unit = {
  //TODO: stop combining if we find that the reduction factor isn't high
val shouldCombine = aggregator.isDefined

  // [1] 是否需要在mapSide的聚合
  if (shouldCombine) {
    // [1.1] 通过aggregator获取mergeValue和createCombiner
    // Combine values in-memory first using our AppendOnlyMap
    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)
    }
   // [2] 如果需要map端聚合，将数据写入map缓存中
    while (records.hasNext) {
      addElementsRead()
      kv = records.next()
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()
   // [2] 如果不需要map端聚合，将数据写入buffer缓存中
buffer.insert(getPartition(kv._1), kv._1, kv._2.asInstanceOf[C])
      // [3] 判断是否需要溢写，并进行溢写
      maybeSpillCollection(usingMap = false)
    }
  }
}

```

   在`AppendOnlyMap` 中时将k-v依次放入到数组中缓存数据。在HASH冲突时，Java原生的`HashMap`是通过拉链法去解决hash冲突的，`AppendOnlyMap`是通过开放地址法–线性探测的方法进行解决冲突的，线性探测间隔总是固定的，通常为1。 `AppendOnlyMap`支持key为null的情况，使用一个变量`nullValue`保存对应的值，`haveNullValue`表示是否存在null的key，如果之前不存在，就将size+1，然后更新值；存在时候只需要更新值即可；另外一点和java的`HashMap`不同的是，`AppendOnlyMap`提供了聚合的方法，来应对[shuffle](<https://so.csdn.net/so/search?q=shuffle&spm=1001.2101.3001.7020>)过程中指定了map-side聚合的情况，使用者需要提供`updateFunc` 。

   由于PartitionedPairBuffer只是一个数据缓冲区，不需要对元素进行聚合操作等，所以添加元素直接将元素append到数组的back即可，不过需要先判断数据容量是否已经满了，满了则需要扩容。然后首先会将<partition, key>作为Tuple放在2*curSize位置上，然后相邻位置2*curSize+1放具体的value，添加完毕后需要进行重采样操作。

// ExternalSorter
private def maybeSpillCollection(usingMap: Boolean): Unit = {
  var estimatedSize = 0L
  if (usingMap) {
    // [1] 估算当前缓存数据结构的size
    estimatedSize =map.estimateSize()
    // [2] 判断是否需要溢写，如果执行溢写后，会重新创建缓存数据结构
    if (maybeSpill(map, estimatedSize)) {
map= new PartitionedAppendOnlyMap[K, C]
    }
  } else {
    estimatedSize =buffer.estimateSize()
    if (maybeSpill(buffer, estimatedSize)) {
buffer= new PartitionedPairBuffer[K, C]
    }
  }
   // [3] 记录当前的峰值内存
  if (estimatedSize >_peakMemoryUsedBytes) {
_peakMemoryUsedBytes= estimatedSize
  }
}

protected def maybeSpill(collection: C, currentMemory: Long): Boolean = {
  var shouldSpill = false
  // [1] 如果当前的记录数是32的倍数， 同时当前内存超过了门限，默认5M
  if (elementsRead % 32 == 0 && currentMemory >=myMemoryThreshold) {
    // Claim up to double our current memory from the shuffle memory pool
    // [2] 尝试申请2倍当前内存，并将门限调整为两倍当前内存
    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
    // [3] 如果没申请下来，则应该进行spill, 或者当前写入的records数大于了强制spill门限，默认是整数的最大值
    shouldSpill = currentMemory >=myMemoryThreshold
}
  shouldSpill = shouldSpill ||_elementsRead>numElementsForceSpillThreshold
 // [4] 进行spill
// Actually spill
  if (shouldSpill) {
_spillCount+= 1
    logSpillage(currentMemory)
    spill(collection)
_elementsRead= 0
_memoryBytesSpilled+= currentMemory
    releaseMemory()
  }
  shouldSpill
}

// ExternalSorter
override protected[this] def spill(collection: WritablePartitionedPairCollection[K, C]): Unit = {
    //【根据指定的比较器comparator进行排序，返回排序结果的迭代器】
    //【如果细看的话，destructiveSortedWritablePartitionedIterator()方法最终采用TimSort算法来排序】
    val inMemoryIterator = collection.destructiveSortedWritablePartitionedIterator(comparator)
    //【将内存数据溢写到磁盘文件】
    val spillFile = spillMemoryIteratorToDisk(inMemoryIterator)
    //【用ArrayBuffer记录所有溢写的磁盘文件】
    spills += spillFile
  }

/**
 * Spill contents of in-memory iterator to a temporary file on disk.
 */
private[this] def spillMemoryIteratorToDisk(inMemoryIterator: WritablePartitionedIterator[K, C])
    : SpilledFile = {
  // Because these files may be read during shuffle, their compression must be controlled by
  // spark.shuffle.compress instead of spark.shuffle.spill.compress, so we need to use
  // createTempShuffleBlock here; see SPARK-3426 for more context.
  // [1] 创建临时的blockid和对应的file
  val (blockId, file) =diskBlockManager.createTempShuffleBlock()

  // These variables are reset after each flush
  var objectsWritten: Long = 0
  val spillMetrics: ShuffleWriteMetrics = new ShuffleWriteMetrics
  // [2] 创建个DiskBlockObjectWriter的写出流
  val writer: DiskBlockObjectWriter =
blockManager.getDiskWriter(blockId, file,serInstance,fileBufferSize, spillMetrics)

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

  // How many elements we have in each partition
  val elementsPerPartition = new Array[Long](numPartitions)

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

  var success = false
  try {
    // [3] 遍历内存数据结构中的数据，在调用writeNext迭代器时会根据comparator按key排序，缓存中的key为（partitionId, key）, 会先按分区排序，再按key排序。
    while (inMemoryIterator.hasNext) {
      val partitionId = inMemoryIterator.nextPartition()
require(partitionId >= 0 && partitionId <numPartitions,
        s"partition Id:${partitionId} should be in the range [0,${numPartitions})")
      inMemoryIterator.writeNext(writer)
      // [3.1] 记录每个分区的元数数
      elementsPerPartition(partitionId) += 1
      objectsWritten += 1
      // [3.2] 默认每1000条生成一个fileSegement
      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.revertPartialWritesAndClose()
      if (file.exists()) {
        if (!file.delete()) {
          logWarning(s"Error deleting${file}")
        }
      }
    }
  }
// [4] 最终将溢写的文件封装为SpilledFile返回
SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition)
}

// ExternalSorter
def writePartitionedMapOutput(
    shuffleId: Int,
    mapId: Long,
    mapOutputWriter: ShuffleMapOutputWriter): Unit = {
  var nextPartitionId = 0
  if (spills.isEmpty) {
    // Case where we only have in-memory data
    val collection = if (aggregator.isDefined)mapelsebuffer
val it = collection.destructiveSortedWritablePartitionedIterator(comparator)
    while (it.hasNext) {
      ...
    }
  } else {
    // We must perform merge-sort; get an iterator by partition and write everything directly.
    // [1] 调用分区迭代器，将分区数据生成(id, elements)二元组
    for ((id, elements) <- this.partitionedIterator) {
      val blockId =ShuffleBlockId(shuffleId, mapId, id)
      var partitionWriter: ShufflePartitionWriter = null
      var partitionPairsWriter: ShufflePartitionPairsWriter = null
      TryUtils.tryWithSafeFinally{
        // 每个分区打开的writer进行并发写入的优化，最终生成一个文件
        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)
}

// ExternalSorter
def partitionedIterator: Iterator[(Int, Iterator[Product2[K, C]])] = {
  val usingMap = aggregator.isDefined
  val collection: WritablePartitionedPairCollection[K, C] = if (usingMap)mapelsebuffer
// [1] 如果没有溢写，直接groupByPartition
if (spills.isEmpty) {
    // Special case: if we have only in-memory data, we don't need to merge streams, and perhaps
    // we don't even need to sort by anything other than partition ID
    if (ordering.isEmpty) {
      // The user hasn't requested sorted keys, so only sort by partition ID, not key
      groupByPartition(destructiveIterator(collection.partitionedDestructiveSortedIterator(None)))
    } else {
      // We do need to sort by both partition ID and key
      groupByPartition(destructiveIterator(
        collection.partitionedDestructiveSortedIterator(Some(keyComparator))))
    }
  } else {
   // [2] 存在溢写，需要先将在内存中和溢写文件中的数据封装为迭代器执行归并排序， 归并排序时通过最小堆实现的
    // Merge spilled and in-memory data
    merge(spills.toSeq, destructiveIterator(
      collection.partitionedDestructiveSortedIterator(comparator)))
  }
}