TiDB 原始碼閱讀系列文章（二十一）基於規則的優化 II

作者：姚珂男

在 TiDB 原始碼閱讀系列文章（七）基於規則的優化 一文中，我們介紹了幾種 TiDB 中的邏輯優化規則，包括列剪裁，最大最小消除，投影消除，謂詞下推和構建節點屬性，本篇將繼續介紹更多的優化規則：聚合消除、外連線消除和子查詢優化。

聚合消除

聚合消除會檢查 SQL 查詢中 Group By 語句所使用的列是否具有唯一性屬性，如果滿足，則會將執行計劃中相應的 LogicalAggregation 運算元替換為 LogicalProjection 運算元。這裡的邏輯是當聚合函式按照具有唯一性屬性的一列或多列分組時，下層運算元輸出的每一行都是一個單獨的分組，這時就可以將聚合函式展開成具體的引數列或者包含引數列的普通函式表示式，具體的程式碼實現在 rule_aggregation_elimination.go 檔案中。下面舉一些具體的例子。

例一：

下面這個 Query 可以將聚合函式展開成列的查詢：

select max(a) from t group by t.pk;
複製程式碼

被等價地改寫成：

select a from t;
複製程式碼

例二：

下面這個 Query 可以將聚合函式展開為包含引數列的內建函式的查詢：

select count(a) from t group by t.pk;
複製程式碼

被等價地改寫成：

select if(isnull(a), 0, 1) from t;
複製程式碼

這裡其實還可以做進一步的優化：如果列 a 具有 Not Null 的屬性，那麼可以將 if(isnull(a), 0, 1) 直接替換為常量 1（目前 TiDB 還沒做這個優化，感興趣的同學可以來貢獻一個 PR）。

另外提一點，對於大部分聚合函式，引數的型別和返回結果的型別一般是不同的，所以在展開聚合函式的時候一般會在引數列上構造 cast 函式做型別轉換，展開後的表示式會儲存在作為替換 LogicalAggregation 運算元的 LogicalProjection 運算元中。

這個優化過程中，有一點非常關鍵，就是如何知道 Group By 使用的列是否滿足唯一性屬性，尤其是當聚合運算元的下層節點不是 DataSource 的時候？我們在 （七）基於規則的優化 一文中的“構建節點屬性”章節提到過，執行計劃中每個運算元節點會維護這樣一個資訊：當前運算元的輸出會按照哪一列或者哪幾列滿足唯一性屬性。因此，在聚合消除中，我們可以通過檢視下層運算元儲存的這個資訊，再結合 Group By 用到的列判斷當前聚合運算元是否可以被消除。

外連線消除

不同於 （七）基於規則的優化 一文中“謂詞下推”章節提到的將外連線轉換為內連線，這裡外連線消除指的是將整個連線操作從查詢中移除。

外連線消除需要滿足一定條件：

• 條件 1 : LogicalJoin 的父親運算元只會用到 LogicalJoin 的 outer plan 所輸出的列
• 條件 2 :
  • 條件 2.1 : LogicalJoin 中的 join key 在 inner plan 的輸出結果中滿足唯一性屬性
  • 條件 2.2 : LogicalJoin 的父親運算元會對輸入的記錄去重

條件 1 和條件 2 必須同時滿足，但條件 2.1 和條件 2.2 只需滿足一條即可。

滿足條件 1 和 條件 2.1 的一個例子：

select t1.a from t1 left join t2 on t1.b = t2.pk;
複製程式碼

可以被改寫成：

select t1.a from t1;
複製程式碼

滿足條件 1 和條件 2.2 的一個例子：

select distinct(t1.a) from t1 left join t2 on t1.b = t2.b;
複製程式碼

可以被改寫成：

select distinct(t1.a) from t1;
複製程式碼

