作者介紹：黃建博，雲計算領域技術開發工程師；金靈， Shopee 軟件研發工程師。
他們的隊伍 huang-b 在性能競賽中斬獲一等獎，本文將介紹 Shuffle 優化 TiDB 算子項目的設計與實踐過程。
在我們往常的印象中，分與合是一對矛盾的概念，但是這次比賽留給我們隊伍一個很深刻的印象是，分與合是一對互相促進的矛盾，只有乾淨利落地分解，才能高效地合併。這種印象一方面來自於我們的參賽思路，我們選擇的方向是使用 Shuffle 操作將算子的數據源拆分爲多個獨立的分區，然後通過並行計算來提升整體的吞吐，優化的過程就是尋找更合適的分解方式，追求更好的擴展性和計算性能，相關技術細節在正文中有詳細介紹。另一方面，則是來自於我們的參賽體驗。
我們的隊伍有兩位成員一位顧問，分散在國內三座不同的城市，比賽從開始到結束，我們沒有機會線下交流過，全都是以 slack 和文檔的形式進行合作，我們開玩笑說，在做一個分佈式的比賽的同時，我們也是一支分佈式的隊伍。在這樣一種受限的條件下，有兩個因素對隊伍的高效合作起了關鍵作用：其一是我們的顧問對任務做了乾淨利落的拆分，同樣的思路應用到兩個不同類型的算子上，使得我們可以花開兩朵，各表一枝；其二就是 TiDB 整體高內聚低耦合的設計，從縱向的角度講，就是 TiDB 清晰的分層設計使得我們的優化可以只關注解析器和執行器層面，而不必深入更底層的 TiKV 存儲，從橫向的角度講，就是 TiDB 對分治和多態的充分實踐使得我們可以只關注被優化的算子，而不必擔心對其它算子產生副作用。
技術背景
我們的優化思路是使用 Shuffle 算子來實現 MergeJoin 算子和 StreamAggregation 算子的並行化。Shuffle 算子最先在 PR https://github.com/pingcap/tidb/pull/14238 中引入，用於並行化 Window 算子。圖 1 展示的就是 Window 算子的並行化過程。
圖中左側是串行的 Window 算子，因爲 Window 算子要求輸入數據有序，所以在數據源和 Window 之間通常有一個 Sort 算子。圖中右側展示的是對應的 Shuffle 算子，爲了完成並行計算，Window 算子和 Sort 算子都被複制了多份，每一份與一 ShuffleWorker 相對應，從數據源流入的數據由 Splitter 按哈希值拆分爲獨立的數據分區，發往不同的 ShuffleWorker ，最終各個 Window 算子的結果彙總後輸出。圖中所畫箭頭就是數據流動的方向，其中數據的分發和結果的彙總是通過 go channel 來實現的，其它數據流動都是父節點通過調用子節點的 Next 方法來獲取的。圖中虛線表示啓動的協程，每個 ShuffleWorker 都會啓動一個協程來完成自身的運算，同時 Splitter 也會啓動一個協程來完成數據的分發。
題目鏈接：
1. ShuffleMergeJoin:https://github.com/pingcap/tidb/issues/14441
2. ShuffleStreamAgg:https://github.com/pingcap/tidb/issues/20651
圖 1 Window 算子並行化
ShuffleMergeJoin
擴展 Shuffle 算子
對 MergeJoin 做並行優化，是不是簡單套用 ShuffleWindow 的框架就可以了呢？不是的，MergeJoin 算子與上文的 Window 算子不同，MergeJoin 需要兩個數據源。那現在的 Shuffle 實現能不能讓每個並行算子對應兩個 ShuffleWorker ，進而對應兩個數據源呢？答案也是不可以，因爲前文提到的 Shuffle 實現把數據分區和計算並行這兩個功能過度耦合在一起了，這種過度耦合使得它無法支持兩個數據源。下面我們對這個問題作具體說明。
過度耦合指的是 ShuffleWorker 充當的角色太多，它既是數據流動的一環，同時也是計算並行的基本單元，於是帶來了這兩個問題：
1. 因爲 ShuffleWorker 時數據分區，所以並行後的每個 MergeJoin 需要兩個 ShuffleWorker 來接收來自兩個數據源的數據，但是 ShuffleWorker 同時又是計算並行的基本單元，於是有 n 個 MergeJoin 算子就會出現 2n 個協程，同一個 MergeJoin 算子的兩個協程還會出現數據競爭。
2. 控制邏輯複雜， ShuffleWorker 作爲數據的一個分區，它必須作爲 Sort 算子的子節點，而它作爲計算並行的基本單元，又必須在協程中調用 Window 算子的 Next 方法來完成計算，所以在原來的實現往 ShuffleWorker 裏放了個指向 Window 算子的指針，這樣的設計一方面存在破壞執行樹有向無環特性的隱患，另一方面也降低了代碼的可讀性。
當然，第一個問題我們可以通過在 ShuffleWorker 中增加一個布爾變量來解決：同一個 MergeJoin 對應的兩個 ShuffleWorker 一個爲 true，一個爲 false，只有爲 true 的那個纔會啓動協程。可是這個方法無疑會使上面提到的複雜的邏輯更加複雜。
我們提出的解決方案是把數據分區和計算並行解耦。如圖 2 所示：計算並行還是由 ShuffleWorker 負責，但是它不再是數據流動過程中的一環，它原來在數據流動過程中的位置由 ShuffleReceiver 來代替。MergeJoin 是 ShuffleWorker 的一個成員，每個 ShuffleWorker 對應一個協程，在協程中調用 MergeJoin 的 Next 方法，並將結果發送給彙總算子，這樣上文中提到的兩個問題都得到了解決。
圖 2 拓展後的 Shuffle 算子
相關 PR：https://github.com/pingcap/tidb/pull/20942
實現與效果
在實現中我們考慮兩個場景：其一是數據源本身無序的情況，這種情況下數據進入 MergeJoin 之前要先經過 Sort 節點；另一是數據源本身有序的情況，這種情況下數據進入 MergeJoin 之前無需排序。
圖 3 數據源無序情況下的 ShuffleMergeJoin
圖 3 展示的就是數據源無序情況下 MergeJoin 的並行化過程，這種情況 MergeJoin 和 Sort 算子的計算開銷都可以分攤到多個協程上。啓動 2 個 worker 的優化效果如表 1 所示，我們在不同規模的數據源上都做了測試，表中前兩列是兩個數據源的行數，表中的後兩列是串行和並行版本的運行性能，單位是 ns/op，越小性能越高。從表中可以看出， Shuffle 是可以明顯加速 MergeJoin 運算的，並且數據量越大的情況下加速效果越好（因爲並行化是會引入管道、協程等額外開銷的，比較大的數據量才能保證並行化的收益大於開銷）。在我們的幾個測試案例中，效果最好的情況下 2 個 worker 的運算時間僅爲串行版本的 56.5% 。
表 1 ShuffleMergeJoin 優化效果
圖 4 展示的是數據源有序情況下 MergeJoin 的並行化過程，區別就是數據不再經過 Sort 算子。這種情況下計算的負載本身比較輕量，相比之下根據哈希值來分發數據的 Splitter 就成爲了系統的性能瓶頸，並行化以後性能提升並不明顯。
圖 4 數據源有序情況下的 ShuffleMergeJoin
相關 PR：
1. ShuffleMergeJoin 實現：https://github.com/pingcap/tidb/pull/21255
2. 控制參數：https://github.com/pingcap/tidb/pull/21332
3. 單元測試與性能測試：https://github.com/pingcap/tidb/pull/21360
ShuffleStreamAggregation
聚合運算是 SQL 語句必不可少的一部分，無論是 OLTP 還是 OLAP 場景，聚合都是經常被使用到的算子。
從系統實現層面來看，聚合的實現一般有兩種，第一種是基於 Hash 的方法，該方法通過構建 Hash table，維護每一個被聚合元素的值，計算得到最後的結果值。另外一種，則是基於有序數據流的方法，該方法要求輸入數據源必須是有序的，然後通過遍歷有序的數據流，並在同時維護相應的聚合值，即可得到最後的計算結果。
一般來說，基於 Hash 的方法具有更快的計算速度，但是它需要維護一個 Hash table，內存空間使用成本較高，當被聚合 key 的可能取值個數非常大的時候，那麼相應 Hash table 中的元素個數也會非常多，對內存是個不小的考驗，存在爆內存的風險，這反而導致計算不能正常地被完成。而基於有序數據流方法的聚合運算實現方式，無需隨時都在內存中維護所有的被聚合 key 的值，因此內存使用量相對較小，但是它的運行速度相對而言更慢一點，且更爲嚴格地要求輸入數據必須是有序的。如果可以提升基於有序數據流方法的聚合算子的運行速度，那麼該方法將會更加適用於大數據量的情況。因此我們選擇對基於有序數據流方法的聚合運算實現方法，即 Stream Aggregation 進行並行加速，以提升該算子的整體運行速度。
實現與效果
在具體的實現過程中，我們利用了之前由其他社區貢獻者提供的 Shuffle 算子，在 StreamAggregation 算子外圍，將輸入數據分割成多個有序的輸入數據流，分別輸入到多個 StreamAggregation 算子當中，然後通過簡單的整合，得到最後的計算結果。簡單地說，初始輸入數據源 DataSource,首先會經過 Shuffle 算子，被分割成多個 Partition，且每個 Partition 都是其內部有序的，然後每個 Partition 分別被作爲一 StreamAggregation 算子的輸入，生成部分結果，最後通過對相同 key 的元素進行整合，即可得到最後的整體計算結果。
此處需要考慮 DataSource 是否有序的情況，如果 DataSource 在被聚合 key 上是無序的，比如普通的 PhysicalTableReader 算子，或其他算子的輸出，那麼需要在被分割之前，使得其有序，因此需要在其上添加一個 Sort 算子（如圖 5 所示）。
圖 5 數據源無序情況下的 ShuffleStreamAggregation
針對這種場景，我們的方法最後取得了非常明顯的性能提升（如表 2 所示）。分析認爲，非並行的情況下，Sort 作用在整個 DataSource 之上，而並行化的版本是作用在每個不同的 Partition 之上，輸入相對較小，且並行執行，因此性能提升較大。
表 2 ShuffleStreamAggregation 優化效果
另外，還需要考慮 DataSource 在被聚合 key 上是有序的情況，比如下面的 SQL 語句，被聚合 key 爲 b，且輸入數據源 t 上剛好有由 b 創建的索引，因此在具體的計算過程中， DataSource 是基於 b 的 PhysicalIndexTableReader ，那麼我們就無序引入 Sort 算子，直接將輸入分割成多個 Partition，然後經過圖 6 所示的計算過程即可得到結果。
create table t(a int, b int, key b(b));
select /*+ stream_agg() */ count(a) from t group by b;
圖 6 數據源有序情況下的 ShuffleStreamAggregation
通過 Benchmark 的結果表明，在該情況下，目前的基於 Shuffle 的實現，運行速度並沒有得到提升，反而有所下降，我們粗淺地認爲，當前的 Shuffle 實現方式是瓶頸點，是後續需要被解決的重點。
相關 PR：
RangeSplitter
上面說了， Shuffle 算子會把數據輸入分割成多個 Partition，最開始的時候只有基於 Hash 方法的 Splitter 被實現，該實現對輸入數據是否有序並不做要求。對於數據源有序的情況，儘管該方法依舊適用，但是使用基於 Range 的方法對數據源進行分割，是一個更爲自然的方式，因爲被聚合 key 相同的多行數據，必然是緊挨在以前的，如果可以直接找到這一塊數據的起始點和結束點，整體一次性分割，則無需構建 HashTable，也不用調用開銷更爲明顯的 HashFunction，使得Partiton 過程的開銷更小。基於這一思路，我們實現了 PartitionRangeSplitter，該方法的計算原理是，將緊挨一起的相同聚合 key 的多行數據，批量地分發到一個 worker 之上。相較於基於 Hash 方法的 Partitioner 而言，基於 Range 方法的實現方式的開銷更小，在同時處理有序輸入數據源的情況下，使用 RangePartitioner 能比使用 HashPartitioner 快上一倍的速度（詳見表 3），由此證明了該算子更加適用於數據源有序的情況。
表 3 RangeSplitter 與 HashSplitter 性能比較
相關 PR：
1. RangeSplitter 實現：https://github.com/pingcap/tidb/pull/21306
2. 相關性能測試：https://github.com/pingcap/tidb/pull/21363
總結
在本次性能挑戰大賽中，我們使用 Shuffle 算子對 MergeJoin 算子和 Stream Aggregation 算子進行了加速，在數據源無序的場景下，取得了明顯的性能提升。在優化 MergeJoin 的過程中，爲了適配多個數據源的算子，我們對現有的 Shuffle 實現做了擴展，提高了可讀性和可擴展性。在優化 StreamAggregation 的過程中，考慮到數據源有序的情況，提出了一個簡單的基於 Range 方法的 Splitter 實現，也證明了其有效性。我們在後續將會考慮如何對現有的 Shuffle 算子進行改造，消除其中存在的性能瓶頸，以期進一步提升基於 Shuffle 的一系列並行算子的性能。