三十分鐘成為 Contributor | 為 TiKV 新增 built-in 函式

PingCAP發表於2018-08-02
UI函式

作者：吳雪蓮

背景知識

SQL 語句傳送到 TiDB 後經過 parser 生成 AST（抽象語法樹），再經過 Query Optimizer 生成執行計劃，執行計劃切分成很多子任務，這些子任務以表示式的方式最後下推到底層的各個 TiKV 來執行。

圖 1

圖 1

如圖 1，當 TiDB 收到來自客戶端的查詢請求

select count(*) from t where a + b > 5

時，執行順序如下：

  1. TiDB 對 SQL 進行解析，組織成對應的表示式，下推給 TiKV

  2. TiKV 收到請求後，迴圈以下過程

    • 獲取下一行完整資料，並按列解析

    • 使用引數中的 where 表示式對資料進行過濾

    • 若上一條件符合，進行聚合計算

  3. TiKV 向 TiDB 返回聚合計算結果

  4. TiDB 對所有涉及的結果進行二次聚合，返回給客戶端

這裡的 where 條件便是以表示式樹的形式下推給 TiKV。在此之前 TiDB 只會向 TiKV 下推一小部分簡單的表示式，比如取出某一個列的某個資料型別的值，簡單資料型別的比較操作，算術運算等。為了充分利用分散式叢集的資源，進一步提升 SQL 在整個叢集的執行速度，我們需要將更多種類的表示式下推到 TiKV 來執行，其中的一大類就是 MySQL built-in 函式。

目前，由於 TiKV 的 built-in 函式尚未全部實現，對於無法下推的表示式，TiDB 只能自行解決。這無疑將成為提升 TiDB 速度的最大絆腳石。好訊息是，TiKV 在實現 built-in 函式時，可以直接參考 TiDB 的對應函式邏輯（順便可以幫 TiDB 找找 Bug），為我們減少了不少工作量。

Built-in 函式無疑是 TiDB 和 TiKV 成長道路上不可替代的一步，如此艱鉅又龐大的任務，我們需要廣大社群朋友們的支援與鼓勵。親愛的朋友們，想玩 Rust 嗎？想給 TiKV 提 PR 嗎？想幫助 TiDB 跑得更快嗎？動動您的小手指，拿 PR 來砸我們吧。您的 PR 一旦被採用，將會有小驚喜哦。

手把手教你實現 built-in 函式

Step 1：準備下推函式

在 TiKV 的 github.com/pingcap/tik… issue 中，找到未實現的函式簽名列表，選一個您想要實現的函式。

Step 2：獲取 TiDB 中可參考的邏輯實現

在 TiDB 的 expression 目錄下查詢相關 builtinXXXSig 物件，這裡 XXX 為您要實現的函式簽名，本例中以 MultiplyIntUnsigned 為例，可以在 TiDB 中找到其對應的函式簽名（builtinArithmeticMultiplyIntUnsignedSig）及 實現。

Step 3：確定函式定義

  1. built-in 函式所在的檔名要求與 TiDB 的名稱對應，如 TiDB 中，expression 目錄下的下推檔案統一以 builtin_XXX 命名，對應到 TiKV 這邊，就是 builtin_XXX.rs。若同名對應的檔案不存在，則需要自行在同級目錄下新建。對於本例，當前函式存放於 TiDB 的 builtin_arithmetic.go 檔案裡，對應到 TiKV 便是存放在 builtin_arithmetic.rs 中。

  2. 函式名稱：函式簽名轉為 Rust 的函式名稱規範，這裡 MultiplyIntUnsigned 將會被定義為 multiply_int_unsigned。

  3. 函式返回值，可以參考 TiDB 中實現的 Eval 函式，對應關係如下：

    TiDB 對應實現的 Eval 函式 TiKV 對應函式的返回值型別
    evalInt Result<Option<i64>>
    evalReal Result<Option<f64>>
    evalString Result<Option<Cow<'a, [u8]>>>
    evalDecimal Result<Option<Cow<'a, Decimal>>>
    evalTime Result<Option<Cow<'a, Time>>>
    evalDuration Result<Option<Cow<'a, Duration>>>
    evalJSON Result<Option<Cow<'a, Json>>>

    可以看到 TiDB 的 builtinArithmeticMultiplyIntUnsignedSig  物件實現了 evalInt 方法，故當前函式（multiply_int_unsigned）的返回型別應該為 Result<Option<i64>>。

  4. 函式的引數, 所有 builtin-in 的引數都與 Expression 的 eval 函式一致，即：

    • 環境配置量 (ctx:&StatementContext)

    • 該行資料每列具體值 (row:&[Datum])

綜上，multiply_int_unsigned 的下推函式定義為：

    pub fn multiply_int_unsigned(
       &self,
       ctx: &mut EvalContext,
       row: &[Datum],
   ) -> Result<Option<i64>>
複製程式碼

Step 4：實現函式邏輯

這一塊相對簡單，直接對照 TiDB 的相關邏輯實現即可。這裡，我們可以看到 TiDB 的 builtinArithmeticMultiplyIntUnsignedSig 的具體實現如下：

