為了加速表達(dá)式計算速度，最近我們對表達(dá)式的計算框架進(jìn)行了重構(gòu)，這篇教程為大家分享如何利用新的計算框架為 TiDB 重寫或新增 built-in 函數(shù)。對于部分背景知識請參考這篇文章，本文將首先介紹利用新的表達(dá)式計算框架重構(gòu) built-in 函數(shù)實現(xiàn)的流程，然后以一個函數(shù)作為示例進(jìn)行詳細(xì)說明，***介紹重構(gòu)前后表達(dá)式計算框架的區(qū)別。

一、重構(gòu) built-in 函數(shù)整體流程

1. 在 TiDB 源碼 expression 目錄下選擇任一感興趣的函數(shù)，假設(shè)函數(shù)名為 XX

2. 重寫 XXFunctionClass.getFunction() 方法

• 該方法參照 MySQL 規(guī)則，根據(jù) built-in 函數(shù)的參數(shù)類型推導(dǎo)函數(shù)的返回值類型
• 根據(jù)參數(shù)的個數(shù)、類型、以及函數(shù)的返回值類型生成不同的函數(shù)簽名，關(guān)于函數(shù)簽名的詳細(xì)介紹見文末附錄

3. 實現(xiàn)該 built-in 函數(shù)對應(yīng)的所有函數(shù)簽名的 evalYY() 方法，此處 YY 表示該函數(shù)簽名的返回值類型

4. 添加測試：

• 在 expression 目錄下，完善已有的 TestXX() 方法中關(guān)于該函數(shù)實現(xiàn)的測試
• 在 executor 目錄下，添加 SQL 層面的測試

5. 運行 make dev，確保所有的 test cast 都能跑過

二、示例

這里以重寫 LENGTH() 函數(shù)的 PR 為例，進(jìn)行詳細(xì)說明

1. 首先看 expression/builtin_string.go：

(1)實現(xiàn) lengthFunctionClass.getFunction() 方法

該方法主要完成兩方面工作： 1. 參照 MySQL 規(guī)則推導(dǎo) LEGNTH 的返回值類型 2. 根據(jù) LENGTH 函數(shù)的參數(shù)個數(shù)、類型及返回值類型生成函數(shù)簽名。由于 LENGTH 的參數(shù)個數(shù)、類型及返回值類型只存在確定的一種情況，因此此處沒有定義新的函數(shù)簽名類型，而是修改已有的 builtinLengthSig，使其組合了 baseIntBuiltinFunc(表示該函數(shù)簽名返回值類型為 int)

type builtinLengthSig struct { 
    baseIntBuiltinFunc 
} 
  
func (c *lengthFunctionClass) getFunction(args []Expression, ctx context.Context) (builtinFunc, error) { 
    // 參照 MySQL 規(guī)則，對 LENGTH 函數(shù)返回值類型進(jìn)行推導(dǎo) 
    tp := types.NewFieldType(mysql.TypeLonglong) 
    tp.Flen = 10 
    types.SetBinChsClnFlag(tp) 
            
    // 根據(jù)參數(shù)個數(shù)、類型及返回值類型生成對應(yīng)的函數(shù)簽名，注意此處與重構(gòu)前不同，使用的是 newBaseBuiltinFuncWithTp 方法，而非 newBaseBuiltinFunc 方法 
    // newBaseBuiltinFuncWithTp 的函數(shù)聲明中，args 表示函數(shù)的參數(shù)，tp 表示函數(shù)的返回值類型，argsTp 表示該函數(shù)簽名中所有參數(shù)對應(yīng)的正確類型 
    // 因為 LENGTH 的參數(shù)個數(shù)為1，參數(shù)類型為 string，返回值類型為 int，因此此處傳入 tp 表示函數(shù)的返回值類型，傳入 tpString 用來標(biāo)識參數(shù)的正確類型。對于多個參數(shù)的函數(shù)，調(diào)用 newBaseBuiltinFuncWithTp 時，需要傳入所有參數(shù)的正確類型 
    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) 實現(xiàn) builtinLengthSig.evalInt() 方法

func (b *builtinLengthSig) evalInt(row []types.Datum) (int64, bool, error) { 
    // 對于函數(shù)簽名 builtinLengthSig，其參數(shù)類型已確定為 string 類型，因此直接調(diào)用 b.args[0].EvalString() 方法計算參數(shù) 
    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() 方法進(jìn)行完善：

func (s *testEvaluatorSuite) TestLength(c *C) { 
    defer testleak.AfterTest(c)() // 監(jiān)測 goroutine 泄漏的工具，可以直接照搬 
    // cases 的測試用例對 length 方法實現(xiàn)進(jìn)行測試 
    // 此處注意，除了正常 case 之外，***能添加一些異常的 case，如輸入值為 nil，或者是多種類型的參數(shù) 
    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 函數(shù)的返回值類型進(jìn)行測試 
        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 函數(shù)的計算結(jié)果進(jìn)行測試 
        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) 
            } 
        } 
    } 
    // 以下測試函數(shù)是否是具有確定性 
    f, err := funcs[ast.Length].getFunction([]Expression{Zero}, s.ctx) 
    c.Assert(err, IsNil) 
    c.Assert(f.isDeterministic(), IsTrue) 
} 

***看 executor/executor_test.go，對 LENGTH 的實現(xiàn)進(jìn)行 SQL 層面的測試：

// 關(guān)于 string built-in 函數(shù)的測試可以在這個方法中添加 
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>")) 
} 

三、重構(gòu)前的表達(dá)式計算框架

