Tungsten-sort 算不得一個(gè)全新的shuffle 方案，它在特定場(chǎng)景下基于類似現(xiàn)有的Sort Based Shuffle處理流程，對(duì)內(nèi)存/CPU/Cache使用做了非常大的優(yōu)化。帶來高效的同時(shí)，也就限定了自己的使用場(chǎng)景。如果Tungsten-sort 發(fā)現(xiàn)自己無法處理，則會(huì)自動(dòng)使用 Sort Based Shuffle進(jìn)行處理。

前言

看這篇文章前，建議你先簡(jiǎn)單看看Spark Sort Based Shuffle內(nèi)存分析。

Tungsten 中文是鎢絲的意思。 Tungsten Project 是 Databricks 公司提出的對(duì)Spark優(yōu)化內(nèi)存和CPU使用的計(jì)劃，該計(jì)劃初期似乎對(duì)Spark SQL優(yōu)化的最多。不過部分RDD API 還有Shuffle也因此受益。

簡(jiǎn)述

Tungsten-sort優(yōu)化點(diǎn)主要在三個(gè)方面:

1. 直接在serialized binary data上sort而不是java objects，減少了memory的開銷和GC的overhead。
2. 提供cache-efficient sorter，使用一個(gè)8bytes的指針，把排序轉(zhuǎn)化成了一個(gè)指針數(shù)組的排序。
3. spill的merge過程也無需反序列化即可完成

這些優(yōu)化的實(shí)現(xiàn)導(dǎo)致引入了一個(gè)新的內(nèi)存管理模型，類似OS的Page，對(duì)應(yīng)的實(shí)際數(shù)據(jù)結(jié)構(gòu)為MemoryBlock,支持off-heap 以及 in-heap 兩種模式。為了能夠?qū)ecord 在這些MemoryBlock進(jìn)行定位，引入了Pointer（指針）的概念。

如果你還記得Sort Based Shuffle里存儲(chǔ)數(shù)據(jù)的對(duì)象PartitionedAppendOnlyMap,這是一個(gè)放在JVM heap里普通對(duì)象，在Tungsten-sort中，他被替換成了類似操作系統(tǒng)內(nèi)存頁的對(duì)象。如果你無法申請(qǐng)到新的Page,這個(gè)時(shí)候就要執(zhí)行spill操作，也就是寫入到磁盤的操作。具體觸發(fā)條件，和Sort Based Shuffle 也是類似的。

開啟條件

Spark 默認(rèn)開啟的是Sort Based Shuffle,想要打開Tungsten-sort ,請(qǐng)?jiān)O(shè)置

spark.shuffle.manager=tungsten-sort

對(duì)應(yīng)的實(shí)現(xiàn)類是：

org.apache.spark.shuffle.unsafe.UnsafeShuffleManager

名字的來源是因?yàn)槭褂昧舜罅縅DK Sun Unsafe API。

當(dāng)且僅當(dāng)下面條件都滿足時(shí)，才會(huì)使用新的Shuffle方式：

• Shuffle dependency 不能帶有aggregation 或者輸出需要排序
• Shuffle 的序列化器需要是 KryoSerializer 或者 Spark SQL's 自定義的一些序列化方式.
• Shuffle 文件的數(shù)量不能大于 16777216
• 序列化時(shí)，單條記錄不能大于 128 MB

可以看到，能使用的條件還是挺苛刻的。

這些限制來源于哪里

參看如下代碼，page的大?。?/p>

this.pageSizeBytes = (int) Math.min(  
                PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES, 
                shuffleMemoryManager.pageSizeBytes());

這就保證了頁大小不超過PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES 的值，該值就被定義成了128M。

而產(chǎn)生這個(gè)限制的具體設(shè)計(jì)原因，我們還要仔細(xì)分析下Tungsten的內(nèi)存模型:

來源于：https://github.com/hustnn/TungstenSecret/tree/master

這張圖其實(shí)畫的是 on-heap 的內(nèi)存邏輯圖，其中 #Page 部分為13bit, Offset 為51bit,你會(huì)發(fā)現(xiàn) 2^51 >>128M的。但是在Shuffle的過程中，對(duì)51bit 做了壓縮，使用了27bit,具體如下：

 [24 bit partition number][13 bit memory page number][27 bit offset in page]

這里預(yù)留出的24bit給了partition number,為了后面的排序用。上面的好幾個(gè)限制其實(shí)都是因?yàn)檫@個(gè)指針引起的：

1. 一個(gè)是partition 的限制，前面的數(shù)字 16777216 就是來源于partition number 使用24bit 表示的。
2. 第二個(gè)是page number
3. 第三個(gè)是偏移量，最大能表示到2^27=128M。那一個(gè)task 能管理到的內(nèi)存是受限于這個(gè)指針的，最多是 2^13 * 128M 也就是1TB左右。

有了這個(gè)指針，我們就可以定位和管理到off-heap 或者 on-heap里的內(nèi)存了。這個(gè)模型還是很漂亮的，內(nèi)存管理也非常高效，記得之前的預(yù)估PartitionedAppendOnlyMap的內(nèi)存是非常困難的，但是通過現(xiàn)在的內(nèi)存管理機(jī)制，是非常快速并且精確的。

