三十分钟成为 Contributor | 为 TiKV 添加 built-in 函数

背景知识

SQL 语句发送到 TiDB 后经过 parser 生成 AST（抽象语法树），再经过 Query Optimizer 生成执行计划，执行计划切分成很多子任务，这些子任务以表达式的方式最后下推到底层的各个 TiKV 来执行。

如图 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 的https://github.com/tikv/tikv/issues/3275issue 中，找到未实现的函数签名列表，选一个您想要实现的函数。

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>
   evalReal	Result>
   evalString	Result>>
   evalDecimal	Result>>
   evalTime	Result>>
   evalDuration	Result>>
   evalJSON	Result>>

   可以看到 TiDB 的builtinArithmeticMultiplyIntUnsignedSig对象实现了 evalInt 方法，故当前函数（multiply_int_unsigned）的返回类型应该为Result>。

4. 函数的参数, 所有 builtin-in 的参数都与 Expression 的eval函数一致，即：

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

   • 该行数据每列具体值 (row:&[Datum])

综上，multiply_int_unsigned的下推函数定义为：

pubfnmultiply_int_unsigned( &self, ctx: &mutEvalContext, row: &[Datum], )->Result<Option<i64>>

Step 4：实现函数逻辑

这一块相对简单，直接对照 TiDB 的相关逻辑实现即可。这里，我们可以看到 TiDB 的builtinArithmeticMultiplyIntUnsignedSig的具体实现如下：

func(s *builtinArithmeticMultiplyIntUnsignedSig)evalInt(row types.Row)(valint64, isNullbool, err error){ a, isNull, err := s.args[0].EvalInt(s.ctx, row)ifisNull || err !=nil{return0, isNull, errors.Trace(err) } unsignedA :=uint64(a) b, isNull, err := s.args[1].EvalInt(s.ctx, row)ifisNull || err !=nil{return0, isNull, errors.Trace(err) } unsignedB :=uint64(b) result := unsignedA * unsignedBifunsignedA !=0&& result/unsignedA != unsignedB {return0,true, types.ErrOverflow.GenByArgs("BIGINT UNSIGNED", fmt.Sprintf("(%s * %s)", s.args[0].String(), s.args[1].String())) }returnint64(result),false,nil}

参考以上代码，翻译到 TiKV 即可，如下：

pubfnmultiply_int_unsigned( &self, ctx: &mutEvalContext, row: &[Datum], )->Result<Option<i64>> {letlhs= try_opt!(self.children[0].eval_int(ctx, row));letrhs= try_opt!(self.children[1].eval_int(ctx, row));letres= (lhsasu64).checked_mul(rhsasu64).map(|t| tasi64);//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]fntest_multiply_int_unsigned() {letcases=vec![ (Datum::I64(1), Datum::I64(2), Datum::U64(2)), ( Datum::I64(i64::MIN), Datum::I64(1), Datum::U64(i64::MINasu64), ), ( Datum::I64(i64::MAX), Datum::I64(1), Datum::U64(i64::MAXasu64), ), (Datum::U64(u64::MAX), Datum::I64(1), Datum::U64(u64::MAX)), ];letmutctx= EvalContext::default();for(left, right, exp)incases {letlhs=datum_expr(left);letrhs=datum_expr(right);letmutop= Expression::build( &mutctx,scalar_func_expr(ScalarFuncSig::MultiplyIntUnsigned, &[lhs, rhs]), ).unwrap(); op.mut_tp().set_flag(types::UNSIGNED_FLAGasu32);letgot= op.eval(&mutctx, &[]).unwrap();assert_eq!(got, exp); }// test overflowletcases=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)incases {letlhs=datum_expr(left);letrhs=datum_expr(right);letmutop= Expression::build( &mutctx,scalar_func_expr(ScalarFuncSig::MultiplyIntUnsigned, &[lhs, rhs]), ).unwrap(); op.mut_tp().set_flag(types::UNSIGNED_FLAGasu32);letgot= op.eval(&mutctx, &[]).unwrap_err();assert!(check_overflow(got).is_ok()); } }

Step 8：运行测试

运行 make expression，确保所有的 test case 都能跑过。

完成以上几个步骤之后,就可以给了下来V 项目提 PR 啦。想要了解提 PR 的基础知识，尝试移步此文，看看是否有帮助。

点击查看更多成为 Contributor 系列文章