func (s *builtinArithmeticMultiplyIntUnsignedSig) evalInt(row types.Row) (val int64, isNull bool, err error) {
  a, isNull, err := s.args[0].EvalInt(s.ctx, row)
  if isNull || err != nil {
     return 0, isNull, errors.Trace(err)
  }
  unsignedA := uint64(a)
  b, isNull, err := s.args[1].EvalInt(s.ctx, row)
  if isNull || err != nil {
     return 0, isNull, errors.Trace(err)
  }
  unsignedB := uint64(b)
  result := unsignedA * unsignedB
  if unsignedA != 0 && result/unsignedA != unsignedB {
     return 0, true, types.ErrOverflow.GenByArgs("BIGINT UNSIGNED", fmt.Sprintf("(%s * %s)", s.args[0].String(), s.args[1].String()))
  }
  return int64(result), false, nil
}
複製程式碼

參考以上程式碼，翻譯到 TiKV 即可，如下：

 pub fn multiply_int_unsigned(
       &self,
       ctx: &mut EvalContext,
       row: &[Datum],
   ) -> Result<Option<i64>> {
       let lhs = try_opt!(self.children[0].eval_int(ctx, row));
       let rhs = try_opt!(self.children[1].eval_int(ctx, row));
       let res = (lhs as u64).checked_mul(rhs as u64).map(|t| t as i64);
       // TODO: output expression in error when column's name pushed down.
       res.ok_or_else(|| Error::overflow("BIGINT UNSIGNED", &format!("({} * {})", lhs, rhs)))
           .map(Some)
   }
複製程式碼

Step 5：新增引數檢查

TiKV 在收到下推請求時，首先會對所有的表示式進行檢查，表示式的引數個數檢查就在這一步進行。

TiDB 中對每個 built-in 函式的引數個數有嚴格的限制，這一部分檢查可參考 TiDB 同目錄下 builtin.go 相關程式碼。

在 TiKV 同級目錄的 scalar_function.rs 檔案裡，找到 ScalarFunc 的 check_args 函式，按照現有的模式，加入引數個數的檢查即可。

Step 6：新增下推支援

TiKV 在對一行資料執行具體的 expression 時，會呼叫 eval 函式，eval 函式又會根據具體的返回型別，執行具體的子函式。這一部分工作在 scalar_function.rs 中以巨集（dispatch_call）的形式完成。

對於 MultiplyIntUnsigned, 我們最終返回的資料型別為 Int，所以可以在 dispatch_call 中找到 INT_CALLS，然後照著加入 MultiplyIntUnsigned => multiply_int_unsigned , 表示當解析到函式簽名 MultiplyIntUnsigned 時，呼叫上述已實現的函式 multiply_int_unsigned。

至此 MultiplyIntUnsigned 下推邏輯已完全實現。

Step 7：新增測試

在函式 multiply_int_unsigned 所在檔案 builtin_arithmetic.rs 底部的 test 模組中加入對該函式簽名的單元測試，要求覆蓋到上述新增的所有程式碼，這一部分也可以參考 TiDB 中相關的測試程式碼。本例在 TiKV 中實現的測試程式碼如下：

    #[test]
   fn test_multiply_int_unsigned() {
       let cases = vec![
           (Datum::I64(1), Datum::I64(2), Datum::U64(2)),
           (
               Datum::I64(i64::MIN),
               Datum::I64(1),
               Datum::U64(i64::MIN as u64),
           ),
           (
               Datum::I64(i64::MAX),
               Datum::I64(1),
               Datum::U64(i64::MAX as u64),
           ),
           (Datum::U64(u64::MAX), Datum::I64(1), Datum::U64(u64::MAX)),
       ];

       let mut ctx = EvalContext::default();
       for (left, right, exp) in cases {
           let lhs = datum_expr(left);
           let rhs = datum_expr(right);

           let mut op = Expression::build(
               &mut ctx,
               scalar_func_expr(ScalarFuncSig::MultiplyIntUnsigned, &[lhs, rhs]),
           ).unwrap();
           op.mut_tp().set_flag(types::UNSIGNED_FLAG as u32);

           let got = op.eval(&mut ctx, &[]).unwrap();
           assert_eq!(got, exp);
       }

       // test overflow
       let cases = vec![
           (Datum::I64(-1), Datum::I64(2)),
           (Datum::I64(i64::MAX), Datum::I64(i64::MAX)),
           (Datum::I64(i64::MIN), Datum::I64(i64::MIN)),
       ];

       for (left, right) in cases {
           let lhs = datum_expr(left);
           let rhs = datum_expr(right);

           let mut op = Expression::build(
               &mut ctx,
               scalar_func_expr(ScalarFuncSig::MultiplyIntUnsigned, &[lhs, rhs]),
           ).unwrap();
           op.mut_tp().set_flag(types::UNSIGNED_FLAG as u32);

           let got = op.eval(&mut ctx, &[]).unwrap_err();
           assert!(check_overflow(got).is_ok());
       }
   }
複製程式碼

Step 8：執行測試

執行 make expression，確保所有的 test case 都能跑過。

完成以上幾個步驟之後，就可以給 TiKV 專案提 PR 啦。想要了解提 PR 的基礎知識，嘗試移步 此文，看看是否有幫助。

相關文章