具體的原理是，對於外連線，outer plan 的每一行記錄肯定會在連線的結果集裡出現一次或多次，當 outer plan 的行不能找到匹配時，或者只能找到一行匹配時，這行 outer plan 的記錄在連線結果中只出現一次；當 outer plan 的行能找到多行匹配時，它會在連線結果中出現多次；那麼如果 inner plan 在 join key 上滿足唯一性屬性，就不可能存在 outer plan 的行能夠找到多行匹配，所以這時 outer plan 的每一行都會且僅會在連線結果中出現一次。同時，上層運算元只需要 outer plan 的資料，那麼外連線可以直接從查詢中被去除掉。同理就可以很容易理解當上層運算元只需要 outer plan 的去重後結果時，外連線也可以被消除。

這部分優化的具體程式碼實現在 rule_join_elimination.go 檔案中。

子查詢優化 / 去相關

子查詢分為非相關子查詢和相關子查詢，例如：

-- 非相關子查詢
select * from t1 where t1.a > (select t2.a from t2 limit 1);
-- 相關子查詢
select * from t1 where t1.a > (select t2.a from t2 where t2.b > t1.b limit 1);
複製程式碼

對於非相關子查詢， TiDB 會在 expressionRewriter 的邏輯中做兩類操作：

• 子查詢展開

  即直接執行子查詢獲得結果，再利用這個結果改寫原本包含子查詢的表示式；比如上述的非相關子查詢，如果其返回的結果為一行記錄 “1” ，那麼整個查詢會被改寫為：

  select * from t1 where t1.a > 1;
  複製程式碼

  詳細的程式碼邏輯可以參考 expression_rewriter.go 中的 handleScalarSubquery 和 handleExistSubquery 函式。

• 子查詢轉為 Join

  對於包含 IN (subquery) 的查詢，比如：

  select * from t1 where t1.a in (select t2.a from t2);
  複製程式碼

  會被改寫成：

  select t1.* from t1 inner join (select distinct(t2.a) as a from t2) as sub on t1.a = sub.a;
  複製程式碼

  如果 t2.a 滿足唯一性屬性，根據上面介紹的聚合消除規則，查詢會被進一步改寫成：

  select t1.* from t1 inner join t2 on t1.a = t2.a;
  複製程式碼

  這裡選擇將子查詢轉化為 inner join 的 inner plan 而不是執行子查詢的原因是：以上述查詢為例，子查詢的結果集可能會很大，展開子查詢需要一次性將 t2 的全部資料從 TiKV 返回到 TiDB 中快取，並作為 t1 掃描的過濾條件；如果將子查詢轉化為 inner join 的 inner plan ，我們可以更靈活地對 t2 選擇訪問方式，比如我們可以對 join 選擇 IndexLookUpJoin 實現方式，那麼對於拿到的每一條 t1 表資料，我們只需拿 t1.a 作為 range 對 t2 做一次索引掃描，如果 t1 表很小，相比於展開子查詢返回 t2 全部資料，我們可能總共只需要從 t2 返回很少的幾條資料。

  注意這個轉換的結果不一定會比展開子查詢更好，其具體情況會受 t1 表和 t2 表資料的影響，如果在上述查詢中， t1 表很大而 t2 表很小，那麼展開子查詢再對 t1 選擇索引掃描可能才是最好的方案，所以現在有引數控制這個轉化是否開啟，詳細的程式碼可以參考 expression_rewriter.go 中的 handleInSubquery 函式。

