Shuffle分类

一个作业经过spark的DAGSchedule调度器划分为多个stage，同时有些下游的stage依赖上游的stage，这样会导致上游的stage做map的工作，下游的stage做reduce的工作。而上下游stage就是通过shuffle连接在一起的。shuffle分为shuffleWriter和ShulleReader，writer即为上游的stage，reader为下游的stage。
　　在这里我们首先介绍一下ShuffleWriter的分类。如下：
　　1.BypassMergeSortShuffleWriter。和Hash Shuffle的实现基本相同，区别在于map task输出汇总一个文件，同时还会产生一个index file。
　　2.UnsafeShuffleWriter。使用Java Unsafe直接操作内存，避免Java对象多余的开销和GC 延迟，效率高。
　　3.SortShullfleWriter。sortShulleWriter和shuffle不同的在于map端的sort。
　　以上三种shuffleWriter都具有各自的应用场景。分别如下：
　　1.没有map端聚合操作且RDD的partition分区数小于200个，使用BypassMergerSortShuffleWriter。
　　2.没有map端聚合，RDD的partitions分区数小于16777216且Serializer支持relocation，使用UnsafeShuffleWriter。
　　3.（1）和（2）不满足使用SortShuffle
上面介绍了ShullfleWriter的分类，下面将详细介绍SortShuffleWriter的具体工作原理。

SortShuffleWriter工作原理

预备知识

org.apache.spark.util.collection.AppendOnlyMap
org.apache.spark.util.collection.PartitionedPairBuffer
org.apache.spark.util.collection.TimSorter
org.apache.spark.shuffle.sort.SortShuffleWriter
　　其中，AppendOnlyMap和PartitionedPairBuffer是shuffle过程中使用的两个数据结构，但是底层都是基于Array实现的。当需要map端的combine或者sort时，使用的是AppendOnlyMap结构；否则，使用的是PartitionedPairBuffer结构。而TimSorter主要是封装了归并排序，用于内存中的数据排序。
　　首先我们看一下SortShullfleWriter源码，如下

/** Write a bunch of records to this task's output */
  override def write(records: Iterator[Product2[K, V]]): Unit = {
    sorter = if (dep.mapSideCombine) {
    #需要map端的聚合
      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端的聚合或者排序
      // 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)
    }
    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).
    val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
    val tmp = Utils.tempFileWith(output)
    try {
      val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
      val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
      shuffleBlockResolver.writeIndexFileAndCommit(dep.shuffleId, mapId, partitionLengths, tmp)
      mapStatus = MapStatus(blockManager.shuffleServerId, partitionLengths)
    } finally {
      if (tmp.exists() && !tmp.delete()) {
        logError(s"Error while deleting temp file ${tmp.getAbsolutePath}")
      }
    }
  }

接下来看一下ExternalSorter类，具体在org.apache.spark.util.collection.ExternalSorter。下面转到了ExternalSort的insert方法中，该方法做了shuffleWriter的绝大部分工作 。我们先看一下insert的源代码，然后介绍该方法所做的主要工作。

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
    if (shouldCombine) {
      // 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)
      }
      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()
        buffer.insert(getPartition(kv._1), kv._1, kv._2.asInstanceOf[C])
       #判断是否会发生溢出
        maybeSpillCollection(usingMap = false)
      }
    }
  }

具体的工作如下：
(1)根据是否需要map端做shouldCombine，决定使用map还是buffer。由于使用的是map，在进行join的过程中，不停地遍历map，当发现相同的时候，直接更新map的value即可。在这里会存在一个问题，就是假如key存在不同的disk的文件上，该如何，这时候当调用Iterator时，会返回一个基于内存和disk的多路归并迭代器。这个迭代器，每次在调用next 方法的时候是基于优先级队列，也就是每个迭代器最小hash值作为比较对象的堆结构，寻找最小的hash值且key值相等的所有元素（因为我们每个map 都是排序过的，所以这总能实现），进行merge，将所有符合要求的元素merge完成后返回。这样便完成了最终的聚合操作。
(2)判断是否为溢出
那么，接下来的重点是看缓存是否会发生溢出，主要在maybeSpillCollection方法中。下面是该方法的源码。

 /**
   * Spill the current in-memory collection to disk if needed.
   *
   * @param usingMap whether we're using a map or buffer as our current in-memory collection
   */
  private def maybeSpillCollection(usingMap: Boolean): Unit = {
    var estimatedSize = 0L
    if (usingMap) {
      estimatedSize = map.estimateSize()
      if (maybeSpill(map, estimatedSize)) {
        map = new PartitionedAppendOnlyMap[K, C]
      }
    } else {
      estimatedSize = buffer.estimateSize()
      if (maybeSpill(buffer, estimatedSize)) {
        buffer = new PartitionedPairBuffer[K, C]
      }
    }

    if (estimatedSize > _peakMemoryUsedBytes) {
      _peakMemoryUsedBytes = estimatedSize
    }
  }

