SparkCore — SortShuffle源码分析下

SortShuffle 源码分析下

  接着上一篇博客，我们分析到了maybeSpill()它会分析是否需要进行溢写到磁盘操作，假如需要溢写那么会调用 spill()方法，这个方法是在Spillable中定义的，它是一个trait，因此我们找它的子类，是ExternalSorter，因此到这里看这个方法的源码：

override protected[this] def spill(collection: WritablePartitionedPairCollection[K, C]): Unit = {
    // 获取临时blockId和一个临时文件
    val (blockId, file) = diskBlockManager.createTempShuffleBlock()

    var objectsWritten: Long = 0
    var spillMetrics: ShuffleWriteMetrics = null
    // 获取writer，用于将数据写入磁盘文件中
    var writer: DiskBlockObjectWriter = null
    def openWriter(): Unit = {
      assert (writer == null && spillMetrics == null)
      spillMetrics = new ShuffleWriteMetrics
      // 这里fileBufferSize大小是32K
      writer = blockManager.getDiskWriter(blockId, file, serInstance, fileBufferSize, spillMetrics)
    }
    // 创建一个writer
    openWriter()

    // 存储写入batch的缓存大小
    val batchSizes = new ArrayBuffer[Long]

    // 存储写入磁盘的每个partition的数据长度
    val elementsPerPartition = new Array[Long](numPartitions)
    // 将数据刷新到磁盘，并且释放writer资源
    def flush(): Unit = {
      val w = writer
      writer = null
      // 提交当前写入缓存中的数据到磁盘文件中
      w.commitAndClose()
      _diskBytesSpilled += spillMetrics.shuffleBytesWritten
      batchSizes.append(spillMetrics.shuffleBytesWritten)
      spillMetrics = null
      objectsWritten = 0
    }

    var success = false
    try {
      // destructiveSortedWritablePartitionedIterator传入之前自定义的排序规则，对key进行排序
      // 其实内部使用的是归并排序（TimSort），这里返回排序完数据的迭代器。
      val it = collection.destructiveSortedWritablePartitionedIterator(comparator)
      while (it.hasNext) {
        // 获取每个partition的ID
        val partitionId = it.nextPartition()
        require(partitionId >= 0 && partitionId < numPartitions,
          s"partition Id: ${partitionId} should be in the range [0, ${numPartitions})")
        // 写数据到磁盘文件中，其实会先写入缓存中，缓存大小是32K
        it.writeNext(writer)
        // 记录每个partition被写入的数据长度
        elementsPerPartition(partitionId) += 1
        // 记录所有partition数据写入长度
        objectsWritten  += 1

        // 或者判断是否达到1万条，每写入1万条，就刷新到磁盘文件中
        if (objectsWritten == serializerBatchSize) {
          flush()
          // 重新获取writer
          openWriter()
        }
      }
      if (objectsWritten > 0) {
        flush()
      } else if (writer != null) {
        val w = writer
        writer = null
        w.revertPartialWritesAndClose()
      }
      success = true
    } finally {
      if (!success) {
        if (writer != null) {
          writer.revertPartialWritesAndClose()
        }
        if (file.exists()) {
          if (!file.delete()) {
            logWarning(s"Error deleting ${file}")
          }
        }
      }
    }
    // 记录被溢写到磁盘的文件信息 
    spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
  }

  上面就是真实的溢写磁盘的操作了，首先获取一个临时的blockId和文件名，其实就是创建一个临时文件，用于保存数据；接着定义一个writer，其内部使用的是Java的BufferedOutputStream，它会将数据写入缓存，当缓存写满了，在溢写到磁盘，缓存大小是fileBufferSize（默认32K）；然后定义了刷新磁盘函数flush，将数据刷新到磁盘。接着就调用了WritablePartitionedPairCollection（它是一个trait，它的子类是上一篇讲的PartitionedAppendOnlyMap）的destructiveSortedWritablePartitionedIterator方法，它的功能主要是先对数据按照map中定义的Comparator进行排序（排序算法是TimSort一种归并排序），接着返回一个写磁盘的WritablePartitionedIterator迭代器。这个方法返回的迭代器中包含的数据已经是按照自定义的排序规则排序后的数据。
  然后依次将数据溢写到缓存，这里还会进行一个判断什么时候flush到磁盘，假设溢写的数据达到10000条，那么也溢写到磁盘。这里做个小总结，我们发现这里将数据刷新到磁盘是按照batch，来刷新的，这么设计的好处是，一次性可以刷新一定大小的数据到磁盘，而不是每次只刷新一条，减少与磁盘IO的操作，这么做可以提升性能。
  接着分析，假设数据的大小依然不够大，那么也会溢写到磁盘，这里我不太能理解，超过缓存32K或10000条，发生刷新磁盘操作;假如数据不够也会发生刷新磁盘操作，那这种情形是不是如下的这种情况：假设数据量较大，先不断的刷新，最后剩下的数据量大小不足32K或者10000条，那么就主动进行刷新操作。
  最后将溢写到磁盘的文件信息加入spills缓存中，它是后面进行磁盘临时文件合并merge的一个重要依据，从这里读取已经刷新到磁盘的临时文件信息，然后从磁盘中读取临时文件，进行合并。