對於相關子查詢，TiDB 會在 expressionRewriter 中將整個包含相關子查詢的表示式轉化為 LogicalApply 運算元。LogicalApply 運算元是一類特殊的 LogicalJoin ，特殊之處體現在執行邏輯上：對於 outer plan 返回的每一行記錄，取出相關列的具體值傳遞給子查詢，再執行根據子查詢生成的 inner plan ，即 LogicalApply 在執行時只能選擇類似迴圈巢狀連線的方式，而普通的 LogicalJoin 則可以在物理優化階段根據代價模型選擇最合適的執行方式，包括 HashJoin，MergeJoin 和 IndexLookUpJoin，理論上後者生成的物理執行計劃一定會比前者更優，所以在邏輯優化階段我們會檢查是否可以應用“去相關”這一優化規則，試圖將 LogicalApply 轉化為等價的 LogicalJoin 。其核心思想是將 LogicalApply 的 inner plan 中包含相關列的那些運算元提升到 LogicalApply 之中或之上，在運算元提升後如果 inner plan 中不再包含任何的相關列，即不再引用任何 outer plan 中的列，那麼 LogicalApply 就會被轉換為普通的 LogicalJoin ，這部分程式碼邏輯實現在 rule_decorrelate.go 檔案中。

具體的運算元提升方式分為以下幾種情況：

• inner plan 的根節點是 LogicalSelection

  則將其過濾條件新增到 LogicalApply 的 join condition 中，然後將該 LogicalSelection 從 inner plan 中刪除，再遞迴地對 inner plan 提升運算元。

  以如下查詢為例：

  select * from t1 where t1.a in (select t2.a from t2 where t2.b = t1.b);
  複製程式碼

  其生成的最初執行計劃片段會是：

  

  LogicalSelection 提升後會變成如下片段：

  

  到此 inner plan 中不再包含相關列，於是 LogicalApply 會被轉換為如下 LogicalJoin ：

  

• inner plan 的根節點是 LogicalMaxOneRow

  即要求子查詢最多輸出一行記錄，比如這個例子：

  select *, (select t2.a from t2 where t2.pk = t1.a) from t1;
  複製程式碼

  因為子查詢出現在整個查詢的投影項裡，所以 expressionRewriter 在處理子查詢時會對其生成的執行計劃在根節點上加一個 LogicalMaxOneRow 限制最多產生一行記錄，如果在執行時發現下層輸出多於一行記錄，則會報錯。在這個例子中，子查詢的過濾條件是 t2 表的主鍵上的等值條件，所以子查詢肯定最多隻會輸出一行記錄，而這個資訊在“構建節點屬性”這一步時會被髮掘出來並記錄在運算元節點的 MaxOneRow 屬性中，所以這裡的 LogicalMaxOneRow 節點實際上是冗餘的，於是我們可以將其從 inner plan 中移除，然後再遞迴地對 inner plan 做運算元提升。

• inner plan 的根節點是 LogicalProjection

  則首先將這個投影運算元從 inner plan 中移除，再根據 LogicalApply 的連線型別判斷是否需要在 LogicalApply 之上再加上一個 LogicalProjection ，具體來說是：對於非 semi-join 這一類的連線（包括 inner join 和 left join ），inner plan 的輸出列會保留在 LogicalApply 的結果中，所以這個投影操作需要保留，反之則不需要。最後，再遞迴地對刪除投影后的 inner plan 提升下層運算元。

• inner plan 的根節點是 LogicalAggregation

