十分鐘成爲 Contributor 系列 | 爲 TiDB 重構 built-in 函數

這是十分鐘成爲 TiDB Contributor 系列的第二篇文章，讓你們能夠無門檻參與大型開源項目，感謝社區爲 TiDB 帶來的貢獻，也但願參與 TiDB Community 能爲你的生活帶來更多有意義的時刻。

爲了加速表達式計算速度，最近咱們對錶達式的計算框架進行了重構，這篇教程爲你們分享如何利用新的計算框架爲 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 爲例，進行詳細說明

首先看 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，能夠保證其全部的參數類型必定是符合該表達式運算規則的數據類型，無需在執行過程當中再對參數類型進行檢查和轉換。

附錄

• 對於一個 built-in 函數，因爲其參數個數、類型以及返回值類型的不一樣，可能會生成多個函數簽名分別用來處理不一樣的狀況。對於大多數 built-in 函數，其每一個參數類型及返回值類型均肯定，此時只須要生成一個函數簽名。
• 對於較爲複雜的返回值類型推導規則，能夠參考 CONCAT 函數的實現和測試。能夠利用 MySQLWorkbench 工具運行查詢語句 select funcName(arg0, arg1, ...) 觀察 MySQL 的 built-in 函數在傳入不一樣參數時的返回值數據類型。
• 在 TiDB 表達式的運算過程當中，只涉及 6 種運算類型(目前正在實現對 JSON 類型的支持)，分別是

1. int (int64)
2. real (float64)
3. decimal
4. string
5. Time

6. Duration

   經過 WrapWithCastAsXX() 方法能夠將一個表達式轉換爲對應的類型。

• 對於一個函數簽名，其返回值類型已經肯定，因此定義時須要組合與該類型對應的 baseXXBuiltinFunc，並實現 evalXX() 方法。(XX 不超過上述 6 種類型的範圍)

---------------------------- 我是 AI 的分割線 ----------------------------------------

回顧三月啓動的《十分鐘成爲 TiDB Contributor 系列 | 添加內建函數》活動，在短短的時間內，咱們收到了來自社區貢獻的超過 200 條新建內建函數，這之中有不少是來自大型互聯網公司的資深數據庫工程師，也不乏在學校或是剛畢業在刻苦鑽研分佈式系統和分佈式數據庫的學生。

TiDB Contributor Club 將你們彙集起來，咱們互相分享、討論，一塊兒成長。

感謝你的參與和貢獻，在開源的道路上咱們將義無反顧地走下去，和你一塊兒。

成爲 New Contributor 贈送限量版馬克杯的活動還在繼續中，任何一個新加入集體的小夥伴都將收到咱們充滿了誠意的禮物，很榮幸可以認識你，也很高興能和你一塊兒堅決地走得更遠。

成爲 New Contributor 獲贈限量版馬克杯，馬克杯獲取流程以下：

1. 提交 PR
2. PR提交以後，請耐心等待維護者進行 Review。目前通常在一到兩個工做日內都會進行 Review，若是當前的 PR 堆積數量較多可能回覆會比較慢。代碼提交後 CI 會執行咱們內部的測試，你須要保證全部的單元測試是能夠經過的。期間可能有其它的提交會與當前 PR 衝突，這時須要修復衝突。維護者在 Review 過程當中可能會提出一些修改意見。修改完成以後若是 reviewer 認爲沒問題了，你會收到 LGTM(looks good to me) 的回覆。當收到兩個及以上的 LGTM 後，該 PR 將會被合併。
3. 合併 PR 後自動成爲 Contributor，會收到來自 PingCAP Team 的感謝郵件，請查收郵件並填寫領取表單
   • 表單填寫地址：cn.mikecrm.com/01wE8tX
4. 後臺 AI 覈查 GitHub ID 及資料信息，確認無誤後隨即使快遞寄出屬於你的限量版馬克杯
5. 期待你分享本身參與開源項目的感想和經驗，TiDB Contributor Club 將和你一塊兒分享開源的力量

瞭解更多關於 TiDB 的資料請登錄咱們的官方網站：pingcap.com

加入 TiDB Contributor Club 請添加咱們的 AI 微信：tidbai