對(duì)于第一個(gè)限制，那是因?yàn)楹罄m(xù)Shuffle Write的sort 部分，只對(duì)前面24bit的partiton number 進(jìn)行排序，key的值沒有被編碼到這個(gè)指針，所以沒辦法進(jìn)行ordering

同時(shí)，因?yàn)檎麄€(gè)過程是追求不反序列化的，所以不能做aggregation。

Shuffle Write

核心類：

 org.apache.spark.shuffle.unsafe.UnsafeShuffleWriter

數(shù)據(jù)會(huì)通過 UnsafeShuffleExternalSorter.insertRecordIntoSorter 一條一條寫入到 serOutputStream 序列化輸出流。

這里消耗內(nèi)存的地方是

 serBuffer = new MyByteArrayOutputStream(1024 * 1024)

默認(rèn)是1M,類似于Sort Based Shuffle 中的ExternalSorter，在Tungsten Sort 對(duì)應(yīng)的為UnsafeShuffleExternalSorter,記錄序列化后就通過sorter.insertRecord方法放到sorter里去了。

這里sorter 負(fù)責(zé)申請(qǐng)Page,釋放Page,判斷是否要進(jìn)行spill都這個(gè)類里完成。代碼的架子其實(shí)和Sort Based 是一樣的。

圖片來源:https://www.ibm.com/developerworks/cn/opensource/os-cn-spark-core/

(另外，值得注意的是，這張圖里進(jìn)行spill操作的同時(shí)檢查內(nèi)存可用而導(dǎo)致的Exeception 的bug 已經(jīng)在1.5.1版本被修復(fù)了,忽略那條路徑)

內(nèi)存是否充足的條件依然shuffleMemoryManager 來決定，也就是所有task shuffle 申請(qǐng)的Page內(nèi)存總和不能大于下面的值：

 ExecutorHeapMemeory * 0.2 * 0.8

上面的數(shù)字可通過下面兩個(gè)配置來更改：

spark.shuffle.memoryFraction=0.2
spark.shuffle.safetyFraction=0.8

UnsafeShuffleExternalSorter 負(fù)責(zé)申請(qǐng)內(nèi)存，并且會(huì)生成該條記錄最后的邏輯地址，也就前面提到的 Pointer。

接著Record 會(huì)繼續(xù)流轉(zhuǎn)到UnsafeShuffleInMemorySorter中，這個(gè)對(duì)象維護(hù)了一個(gè)指針數(shù)組：

  private long[] pointerArray;

數(shù)組的初始大小為 4096，后續(xù)如果不夠了，則按每次兩倍大小進(jìn)行擴(kuò)充。

假設(shè)100萬條記錄，那么該數(shù)組大約是8M 左右，所以其實(shí)還是很小的。一旦spill后該UnsafeShuffleInMemorySorter就會(huì)被賦為null,被回收掉。

我們回過頭來看spill,其實(shí)邏輯上也異常簡(jiǎn)單了，UnsafeShuffleInMemorySorter 會(huì)返回一個(gè)迭代器，該迭代器粒度每個(gè)元素就是一個(gè)指針，然后到根據(jù)該指針可以拿到真實(shí)的record,然后寫入到磁盤，因?yàn)檫@些record 在一開始進(jìn)入UnsafeShuffleExternalSorter 就已經(jīng)被序列化了，所以在這里就純粹變成寫字節(jié)數(shù)組了。形成的結(jié)構(gòu)依然和Sort Based Shuffle 一致，一個(gè)文件里不同的partiton的數(shù)據(jù)用fileSegment來表示，對(duì)應(yīng)的信息存在一個(gè)index文件里。

另外寫文件的時(shí)候也需要一個(gè) buffer :

 spark.shuffle.file.buffer = 32k

另外從內(nèi)存里拿到數(shù)據(jù)放到DiskWriter,這中間還要有個(gè)中轉(zhuǎn)，是通過

 final byte[] writeBuffer = new byte[DISK_WRITE_BUFFER_SIZE=1024 * 1024];

來完成的，都是內(nèi)存，所以很快。

Task結(jié)束前，我們要做一次mergeSpills操作，然后形成一個(gè)shuffle 文件。這里面其實(shí)也挺復(fù)雜的，
如果開啟了

 `spark.shuffle.unsafe.fastMergeEnabled=true`

并且沒有開啟

`spark.shuffle.compress=true`

或者壓縮方式為：

 LZFCompressionCodec

則可以非常高效的進(jìn)行合并,叫做transferTo。不過無論是什么合并，都不需要進(jìn)行反序列化。

Shuffle Read

Shuffle Read 完全復(fù)用HashShuffleReader,具體參看 Sort-Based Shuffle。

總結(jié)

我個(gè)人感覺，Tungsten-sort 實(shí)現(xiàn)了內(nèi)存的自主管理，管理方式模擬了操作系統(tǒng)的方式，通過Page可以使得大量的record被順序存儲(chǔ)在內(nèi)存，整個(gè)shuffle write 排序的過程只需要對(duì)指針進(jìn)行運(yùn)算(二進(jìn)制排序)，并且無需反序列化，整個(gè)過程非常高效，對(duì)于減少GC,提高內(nèi)存訪問效率，提高CPU使用效率確實(shí)帶來了明顯的提升。