1. 首先我們會檢查這個聚合運算元是否可以被提升到 LogicalApply 之上再執行。以如下查詢為例：

   select *, (select sum(t2.b) from t2 where t2.a = t1.pk) from t1;
   複製程式碼

   其最初生成的執行計劃片段會是：

   

   將聚合提升到 LogicalApply 後的執行計劃片段會是：

   

   即先對 t1 和 t2 做連線，再在連線結果上按照 t1.pk 分組後做聚合。這裡有兩個關鍵變化：第一是不管提升前 LogicalApply 的連線型別是 inner join 還是 left join ，提升後必須被改為 left join ；第二是提升後的聚合新增了 Group By 的列，即要按照 outer plan 傳進 inner plan 中的相關列做分組。這兩個變化背後的原因都會在後面進行闡述。因為提升後 inner plan 不再包含相關列，去相關後最終生成的執行計劃片段會是：

   

   聚合提升有很多限定條件：

   • LogicalApply 的連線型別必須是 inner join 或者 left join 。 LogicalApply 是根據相關子查詢生成的，只可能有 3 類連線型別，除了 inner join 和 left join 外，第三類是 semi join （包括 SemiJoin，LeftOuterSemiJoin，AntiSemiJoin，AntiLeftOuterSemiJoin），具體可以參考 expression_rewriter.go 中的程式碼，限於篇幅在這裡就不對此做展開了。對於 semi join 型別的 LogicalApply ，因為 inner plan 的輸出列不會出現在連線的結果中，所以很容易理解我們無法將聚合運算元提升到 LogicalApply 之上。

   • LogicalApply 本身不能包含 join condition 。以上面給出的查詢為例，可以看到聚合提升後會將子查詢中包含相關列的過濾條件 (t2.a = t1.pk) 新增到 LogicalApply 的 join condition 中，如果 LogicalApply 本身存在 join condition ，那麼聚合提升後聚合運算元的輸入（連線運算元的輸出）就會和在子查詢中時聚合運算元的輸入不同，導致聚合運算元結果不正確。

   • 子查詢中用到的相關列在 outer plan 輸出裡具有唯一性屬性。以上面查詢為例，如果 t1.pk 不滿足唯一性，假設 t1 有兩條記錄滿足 t1.pk = 1，t2 只有一條記錄 { (t2.a: 1, t2.b: 2) }，那麼該查詢會輸出兩行結果 { (sum(t2.b): 2), (sum(t2.b): 2) }；但對於聚合提升後的執行計劃，則會生成錯誤的一行結果{ (sum(t2.b): 4) }。當 t1.pk 滿足唯一性後，每一行 outer plan 的記錄都對應連線結果中的一個分組，所以其聚合結果會和在子查詢中的聚合結果一致，這也解釋了為什麼聚合提升後需要按照 t1.pk 做分組。

   • 聚合函式必須滿足當輸入為 null 時輸出結果也一定是 null 。這是為了在子查詢中沒有匹配的特殊情況下保證結果的正確性，以上面查詢為例，當 t2 表沒有任何記錄滿足 t2.a = t1.pk 時，子查詢中不管是什麼聚合函式都會返回 null 結果，為了保留這種特殊情況，在聚合提升的同時， LogicalApply 的連線型別會被強制改為 left join（改之前可能是 inner join ），所以在這種沒有匹配的情況下，LogicalApply 輸出結果中 inner plan 部分會是 null ，而這個 null 會作為新新增的聚合運算元的輸入，為了和提升前結果一致，其結果也必須是 null 。

2. 對於根據上述條件判定不能提升的聚合運算元，我們再檢查這個聚合運算元的子節點是否為 LogicalSelection ，如果是，則將其從 inner plan 中移除並將過濾條件新增到 LogicalApply 的 join condition 中。這種情況下 LogicalAggregation 依然會被保留在 inner plan 中，但會將 LogicalSelection 過濾條件中涉及的 inner 表的列新增到聚合運算元的 Group By 中。比如對於查詢：

   select *, (select count(*) from t2 where t2.a = t1.a) from t1;
   複製程式碼

   其生成的最初的執行計劃片段會是：

   

   因為聚合函式是 count(*) ，不滿足當輸入為 null 時輸出也為 null 的條件，所以它不能被提升到 LogicalApply 之上，但它可以被改寫成：

   

   注意 LogicalAggregation 的 Group By 新加了 t2.a ，這一步將原本的先做過濾再做聚合轉換為了先按照 t2.a 分組做聚合，再將聚合結果與 t1 做連線。 LogicalSelection 提升後 inner plan 已經不再依賴 outer plan 的結果了，整個查詢去相關後將會變為：

   

總結

這是基於規則優化的第二篇文章，後續我們還將介紹更多邏輯優化規則：聚合下推，TopN 下推和 Join Reorder 。