writePartitionedFile将磁盘中的临时文件合并

  接着我们分析下一这个方法，它的功能就是将磁盘中的文件进行合并到一个新的临时文件中，这个方法也比较复杂，我们来进行分析：

def writePartitionedFile(
      blockId: BlockId,
      outputFile: File): Array[Long] = {
    // 记录每个partition数据的长度
    val lengths = new Array[Long](numPartitions)

    // 如果当前磁盘没有临时文件
    if (spills.isEmpty) {
      // 数据量比较小，都存储在内存中了
      // 获取数据集合，是需要聚合还是只是普通数据
      val collection = if (aggregator.isDefined) map else buffer
      // 如果需要聚合，那么返回聚合后的数据，这时数据已经按照partition ID和key，进行排序了
      val it = collection.destructiveSortedWritablePartitionedIterator(comparator)
      // 将数据溢写到磁盘文件
      while (it.hasNext) {
        // 获取BlockManager的writer
        val writer = blockManager.getDiskWriter(blockId, outputFile, serInstance, fileBufferSize,
          context.taskMetrics.shuffleWriteMetrics.get)
        val partitionId = it.nextPartition()
        // 将数据写入磁盘文件，
        while (it.hasNext && it.nextPartition() == partitionId) {
          it.writeNext(writer)
        }
        // 刷新到磁盘中，这里是真正的写磁盘。
        writer.commitAndClose()
        // 获取写入磁盘的位置切片，这个就是索引
        val segment = writer.fileSegment()
        // 记录每个partition写入的数据索引
        lengths(partitionId) = segment.length
      }
    } else {
      // 获取分区迭代器，可以访问每个分区，以及分区中的数据，这里会从临时文件中读取数据，放入内存缓存中。
      // id是partition的ID，elements是每个分区中对应的元素
      for ((id, elements) <- this.partitionedIterator) {
        // 假设分区有数据
        if (elements.hasNext) {
          // 按照分区来写数据
          // 获取writer
          val writer = blockManager.getDiskWriter(blockId, outputFile, serInstance, fileBufferSize,
            context.taskMetrics.shuffleWriteMetrics.get)
          // 将数据key和value，分别写入缓存中。
          for (elem <- elements) {
            writer.write(elem._1, elem._2)
          }
          // 刷新到磁盘
          writer.commitAndClose()
          // 获取分区偏移量
          val segment = writer.fileSegment()
          // 存储分区偏移量
          lengths(id) = segment.length
        }
      }
    }

    // 记录写入溢写到磁盘文件的数据长度
    context.taskMetrics().incMemoryBytesSpilled(memoryBytesSpilled)
    context.taskMetrics().incDiskBytesSpilled(diskBytesSpilled)
    context.internalMetricsToAccumulators(
      InternalAccumulator.PEAK_EXECUTION_MEMORY).add(peakMemoryUsedBytes)

    // 返回分区及其对应的偏移量
    lengths
  }

  这里依然区分是有溢写还是没有溢写文件，没有溢写文件就比较简单了，对分区数据进行排序，然后写入磁盘即可；假设有溢写文件，首先获取一个分区迭代器partitionedIterator，这个非常重要，它将所有临时文件的数据读到当前缓存中，进行merge操作（如果有需要的话。我觉得这里可能会发生OOM，假设临时文件比较多，那么一次性读取到内存中，很容易将Task的内存占满。），下面会详细分析这个迭代器，接着将临时文件合并后的在缓存中的数据（已经是排序后的数据），按照分区，写入到创建好的文件中，并且保存该分区的切片位置，也就是它在文件中的offset，最后返回每个分为的写入位置。
  下面我们分析分区迭代器partitionedIterator

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) {
      if (!ordering.isDefined) {
        // 假设不需要对key进行排序，那么只对partition ID进行排序，然后对partition进行聚合
        groupByPartition(collection.partitionedDestructiveSortedIterator(None))
      } else {
        // 对key和partitionID，进行排序，并按照partition进行聚合
        groupByPartition(collection.partitionedDestructiveSortedIterator(Some(keyComparator)))
      }
    } else {
      // 对溢写到磁盘文件进行合并操作
      merge(spills, collection.partitionedDestructiveSortedIterator(comparator))
    }
  }

  这里依然判断，是否需要聚合，以及是否有溢写文件，我们就不对没有溢写文件的情况进行分析了，这里假设有溢写文件，那么就调用merge方法，将溢写文件数据合并，并读到内存中。具体方法如下：

