作者:吳雪蓮
背景知識
SQL 語句發送到 TiDB 後經過 parser 生成 AST(抽象語法樹),再經過 Query Optimizer 生成執行計劃,執行計劃切分成很多子任務,這些子任務以表達式的方式最後下推到底層的各個 TiKV 來執行。
<center>圖 1</center>
如圖 1,當 TiDB 收到來自客戶端的查詢請求
select count(*) from t where a + b > 5
時,執行順序如下:
- TiDB 對 SQL 進行解析,組織成對應的表達式,下推給 TiKV
-
TiKV 收到請求後,循環以下過程
- 獲取下一行完整數據,並按列解析
- 使用參數中的 where 表達式對數據進行過濾
- 若上一條件符合,進行聚合計算
- TiKV 向 TiDB 返回聚合計算結果
- 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 的 https://github.com/pingcap/tikv/issues/3275 issue 中,找到未實現的函數簽名列表,選一個您想要實現的函數。
Step 2:獲取 TiDB 中可參考的邏輯實現
在 TiDB 的 expression 目錄下查找相關 builtinXXXSig 對象,這裏 XXX 爲您要實現的函數簽名,本例中以 MultiplyIntUnsigned 爲例,可以在 TiDB 中找到其對應的函數簽名(builtinArithmeticMultiplyIntUnsignedSig
)及 實現。
Step 3:確定函數定義
- built-in 函數所在的文件名要求與 TiDB 的名稱對應,如 TiDB 中,expression 目錄下的下推文件統一以 builtin_XXX 命名,對應到 TiKV 這邊,就是
builtin_XXX.rs
。若同名對應的文件不存在,則需要自行在同級目錄下新建。對於本例,當前函數存放於 TiDB 的 builtin_arithmetic.go 文件裏,對應到 TiKV 便是存放在 builtin_arithmetic.rs 中。 - 函數名稱:函數簽名轉爲 Rust 的函數名稱規範,這裏
MultiplyIntUnsigned
將會被定義爲multiply_int_unsigned
。 -
函數返回值,可以參考 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>>
。 -
函數的參數, 所有 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 的基礎知識,嘗試移步 此文,看看是否有幫助。