Databend 内置标量函数开发指南

Databend 内置标量函数开发指南
文章图片

原文: https://databend.rs/development/how-to-write-scalar-functions/
什么是标量函数?? 标量函数(有时被称为用户自定义函数/UDF)为每条记录返回一个单一的值,而不是作为一个结果集,并且可以在查询或SET语句中的大多数地方使用,除了FROM子句。
One to One Mapping execution┌─────┐┌──────┐ │a││x│ ├─────┤├──────┤ │b││y│ ├─────┤ScalarFunction├──────┤ │c││z│ ├─────┼────────────────────?──────┤ │d│Exec│u│ ├─────┤├──────┤ │e││v│ ├─────┤├──────┤ │f││w│ └─────┘└──────┘

trait介绍 ? 所有的标量函数都实现了 Function trait,我们把这些函数注册到一个全局静态的FunctionFactory 中,这个工厂只是一个索引map,key: 标量函数名(function name)。
?? Databend 中的函数名称是不区分大小写的。
pub trait Function: fmt::Display + Sync + Send + DynClone { fn name(&self) -> &str; fn num_arguments(&self) -> usize { 0 }// (1, 2) means we only accept [1, 2] arguments // None means it's not variadic function fn variadic_arguments(&self) -> Option<(usize, usize)> { None }// return monotonicity node, should always return MonotonicityNode::Function fn get_monotonicity(&self, _args: &[MonotonicityNode]) -> Result { Ok(MonotonicityNode::Function(Monotonicity::default(), None)) }fn return_type(&self, args: &[DataType]) -> Result; fn nullable(&self, _input_schema: &DataSchema) -> Result; fn eval(&self, columns: &DataColumnsWithField, _input_rows: usize) -> Result; }

如何理解?
剖析一下上述 trait 中函数的意义:
  • name ? 表示这个函数的名称,如 logsign。不过有时我们应该将名称存储在函数内部,因为不同的名称可能共享同一个函数,如 powpower,我们可以将 power 作为 pow 的别名(同义词)函数。
  • num_arguments ? 表示 编写的标量函数 可以接受多少个参数。
  • variadic_arguments ? 标记该函数可以接受可变参数。例如,round() 接受一个或两个函数,它的范围是 [1,2],我们在这里使用封闭区间。
  • get_monotonicity ? 表示这个函数的单调性,标明它可以用来优化执行。
  • return_type ? 表示该函数的返回类型,我们也可以在该函数中校验 args
  • nullable ? 表示是否可以返回一个可空字段的列(目前来说,返回true/false即可)。
  • eval ? eval是执行 ScalarFunction 的主函数:
    1. columns → 输入列
    2. input_rows → 输入行数
我们将在下面解释如何编写eval函数。
前置知识 ? 在编写eval函数之前,你可能需要以下这些知识。
数据类型
在 Databend 中数据类型分为:逻辑类型和物理类型两个形态。
逻辑数据类型是我们在Databend中使用的数据类型,物理数据类型是我们在执行/计算引擎中使用的数据类型。例如 Date32,它是一种逻辑数据类型,但是它的物理类型是 Int32,所以它的列由 DFInt32Array 表示。
内部有几种方式返回上述的两种数据类型:
  1. 通过 DataFielddata_type() 获取逻辑数据类型
  2. 通过 DataColumndata_type() 获取物理数据类型
  3. DataColumnsWithFielddata_type(),该函数返回逻辑数据类型
内存布局
Databend 的内存布局是基于 Arrow 的。关于 Arrow 的内存布局,你可以在[这里]了解。
就拿原始类型 int32 组成的数组举个例。[1, null, 2, 4, 8] 看起来像这样:
* Length: 5, Null count: 1 * Validity bitmap buffer:|Byte 0 (validity bitmap) | Bytes 1-63| |-------------------------|-----------------------| | 00011101| 0 (padding)|* Value Buffer:|Bytes 0-3| Bytes 4-7| Bytes 8-11| Bytes 12-15 | Bytes 16-19 | Bytes 20-63 | |------------|-------------|-------------|-------------|-------------|-------------| | 1| unspecified | 2| 4| 8| unspecified |

常量列
有时候column在block中是常量,例如:select 3 from table, column: 3 恒为3,所以我们可以用一个常量列来表示它。这在计算期间对于节省内存非常有用。
因此,Databend 的 DataColumn 表示为:
pub enum DataColumn { // Array of values. Series is wrap of arrow's array Array(Series), // A Single value. Constant(DataValue, usize), }

一些指导方针 ?
  1. 列转换
为了执行标量函数,我们肯定需要对参数的输入列进行遍历。由于我们已经在 return_type 函数中校验了数据类型,所以我们可以使用 i32 函数将输入列强制转换为特定类型的列,比如DFInt32Array
  1. 常量列
我们前面提到:应该注意 常量列匹配情况,以提高内存使用。
  1. 列迭代和有效位图的结合
当我们需要迭代列,我们可以使用 column.iter() 来生成一个迭代器,迭代项为 Option,None 说明 null。不过这种办法比较低效,因为我们每次在循环内迭代列时都需要检查空值,这会污染CPU缓存。
根据 Arrow 的内存布局,我们可以直接利用 原始列的有效性位图 来表示空值。所以我们有ArrayApply trait 来协助你迭代列。如果有两个zip迭代器,我们可以使用 binary 函数来合并两列的有效性位图:
  • ArrayApply
let array: DFUInt8Array = self.apply_cast_numeric(|a| { AsPrimitive::::as_(a - (a / rhs) * rhs) });

  • binary
binary(x_series.f64()?, y_series.f64()?, |x, y| x.pow(y))

  1. Nullable 检查
【Databend 内置标量函数开发指南】有时 Nullable 很令人讨厌,但在大多数情况下我们可以接受 DataType::Null 参数。
  1. 隐式转换
Databend 可以接受隐式转换,例如:pow('3', 2), sign('1232') 我们可以使用cast_with_type 将参数转换到特定的列。
参考 ? 正如你在上面所看到的, 在 Databend 中添加一个新的标量函数并不像你想象的那么难。不过在你开始添加之前, 也可以参考其他标量函数的例子, 如 sign, expr, tan, atan
测试 ? 为了成为一名优秀的工程师,不要忘记测试你的代码,请在你完成新的标量函数后添加单元测试和无状态测试。
总结 ? 我们欢迎所有社区用户为 Databend 贡献更强大的功能。如果您发现任何问题,也请随时在GitHub上 给我们提一个issue,我们会尽最大努力帮助你。

    推荐阅读