TiDB 源碼閱讀（十五） Sort Merge Join

08-23

來自專欄 TiDB 的後花園30 人贊了文章

作者： @姚維

什麼是 Sort Merge Join

在開始閱讀源碼之前, 我們來看看什麼是 Sort Merge Join (SMJ)，定義可以看 wikipedia。簡單說來就是將 Join 的兩個表，首先根據連接屬性進行排序，然後進行一次掃描歸併, 進而就可以得出最後的結果。這個演算法最大的消耗在於對內外表數據進行排序，而當連接列為索引列時，我們可以利用索引的有序性避免排序帶來的消耗, 所以通常在查詢優化器中，連接列為索引列的情況下可以考慮選擇使用 SMJ。

TiDB Sort Merge Join 實現

執行過程

TiDB 的實現代碼在 tidb/executor/merge_join.go 中 MergeJoinExec.NextChunk 是這個運算元的入口。下面以 SELECT * FROM A JOIN B ON A.a = B.a 為例，對 SMJ 執行過程進行簡述，假設此時外表為 A，內表為 B，join-keys 為 a，A，B 表的 a 列上都有索引：

1.順序讀取外表 A 直到 join-keys 中出現另外的值，把相同 keys 的行放入數組 a1，同樣的規則讀取內表 B，把相同 keys 的行放入數組 a2。如果外表數據或者內表數據讀取結束，退出。

2. 從 a1 中讀取當前第一行數據，設為 v1。從 a2 中讀取當前第一行數據，設為 v2。

3. 根據 join-keys 比較 v1，v2，結果分為幾種情況：

cmpResult > 0, 表示 v1 大於 v2，把當前 a2 的數據丟棄，從內表讀取下一批數據，讀取方法同 1。重複 2。
cmpResult < 0, 表示 v1 小於 v2，說明外表的 v1 沒有內表的值與之相同，把外表數據輸出給 resultGenerator（不同的連接類型會有不同的結果輸出，例如外連接會把不匹配的外表數據輸出）。
cmpResult == 0, 表示 v1 等於 v2。那麼遍歷 a1 裡面的數據，跟 a2 的數據，輸出給 resultGenerator 作一次連接。

4. 回到步驟 1。

下面的圖展示了 SMJ 的過程：

讀取內表 / 外表數據

我們分別通過 fetchNextInnerRows 或者 fetchNextOuterRows 讀取內表和外表的數據。這兩個函數實現的功能類似，這裡只詳述函數 fetchNextInnerRows 的實現。

MergeSortExec 運算元讀取數據，是通過迭代器 readerIterator 完成，readerIterator 可以順序讀取數據。MergeSortExec 運算元維護兩個 readerIterator：outerIter 和 innerIter，它們在 buildMergeJoin 函數中被構造。

真正讀取數據的操作是在 readerIterator.nextSelectedRow 中完成, 這裡會通過 ri.reader.NextChunk 每次讀取一個 Chunk 的數據，關於 Chunk 的相關內容，可以查看我們之前的文章 TiDB 源碼閱讀系列文章（十）Chunk 和執行框架簡介。

這裡值得注意的是，我們通過 expression.VectorizedFilter 對外表數據進行過濾，返回一個 curSelected 布爾數組，用於外表的每一行數據是否是滿足 filter 過濾條件。以 select * from t1 left outer join t2 on t1.a=100; 為例, 這裡的 filter 是 t1.a=100, 對於沒有通過這個過濾條件的行，我們通過 ri.joinResultGenerator.emitToChunk 函數發送給 resultGenerator, 這個 resultGenerator 是一個 interface，具體是否輸出這行數據，會由 join 的類型決定，比如外連接則會輸出，內連接則會忽略。具體關於 resultGenerator, 可以參考之前的文章：TiDB 源碼閱讀系列文章（九）Hash Join

rowsWithSameKey 通過 nextSelectedRow 不斷讀取下一行數據，並通過對每行數據的 join-keys 進行判斷是不是屬於同一個 join-keys，如果是，會把相同 join-keys 的行分別放入到 innerChunkRows 和 outerIter4Row 數組中。然後對其分別建立迭代器 innerIter4Row 和 outerIter4Row。在 SMJ 中的執行過程中，會利用這兩個迭代器來獲取數據進行真正的比較得出 join result。

Merge-Join

實現 Merge-Join 邏輯的代碼在函數 MergeJoinExec.joinToChunk, 對內外表迭代器的當前數據根據各自的 join-keys 作對比，有如下幾個結果：

cmpResult > 0，代表外表當前數據大於內表數據，那麼通過 fetchNextInnerRows 直接讀取下一個內表數據，然後重新比較即可。
cmpResult < 0，代表外表當前數據小於內表數據，這個時候就分幾種情況了，如果是外連接，那麼需要輸出外表數據 + NULL，如果是內連接，那麼這個外表數據就被忽略，對於這個不同邏輯的處理，統一由 e.resultGenerator 來控制，我們只需要把外表數據通過 e.resultGenerator.emitToChunk 調用它即可。然後通過 fetchNextOuterRows 讀取下一個外表數據，重新比較。
cmpResult == 0，代表外表當前數據等於內表當前數據，這個時候就把外表數據跟內表當前數據做一次連接，通過 e.resultGenerator.emitToChunk 生成結果。之後外表跟內表分別獲取下一個數據，重新開始比較。

重複上面的過程，直到外表或者內表數據被遍歷完，退出 Merge-Join 的過程。

我們上面的分析代碼基於 Source-code 分支，可能大家已經發現了一些問題，比如我們會一次性讀取內外表的 Join group（相同的 key）。這裡如果相同的 key 比較多，是有內存 OOM 的風險的。針對這個問題，我們在最新的 master 分支做了幾個事情來優化：

外表其實不需要把相同的 keys 一次性都讀取上來，它只需要按次迭代外表數據，再跟內表逐一對比作連接即可。這裡至少可以減少外表發生 OOM 的問題，可以大大減少 OOM 的概率。
對於內表，我們對 OOM 也不是沒有辦法，我們用 memory.Tracker 這個內存追蹤器來記錄當前內表已經使用的中間結果的內存大小，如果它超過我們設置的閾值，我們會採取輸出日誌或者終止 SQL 繼續運行的方法來規避 OOM 的發生。關於 memory.Tracker 我們不在此展開，可以留意我們後續的源碼分析文章。

後續我們還會在 Merge-Join 方面做一些優化，比如我們可以做多路歸併，中間結果存外存等等，敬請期待。