為 TiDB 重構 built-in 函數
為了加速表達式計算速度，最近我們對表達式的計算框架進行了重構，這篇教程為大家分享如何利用新的計算框架為 TiDB 重寫或新增 built-in 函數。對于部分背景知識請參考這篇文章，本文將首先介紹利用新的表達式計算框架重構 built-in 函數實現的流程，然后以一個函數作為示例進行詳細說明，***介紹重構前后表達式計算框架的區別。
一、重構 built-in 函數整體流程
1. 在 TiDB 源碼 expression 目錄下選擇任一感興趣的函數，假設函數名為 XX
2. 重寫 XXFunctionClass.getFunction() 方法
- 該方法參照 MySQL 規則，根據 built-in 函數的參數類型推導函數的返回值類型
- 根據參數的個數、類型、以及函數的返回值類型生成不同的函數簽名，關于函數簽名的詳細介紹見文末附錄
3. 實現該 built-in 函數對應的所有函數簽名的 evalYY() 方法，此處 YY 表示該函數簽名的返回值類型
4. 添加測試：
- 在 expression 目錄下，完善已有的 TestXX() 方法中關于該函數實現的測試
- 在 executor 目錄下，添加 SQL 層面的測試
5. 運行 make dev，確保所有的 test cast 都能跑過
二、示例
這里以重寫 LENGTH() 函數的 PR 為例，進行詳細說明
1. 首先看 expression/builtin_string.go：
(1)實現 lengthFunctionClass.getFunction() 方法
該方法主要完成兩方面工作： 1. 參照 MySQL 規則推導 LEGNTH 的返回值類型 2. 根據 LENGTH 函數的參數個數、類型及返回值類型生成函數簽名。由于 LENGTH 的參數個數、類型及返回值類型只存在確定的一種情況，因此此處沒有定義新的函數簽名類型，而是修改已有的 builtinLengthSig，使其組合了 baseIntBuiltinFunc(表示該函數簽名返回值類型為 int)
- type builtinLengthSig struct {
- baseIntBuiltinFunc
- }
- func (c *lengthFunctionClass) getFunction(args []Expression, ctx context.Context) (builtinFunc, error) {
- // 參照 MySQL 規則，對 LENGTH 函數返回值類型進行推導
- tp := types.NewFieldType(mysql.TypeLonglong)
- tp.Flen = 10
- types.SetBinChsClnFlag(tp)
- // 根據參數個數、類型及返回值類型生成對應的函數簽名，注意此處與重構前不同，使用的是 newBaseBuiltinFuncWithTp 方法，而非 newBaseBuiltinFunc 方法
- // newBaseBuiltinFuncWithTp 的函數聲明中，args 表示函數的參數，tp 表示函數的返回值類型，argsTp 表示該函數簽名中所有參數對應的正確類型
- // 因為 LENGTH 的參數個數為1，參數類型為 string，返回值類型為 int，因此此處傳入 tp 表示函數的返回值類型，傳入 tpString 用來標識參數的正確類型。對于多個參數的函數，調用 newBaseBuiltinFuncWithTp 時，需要傳入所有參數的正確類型
- bf, err := newBaseBuiltinFuncWithTp(args, tp, ctx, tpString)
- if err != nil {
- return nil, errors.Trace(err)
- }
- sig := &builtinLengthSig{baseIntBuiltinFunc{bf}}
- return sig.setSelf(sig), errors.Trace(c.verifyArgs(args))
- }
(2) 實現 builtinLengthSig.evalInt() 方法
- func (b *builtinLengthSig) evalInt(row []types.Datum) (int64, bool, error) {
- // 對于函數簽名 builtinLengthSig，其參數類型已確定為 string 類型，因此直接調用 b.args[0].EvalString() 方法計算參數
- val, isNull, err := b.args[0].EvalString(row, b.ctx.GetSessionVars().StmtCtx)
- if isNull || err != nil {
- return 0, isNull, errors.Trace(err)
- }
- return int64(len([]byte(val))), false, nil
- }
然后看 expression/builtin_string_test.go，對已有的 TestLength() 方法進行完善：
- func (s *testEvaluatorSuite) TestLength(c *C) {
- defer testleak.AfterTest(c)() // 監測 goroutine 泄漏的工具，可以直接照搬
- // cases 的測試用例對 length 方法實現進行測試
- // 此處注意，除了正常 case 之外，***能添加一些異常的 case，如輸入值為 nil，或者是多種類型的參數
- cases := []struct {
- args interface{}
- expected int64
- isNil bool
- getErr bool
- }{
- {"abc", 3, false, false},
- {"你好", 6, false, false},
- {1, 1, false, false},
- ...
- }
- for _, t := range cases {
- f, err := newFunctionForTest(s.ctx, ast.Length, primitiveValsToConstants([]interface{}{t.args})...)
- c.Assert(err, IsNil)
- // 以下對 LENGTH 函數的返回值類型進行測試
- tp := f.GetType()
- c.Assert(tp.Tp, Equals, mysql.TypeLonglong)
- c.Assert(tp.Charset, Equals, charset.CharsetBin)
- c.Assert(tp.Collate, Equals, charset.CollationBin)
- c.Assert(tp.Flag, Equals, uint(mysql.BinaryFlag))
- c.Assert(tp.Flen, Equals, 10)
- // 以下對 LENGTH 函數的計算結果進行測試
- d, err := f.Eval(nil)
- if t.getErr {
- c.Assert(err, NotNil)
- } else {
- c.Assert(err, IsNil)
- if t.isNil {
- c.Assert(d.Kind(), Equals, types.KindNull)
- } else {
- c.Assert(d.GetInt64(), Equals, t.expected)
- }
- }
- }
- // 以下測試函數是否是具有確定性
- f, err := funcs[ast.Length].getFunction([]Expression{Zero}, s.ctx)
- c.Assert(err, IsNil)
- c.Assert(f.isDeterministic(), IsTrue)
- }
***看 executor/executor_test.go，對 LENGTH 的實現進行 SQL 層面的測試：
- // 關于 string built-in 函數的測試可以在這個方法中添加
- func (s *testSuite) TestStringBuiltin(c *C) {
- defer func() {
- s.cleanEnv(c)
- testleak.AfterTest(c)()
- }()
- tk := testkit.NewTestKit(c, s.store)
- tk.MustExec("use test")
- // for length
- // 此處的測試***也能覆蓋多種不同的情況
- tk.MustExec("drop table if exists t")
- tk.MustExec("create table t(a int, b double, c datetime, d time, e char(20), f bit(10))")
- tk.MustExec(`insert into t values(1, 1.1, "2017-01-01 12:01:01", "12:01:01", "abcdef", 0b10101)`)
- result := tk.MustQuery("select length(a), length(b), length(c), length(d), length(e), length(f), length(null) from t")
- result.Check(testkit.Rows("1 3 19 8 6 2 <nil>"))
- }
三、重構前的表達式計算框架
TiDB 通過 Expression 接口(在 expression/expression.go 文件中定義)對表達式進行抽象，并定義 eval 方法對表達式進行計算：
- type Expression interface{
- ...
- eval(row []types.Datum) (types.Datum, error)
- ...
- }
實現 Expression 接口的表達式包括：
- Scalar Function：標量函數表達式
- Column：列表達式
- Constant：常量表達式
下面以一個例子說明重構前的表達式計算框架。
例如：
- create table t (
- c1 int,
- c2 varchar(20),
- c3 double
- )
- select * from t where c1 + CONCAT( c2, c3 < “1.1” )
對于上述 select 語句 where 條件中的表達式： 在編譯階段，TiDB 將構建出如下圖所示的表達式樹:
在執行階段，調用根節點的 eval 方法，通過后續遍歷表達式樹對表達式進行計算。
對于表達式 ‘<’，計算時需要考慮兩個參數的類型，并根據一定的規則，將兩個參數的值轉化為所需的數據類型后進行計算。上圖表達式樹中的 ‘<’，其參數類型分別為 double 和 varchar，根據 MySQL 的計算規則，此時需要使用浮點類型的計算規則對兩個參數進行比較，因此需要將參數 “1.1” 轉化為 double 類型，而后再進行計算。
同樣的，對于上圖表達式樹中的表達式 CONCAT，計算前需要將其參數分別轉化為 string 類型;對于表達式 ‘+’，計算前需要將其參數分別轉化為 double 類型。
因此，在重構前的表達式計算框架中，對于參與運算的每一組數據，計算時都需要大量的判斷分支重復地對參數的數據類型進行判斷，若參數類型不符合表達式的運算規則，則需要將其轉換為對應的數據類型。
此外，由 Expression.eval() 方法定義可知，在運算過程中，需要通過 Datum 結構不斷地對中間結果進行包裝和解包，由此也會帶來一定的時間和空間開銷。
為了解決這兩點問題，我們對表達式計算框架進行重構。
四、重構后的表達式計算框架
重構后的表達式計算框架，一方面，在編譯階段利用已有的表達式類型信息，生成參數類型“符合運算規則”的表達式，從而保證在運算階段中無需再對類型增加分支判斷;另一方面，運算過程中只涉及原始類型數據，從而避免 Datum 帶來的時間和空間開銷。
繼續以上文提到的查詢為例，在編譯階段，生成的表達式樹如下圖所示，對于不符合函數參數類型的表達式，為其加上一層 cast 函數進行類型轉換;
這樣，在執行階段，對于每一個 ScalarFunction，可以保證其所有的參數類型一定是符合該表達式運算規則的數據類型，無需在執行過程中再對參數類型進行檢查和轉換。
五、附錄
1. 對于一個 built-in 函數，由于其參數個數、類型以及返回值類型的不同，可能會生成多個函數簽名分別用來處理不同的情況。對于大多數 built-in 函數，其每個參數類型及返回值類型均確定，此時只需要生成一個函數簽名。
2. 對于較為復雜的返回值類型推導規則，可以參考 CONCAT 函數的實現和測試。可以利用 MySQLWorkbench 工具運行查詢語句 select funcName(arg0, arg1, ...) 觀察 MySQL 的 built-in 函數在傳入不同參數時的返回值數據類型。
3. 在 TiDB 表達式的運算過程中，只涉及 6 種運算類型(目前正在實現對 JSON 類型的支持)，分別是
- int (int64)
- real (float64)
- decimal
- string
- Time
- Duration
4. 通過 WrapWithCastAsXX() 方法可以將一個表達式轉換為對應的類型。
5. 對于一個函數簽名，其返回值類型已經確定，所以定義時需要組合與該類型對應的 baseXXBuiltinFunc，并實現 evalXX() 方法。(XX 不超過上述 6 種類型的范圍)
【本文是51CTO專欄機構“PingCAP”的原創文章，轉載請聯系作者本人獲取授權】
