大數據：Spark Shuffle（一）ShuffleWrite:Executor如何將Shuffle的結果進行歸併寫到數據文件中去

1. 前序

關於Executor如何運行算子，請參考前面博文：大數據：Spark Core（四）用LogQuery的例子來說明Executor是如何運算RDD的算子，當Executor進行reduce運算的時候，生成運算結果的臨時Shuffle數據，並保存在磁盤中，被最後的Action算子調用，而這個階段就是在ShuffleMapTask裏執行的。

前面博客中也提到了，用什麼ShuffleWrite是由ShuffleHandler來決定的，在這篇博客裏主要介紹最常見的SortShuffleWrite的核心算法ExternalSorter.

2. 結構AppendOnlyMap

在前面博客中介紹了SortShuffleWriter調用ExternalSorter.insertAll進行數據插入和數據合併的，ExternalSorted裏使用了PartitionedAppendOnlyMap作爲數據的存儲方式

先來看PartitionedAppendOnlyMap的結構

雖然名字爲Map,但是在這裏和常見的Map的結構並不太一樣，裏面並沒有使用鏈表結果保存相同的hash值的key，當插入的key的hashcode相同的時但key不相同，會通過i的疊加一直找到數組裏空閒的位置。

這裏有幾個注意點：

• Key 注意這裏的Key並不是通過Map裏拆分的Key, 而是Tuple2(PartitionId,Key)，由分片的段和key組合的聯合key
• 如何計算PartitionId? 這是由Partitioner來決定的

2.1 Partitioner

Partitioner的方法

abstract class Partitioner extends Serializable {
  def numPartitions: Int
  def getPartition(key: Any): Int
}

通過調用getPartition方法找到對應的partition相應的塊，而常用的是HashPartitioner

 def getPartition(key: Any): Int = key match {
    case null => 0
    case _ => Utils.nonNegativeMod(key.hashCode, numPartitions)
  }

計算 key的hashCode，進行總的分片數求餘，分配到對應的片區

3. Spill

在大數據的情況下進行歸併，由於合併的數據量非常大，僅僅使用AppendOnlyMap進行數據的歸併內存顯然是不足夠的，在這種情況下需要對講內存裏的已經歸併的數據刷到磁盤上避免OOM的風險。

控制Spill到磁盤的閥值

• 內存：雖然Java的堆內存管理是由JVM虛擬機管控，但是Spark自己實現了一個簡單的但不精準的內存管理，內存的申請在TaskMemoryManager裏進行管理

 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
    }

在每添加32個元素的時候，檢查一下當前的內存狀況，currentMemory是Map當前大概使用的內存，myMemoryThreshold是可以使用的內存址，初始的時候受參數控制：

spark.shuffle.spill.initialMemoryThreshold

爲何要嘗試申請1倍的當前內存？AppendOnlyMap的每次擴容是1倍數組

• 數據的數量：有的時候每條數據量比較小，但是數據的數量非常大，爲了避免在AppendOnlyMap裏有大量的數據，在Spill的時候同時還可以使用數量的控制：

spark.shuffle.spill.numElementsForceSpillThreshold

3.1 如何Spill?

當從AppendOnlyMap到SpilledFile磁盤總共有3個過程

1. 整理數組，將數組裏的不存在KV的空間移除
2. 按照區塊排序，對同一區塊裏的Key使用TimeSort進行排序，TimeSort不在此處討論
3. Spill到文件的時候，只是保存了序列化了Key,Value並沒有保存Key的區塊信息，但在SpilledFile的對象中有記錄每個partitionkey的數量的數組

SpilledFile的命名：temp_shuffle_UUID

4. 生成ShuffleWrite的數據文件

在3章節的時候，有沒有考慮過爲何要排序完才Spill到臨時文件中？

Spark中是不要求在reduce端進行排序的，生成Shuffle的結果文件並不要求排序，但是因爲Spill到文件中後，有可能相同的Key會分佈在不同的文件中，所以需要對不同的文件進行相同的Key的值的計算。如果Spill到文件是亂序的，那代表在最後生成Shuffle結果的時候，還是要Load所有文件才能確定哪些Key是重複的需要做合併，這樣依然面對着內存不夠的情況。

生成Shuffle文件過程實際上就是個外排序的過程。

• 首先對AppendOnlyMap進行歸併，排序
• 開始對同一區塊的進行歸併
• 將AppendOnlyMap，SpilledFile的文件進行優先級的Queue的迭代，每次迭代出所有Queue中一個最小的Key，最小的Key就是HashCode最小

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

• 當找到一個最小的Key的時候，並不能保存到ShuffleWrite文件中，因爲有可能存在相同的最小的key，所以還需要在迭代找到下一個最小的Key，如果key的hashcode相同的時候，要進行相同的Key進行合併（因爲Key的排序是依賴於HashCode的大小，所以相同的最小的Key代表的是HashCode相同的Key），如果不同則保存成相同HashCode的數組，進行下一次的優先queue的查找，直到找到的Key的hashcode大於最小的Key結束

   if (!hasNext) {
            throw new NoSuchElementException
          }
          keys.clear()
          combiners.clear()
          val firstPair = sorted.next()
          keys += firstPair._1
          combiners += firstPair._2
          val key = firstPair._1
          while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
            val pair = sorted.next()
            var i = 0
            var foundKey = false
            while (i < keys.size && !foundKey) {
              if (keys(i) == pair._1) {
                combiners(i) = mergeCombiners(combiners(i), pair._2)
                foundKey = true
              }
              i += 1
            }
            if (!foundKey) {
              keys += pair._1
              combiners += pair._2
            }
          }

• 將k,v內容寫到shufflewrite的文件Shuffle_shuffleId_mapId_reduceId.data中去

• 重複前面的行爲直到所有的key被迭代結束
• 前面的歸併是以區塊(Partition)爲單位的，而data的文件裏並沒有保存區塊的相關信息，但在每迭代完一個Partition的時候(SpilledFile文件裏面也沒有Partition的信息，但是是通過SpilledFile結構中的numPartition的數量來判斷Partition的數據是否已經讀完)，會生成一個Segement，Segement 裏記錄了這個塊保存在data文件裏的長度
• 最後生成Shuffle_shuffleId_mapId_reduceId.index文件，文件裏記錄了每個Partition在data文件中的位移

這樣一個完整的Shuffle結果寫入data的邏輯執行完了

5 總結

• 使用AppendOnlyMap數據結構進行輸入數據的合併計算
• 輸入的數據是進行分區合併計算，分區的方式是由Partitioner決定的
• 當內存不夠的時候，會進行相同區塊下的數據整理排序，Spill到臨時文件temp_shuffle_UUID
• 最後對所有的數據集合（AppendOnlyMap裏的數據和多個Spill的臨時文件）進行區塊的數據合併
• 生成Shuffle_shuffleId_mapId_reduceId.data 分區的數據文件，Shuffle_shuffleId_mapId_reduceId.index記錄分區的位置