private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
      : Iterator[(Int, Iterator[Product2[K, C]])] = {

    // readers获取每个溢写文件
    val readers = spills.map(new SpillReader(_))
    val inMemBuffered = inMemory.buffered
    (0 until numPartitions).iterator.map { p =>
      // 将临时文件中的数据从磁盘读出放入内存，并获取每个partition的迭代器
      val inMemIterator = new IteratorForPartition(p, inMemBuffered)
      // 读取数据，并返回一个迭代器
      val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
      if (aggregator.isDefined) {
        // Perform partial aggregation across partitions
        // 跨分区进行聚合操作，RDD的不同分区数据进行聚合操作
        (p, mergeWithAggregation(
          iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
      } else if (ordering.isDefined) {
        // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
        // sort the elements without trying to merge them
        //  如果没有定义聚合器，那么只对数据进行排序
        (p, mergeSort(iterators, ordering.get))
      } else {
        // 假如既不排序也不聚合，那么只返回一个数据迭代器
        (p, iterators.iterator.flatten)
      }
    }
  }

  这个方法就是对溢写文件读取，并放入内存，假如需要进行聚合，这里注意使用到了aggregator的mergeCombiner，这个是跨分区进行数据聚合操作，对RDD的所有分区的数据进行聚合，那么怎么样才能读到所有分区的数据？临时文件是按照分区为单位存储的，我觉得主要是靠这里的iterators，它里面每次读取一个新分区数据的时候会将之前读取分区的放入inMemIterator中的inMemBuffered数据也添加进去（这是我个人的理解，不知道有没有问题，后续有问题再更新），这样每次都可以进行迭代，直到所有的分区数据都读取了。这里mergeWithAggregation和mergeSort就不做详细的分析了，后续我会再添加进去，感觉对这块理解不够深入。
  上述这个方法返回分区ID和聚合后的数据。
  至此，writePatitionedFile方法就分析结束了，这个方法主要就是将溢写到磁盘的临时文件进行合并，写入到一个文件中，并返回每个分区及其对应的offset，以便生成索引文件供reduce task进行数据拉取。
  我们总结一下整个SortShuffle的过程，首先它创建了一个非常重要的组件ExternalSorter，这个组件负责对RDD的partition数据进行操作，主要包括聚合、排序和溢写到磁盘，这里面包含了将数据分批次写入磁盘的设计思想，提升效率，以及这里面使用到了一个缓存结构PartitionedAppendOnlyMap，它里面定义了一个Comparator，定义了排序规则，这有点类似二次排序，这里使用的排序算法是TimSort；接着将溢写到磁盘的spills，进行聚合，这里就会对RDD的所有分区进行聚合操作了，之前的聚合操作都是分区内部，这里可以对RDD的整个分区进行聚合，最后将聚合后的数据写入到一个输出文件中，并且记录每个partition的写入位置offset，并封装成一个索引文件，一起发送到DAGScheduler的MapOutputTracker上。整个Shuffle的map端的写就算是结束了。