在该方法中首先需要评估一下是否会发生OOM。由于PartitionedAppendOnlyMap或者PartitionedPairBuffer都继承了SizeTracker类，所以不管是buffer还是map都是使用了SizeTracker中的estimateSize方法。在处理每条record时，都会进行一次estimateSize，评估出已经使用了的缓存大小。下面是该方法

/**
   * Estimate the current size of the collection in bytes. O(1) time.
   */
  def estimateSize(): Long = {
    assert(samples.nonEmpty)
    val extrapolatedDelta = bytesPerUpdate * (numUpdates - samples.last.numUpdates)
    (samples.last.size + extrapolatedDelta).toLong
  }

其中bytesPerUpdate是两次采样之间每次update这个集合，size增长了多少。也就是说每操作集合一下，集合的size增长了多少。

bytesPerUpdate=(latest.size - previous.size).toDouble / (latest.numUpdates - previous.numUpdates)

然后以bytePerUpdate作为最近平均每次更新时的bytePerUpdate，用当前的update次数减去最后一个Sample的update次数，然后乘以bytePerUpdate，结果加上最后一个Sample记录的大小。

val extrapolatedDelta = bytesPerUpdate * (numUpdates - samples.last.numUpdates)
    (samples.last.size + extrapolatedDelta).toLong

我们回到maybeSpillCollection方法中，当评估出map或者buffer的大小以后，直接调用方法maybeSpill，判断是否发生溢写操作。该方法主要存在org.apache.spark.util.collection.Spillable类中，方法如下

/**
   * Spills the current in-memory collection to disk if needed. Attempts to acquire more
   * memory before spilling.
   *
   * @param collection collection to spill to disk
   * @param currentMemory estimated size of the collection in bytes
   * @return true if `collection` was spilled to disk; false otherwise
   */
  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
  }

首先是传入这个集合以及该集合的使用大小，在这个方法中，我们每32个记录判断一次。当当前的集合大小大于阈值时，需要重新扩大能存的大小，同时把集合中的数据写入磁盘上。
　　接下来是调用ExternalSort的spill方法，该方法的主要作用是集合写到磁盘，并且把添加到一个spills数组中。

/**
   * Spill our in-memory collection to a sorted file that we can merge later.
   * We add this file into `spilledFiles` to find it later.
   *
   * @param collection whichever collection we're using (map or buffer)
   */
  override protected[this] def spill(collection: WritablePartitionedPairCollection[K, C]): Unit = {
    val inMemoryIterator = collection.destructiveSortedWritablePartitionedIterator(comparator)
    val spillFile = spillMemoryIteratorToDisk(inMemoryIterator)
    spills += spillFile
  }

接下来就是meger这些磁盘文件了，在这里我们把文件记录在spills数组中了。

merge过程

当我们请求一个iterator或者文件时，会将所有的SpilledFile和在内存当中未进行溢写的数据进行合并。
　　从上面可以看出接下来执行的是ExternalSort的writePartitionedFile方法，其主要作用就是把数据写入到磁盘，并返回一个索引。