TiDB 通過 Expression 接口(在 expression/expression.go 文件中定義)對表達(dá)式進(jìn)行抽象，并定義 eval 方法對表達(dá)式進(jìn)行計算：

type Expression interface{ 
    ... 
    eval(row []types.Datum) (types.Datum, error) 
    ... 
} 

實現(xiàn) Expression 接口的表達(dá)式包括：

• Scalar Function：標(biāo)量函數(shù)表達(dá)式
• Column：列表達(dá)式
• Constant：常量表達(dá)式

下面以一個例子說明重構(gòu)前的表達(dá)式計算框架。

例如：

create table t ( 
    c1 int, 
    c2 varchar(20), 
    c3 double 
) 
  
select * from t where c1 + CONCAT( c2, c3 < “1.1” ) 

對于上述 select 語句 where 條件中的表達(dá)式： 在編譯階段，TiDB 將構(gòu)建出如下圖所示的表達(dá)式樹:

在執(zhí)行階段，調(diào)用根節(jié)點的 eval 方法，通過后續(xù)遍歷表達(dá)式樹對表達(dá)式進(jìn)行計算。

對于表達(dá)式 ‘<’，計算時需要考慮兩個參數(shù)的類型，并根據(jù)一定的規(guī)則，將兩個參數(shù)的值轉(zhuǎn)化為所需的數(shù)據(jù)類型后進(jìn)行計算。上圖表達(dá)式樹中的 ‘<’，其參數(shù)類型分別為 double 和 varchar，根據(jù) MySQL 的計算規(guī)則，此時需要使用浮點類型的計算規(guī)則對兩個參數(shù)進(jìn)行比較，因此需要將參數(shù) “1.1” 轉(zhuǎn)化為 double 類型，而后再進(jìn)行計算。

同樣的，對于上圖表達(dá)式樹中的表達(dá)式 CONCAT，計算前需要將其參數(shù)分別轉(zhuǎn)化為 string 類型;對于表達(dá)式 ‘+’，計算前需要將其參數(shù)分別轉(zhuǎn)化為 double 類型。

因此，在重構(gòu)前的表達(dá)式計算框架中，對于參與運算的每一組數(shù)據(jù)，計算時都需要大量的判斷分支重復(fù)地對參數(shù)的數(shù)據(jù)類型進(jìn)行判斷，若參數(shù)類型不符合表達(dá)式的運算規(guī)則，則需要將其轉(zhuǎn)換為對應(yīng)的數(shù)據(jù)類型。

此外，由 Expression.eval() 方法定義可知，在運算過程中，需要通過 Datum 結(jié)構(gòu)不斷地對中間結(jié)果進(jìn)行包裝和解包，由此也會帶來一定的時間和空間開銷。

為了解決這兩點問題，我們對表達(dá)式計算框架進(jìn)行重構(gòu)。

四、重構(gòu)后的表達(dá)式計算框架

重構(gòu)后的表達(dá)式計算框架，一方面，在編譯階段利用已有的表達(dá)式類型信息，生成參數(shù)類型“符合運算規(guī)則”的表達(dá)式，從而保證在運算階段中無需再對類型增加分支判斷;另一方面，運算過程中只涉及原始類型數(shù)據(jù)，從而避免 Datum 帶來的時間和空間開銷。

繼續(xù)以上文提到的查詢?yōu)槔?，在編譯階段，生成的表達(dá)式樹如下圖所示，對于不符合函數(shù)參數(shù)類型的表達(dá)式，為其加上一層 cast 函數(shù)進(jìn)行類型轉(zhuǎn)換;

這樣，在執(zhí)行階段，對于每一個 ScalarFunction，可以保證其所有的參數(shù)類型一定是符合該表達(dá)式運算規(guī)則的數(shù)據(jù)類型，無需在執(zhí)行過程中再對參數(shù)類型進(jìn)行檢查和轉(zhuǎn)換。

五、附錄

1. 對于一個 built-in 函數(shù)，由于其參數(shù)個數(shù)、類型以及返回值類型的不同，可能會生成多個函數(shù)簽名分別用來處理不同的情況。對于大多數(shù) built-in 函數(shù)，其每個參數(shù)類型及返回值類型均確定，此時只需要生成一個函數(shù)簽名。

2. 對于較為復(fù)雜的返回值類型推導(dǎo)規(guī)則，可以參考 CONCAT 函數(shù)的實現(xiàn)和測試?？梢岳?MySQLWorkbench 工具運行查詢語句 select funcName(arg0, arg1, ...) 觀察 MySQL 的 built-in 函數(shù)在傳入不同參數(shù)時的返回值數(shù)據(jù)類型。

3. 在 TiDB 表達(dá)式的運算過程中，只涉及 6 種運算類型(目前正在實現(xiàn)對 JSON 類型的支持)，分別是

• int (int64)
• real (float64)
• decimal
• string
• Time
• Duration

4. 通過 WrapWithCastAsXX() 方法可以將一個表達(dá)式轉(zhuǎn)換為對應(yīng)的類型。

5. 對于一個函數(shù)簽名，其返回值類型已經(jīng)確定，所以定義時需要組合與該類型對應(yīng)的 baseXXBuiltinFunc，并實現(xiàn) evalXX() 方法。(XX 不超過上述 6 種類型的范圍)

【本文是51CTO專欄機(jī)構(gòu)“PingCAP”的原創(chuàng)文章，轉(zhuǎn)載請聯(lián)系作者本人獲取授權(quán)】

戳這里，看該作者更多好文