Spark Shuffle 模塊③ - Sort Based Shuffle write
本文為 Spark 2.0 源碼剖析,其他版本可能有所不同
自 Spark 1.2 起,Sort Based Shuffle 替代 Hash Based Shuffle 成為 Spark 默認的 Shuffle 策略。
Shuffle Map Task 會按照 key 相對應的 partition id 進行排序,對于屬于同一個 partition 的 keys 可選的進行或不進行排序。因為對于不需要排序的操作來說,這個排序是有損性能的。對于那些需要 Sort 的操作,比如 sortByKey,這個排序是由 Reducer 完成的。整個過程的如圖所示:
records 的插入、排序、spill 及最終的合并操作都封裝在 SortShuffleWriter#write(records: Iterator[Product2[K, V]]): Unit
中,其核心流程如下(注意:下面的流程針對于一個 partition):
我們自上而下進行剖析
插入 records 至 sorter
Step1:創建 sorter: ExternalSorter
創建 ExternalSorter 對象 sorter 分兩種情況:
- 需要在 map 端做 combine 操作:要提供聚合(Aggregator)和可選的 key 排序器(keyOrdering)來構造 ExternalSorter 對象,下文稱 “mapCombine ExternalSorter”
- 不需要再 map 端做 combine 操作:不需要 map 端聚合,也不需要 key 排序器來構造 ExternalSorter 對象,下文稱 “非 mapCombine ExternalSorter”
這兩種不同的 ExternalSorter 對象的不同行為會在下文中說明
Step2:調用 ExternalSorter#insertAll( records: Iterator... ) 插入一個 partition 數據
這一步也分為兩種情況
需要在 map 端做 combine
核心流程如下:
以 WordCount 中的 reduceByKey 為例:
textFile.map(word => (word, 1)).reduceByKey( (p1, p2) => p1 + p2 )
首先得到以下幾個函數變量(這一小節主要關注 insertAll 行為,暫不關注 maybeSpillCollection 方法):
- createCombiner 用于創建某個 key 對應的初始聚合值。在這里為
(p: Int) => p
,即 key 第一次出現的 record 的 value 即為初始聚合值 - mergeValue 用于在 map 端將同一個 partition 中 key 相同的 value 合并到一起。在這里,為
(p1: Int, p2: Int) => p1 + p2
- mergeCombiners 用于在 reduce 端將不同 partition 中 key 相同的 value 合并到一起。在這里,為
(p1: Int, p2: Int) => p1 + p2
之后,遍歷 records 并調用 map.changeValue((getPartition(kv._1), kv._1), update)
進行 combine,這里的 map 是 PartitionedAppendOnlyMap
類型,是一個只能添加數據不能移除數據的 HashMap,該 HashMap 最終也是基于數組來實現,數組中的 a[i] = key i, a[i+1]=value i
,當添加的元素超過一定的指標時,數組會進行 grow,容量翻倍,同時所有的 key、value 會進行 rehash 重新分配位置。combine 的過程如 update 定義,對于已經有聚合結果的新值進行合并,將沒有聚合結果的新值作為初始的聚合結果值。
在每 insert 一條 record 之后都會調用 maybeSpillCollection(usingMap = true)
來判斷是否需要進行 spill,這將在后面的內容中介紹。
不需要在 map 端做 combine
核心流程如下:
這里的 buffer 是 PartitionedPairBuffer
類型,是一個只支持添加的 PairBuffer,底層實現是一個數組,存儲的 pair,即 key、value 對,key 的類型為 (Int, K)
,即 Partition Id 和 key。同一個 k、v 對被放在數組相鄰的位置,與 PartitionedAppendOnlyMap
相同。
每插入一條 record 會調用 maybeSpillCollection(usingMap = false)
方法,這同樣將在后文中說明。
spill 內存 buffer 至文件
主要流程如下:
第一個值得關注的點是:map 和 buffer 是如何估算占用了多少內存?實現如下:
val extrapolatedDelta = bytesPerUpdate * (numUpdates - samples.last.numUpdates)
(samples.last.size + extrapolatedDelta).toLong
即:以 bytePerUpdate 作為最近平均每次更新時的大小,用當前的 update 次數減去最后一個Sample 的 update 次數,然后乘以 bytePerUpdate,結果加上最后一個 Sample 記錄的大小
另一個重要的點是 maybeSpill 方法,該方法首先確定是否有必要進行 spill,如果有,則將當前內存集合中溢寫到磁盤上。在執行溢寫前會嘗試申請更多內存。有兩種情況會需要進行 spill:
- 當前集合包含的 records 數超過
spark.shuffle.spill.numElementsForceSpillThreshold
指定值,該值默認大小為Long.MaxValue
- 這種情況復雜些,其判斷流程如下:
主要是在進行 spill 之前會嘗試申請更多的內存來存放 records 來避免 spill。
若需要進行 spill,則首先會進行 spill 操作,然后釋放已 spill 的集合對應的內存,釋放主要是釋放 execution pool 內存以供其他 tasks 使用,并將 myMemoryThreshold
賦值為初始值 spark.shuffle.spill.initialMemoryThreshold
對應值,即初始值
spill 操作
spill 的操作要考慮到之后要對之后生成的 spill 文件做 merge,因為最終一個 Shuffle Map Task 只生成一個輸出文件和 Index 文件。
如果是需要做 map 端 combine,spill 時會對 map 中的數據先按 partition id 進行排序,若也提供了 key comparator,則會對屬于同一個 partition 的 records 按 key 進行排序。做完排序后,先進行序列化再寫入磁盤文件。
如果是不需要做 map 端的 combine,則只需對 buffer 按 partition id 進行排序即可,不需要對同一partition 的 records 按 key 進行排序。排序后,同樣先序列化,再寫入磁盤文件。
之后做 merge 時,使用 SpillReader 來讀取 spill 數據又要先反序列化,再做最終排序,再寫入最終文件,這一過程是 shuffle 過程中消耗比較大的一部分。
合并數據
合并的核心流程如下,由 ExternalSorter#writePartitionedFile(...)
方法實現
其中,最關鍵的 merge 流程如下:
- 為每個 spill 出來的文件生成一個 reader: SpillReader,得到
readers: Seq[SpillReader]
(reader 讀取 spilled 文件要先反序列化) - 將內存集合進行 buffered,得到 inMemBuffered
- 針對每個 partition p,執行以下操作:
- 取出 inMemBuffered 中屬于 partition p 的 records,得到迭代器 inMemIterator
- 使用 readers 取出屬于 partition p 的 records 對應的迭代器 Seq 與上一步中得到的 inMemIterator 合并得到最終的迭代器序列
- 如果定義了 map 端聚合,則先對上一步得到的迭代器序列中的 records 進行聚合,若還定義了 key comparator,則使用該 comparator 對 records 進行排序
- 若沒有定義 map 端聚合但定義了 key comparator,則不做聚合而直接對 records 進行排序
- 若既沒有定義 map 端聚合也沒有定義 key comparator,則直接返回總體 records 對應的迭代器