/**
   * Write all the data added into this ExternalSorter into a file in the disk store. This is
   * called by the SortShuffleWriter.
   *
   * @param blockId block ID to write to. The index file will be blockId.name + ".index".
   * @return array of lengths, in bytes, of each partition of the file (used by map output tracker)
   */
  def writePartitionedFile(
      blockId: BlockId,
      outputFile: File): Array[Long] = {

    // Track location of each range in the output file
    val lengths = new Array[Long](numPartitions)
    val writer = blockManager.getDiskWriter(blockId, outputFile, serInstance, fileBufferSize,
      context.taskMetrics().shuffleWriteMetrics)
      #所有的数据都存储在内存中，直接写入文件
    if (spills.isEmpty) {
      // Case where we only have in-memory data
      val collection = if (aggregator.isDefined) map else buffer
      val it = collection.destructiveSortedWritablePartitionedIterator(comparator)
      while (it.hasNext) {
        val partitionId = it.nextPartition()
        while (it.hasNext && it.nextPartition() == partitionId) {
          it.writeNext(writer)
        }
        val segment = writer.commitAndGet()
        lengths(partitionId) = segment.length
      }
    } else {
      // We must perform merge-sort; get an iterator by partition and write everything directly.
    # 当发生了溢写操作
      for ((id, elements) <- this.partitionedIterator) {
        if (elements.hasNext) {
          for (elem <- elements) {
            writer.write(elem._1, elem._2)
          }
          val segment = writer.commitAndGet()
          lengths(id) = segment.length
        }
      }
    }

    writer.close()
    context.taskMetrics().incMemoryBytesSpilled(memoryBytesSpilled)
    context.taskMetrics().incDiskBytesSpilled(diskBytesSpilled)
    context.taskMetrics().incPeakExecutionMemory(peakMemoryUsedBytes)

    lengths
  }

当发生了溢写操作，这时候需要合并多个小文件，这时候主要是通过调用了partitionedIterator方法，在partitionedIterator主要调用了是merge方法，下面看一下merge方法。

/**
   * Merge a sequence of sorted files, giving an iterator over partitions and then over elements
   * inside each partition. This can be used to either write out a new file or return data to
   * the user.
   *
   * Returns an iterator over all the data written to this object, grouped by partition. For each
   * partition we then have an iterator over its contents, and these are expected to be accessed
   * in order (you can't "skip ahead" to one partition without reading the previous one).
   * Guaranteed to return a key-value pair for each partition, in order of partition ID.
   */
  private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
      : Iterator[(Int, Iterator[Product2[K, C]])] = {
     #每一个spilledFile创建一个reader
    val readers = spills.map(new SpillReader(_))
    val inMemBuffered = inMemory.buffered
    (0 until numPartitions).iterator.map { p =>
      val inMemIterator = new IteratorForPartition(p, inMemBuffered)
        #合并内存与spill文件
      val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
      if (aggregator.isDefined) {
        ##聚合
        // Perform partial aggregation across partitions
        (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)
      }
    }
  }

该方法将属于同一个reduce端的partition的内存数据和spill文件数据合并起来，再进行聚合排序（有需要的话），最后返回（reduce对应的partitionId，该分区数据迭代器）。将数据merge-sort后写入最终的文件后，需要将每个partition的偏移量持久化到文件以供后续每个reduce根据偏移量获取自己的数据，写偏移量的逻辑很简单，就是根据前面得到的partition长度的数组将偏移量写到index文件中。接下来看一下mergeSort方法。

 /**
   * Merge-sort a sequence of (K, C) iterators using a given a comparator for the keys.
   */
  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]]
    val heap = new mutable.PriorityQueue[Iter]()(new Ordering[Iter] {
      // Use the reverse of comparator.compare because PriorityQueue dequeues the max
      override def compare(x: Iter, y: Iter): Int = -comparator.compare(x.head._1, y.head._1)
    })
    heap.enqueue(bufferedIters: _*)  // Will contain only the iterators with hasNext = true
    new Iterator[Product2[K, C]] {
      override def hasNext: Boolean = !heap.isEmpty

      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
      }
    }
  }

该方法主要是利用了java collection的优先队列，这样每次遍历的时候都是取最小或者最大的一个，这样就可以把spill文件按照(partitionid，key)排序了，最后生成一个文件，同时会生成一个索引文件，索引文件是由上面的writePartitionedFile方法生成的。
　　　下一节，我们将介绍基于sort shuffle的shufflereader。