![file]()
原文：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<MonotonicityNode> {
Ok(MonotonicityNode::Function(Monotonicity::default(), None))
}
fn return_type(&self, args: &[DataType]) -> Result<DataType>;
fn nullable(&self, _input_schema: &DataSchema) -> Result<bool>;
fn eval(&self, columns: &DataColumnsWithField, _input_rows: usize) -> Result<DataColumn>;
}
如何理解？
剖析一下上述 trait 中函数的意义：
- name ⇒ 表示这个函数的名称，如
log、sign。不过有时我们应该将名称存储在函数内部，因为不同的名称可能共享同一个函数，如 pow 和 power，我们可以将 power 作为 pow 的别名(同义词)函数。
- num_arguments ⇒ 表示 编写的标量函数 可以接受多少个参数。
- variadic_arguments ⇒ 标记该函数可以接受可变参数。例如，
round() 接受一个或两个函数，它的范围是 [1,2]，我们在这里使用封闭区间。
- get_monotonicity ⇒ 表示这个函数的单调性，标明它可以用来优化执行。
- return_type ⇒ 表示该函数的返回类型，我们也可以在该函数中校验
args 。
- nullable ⇒ 表示是否可以返回一个可空字段的列(目前来说，返回true/false即可)。
- eval ⇒ eval是执行 ScalarFunction 的主函数：
columns → 输入列input_rows → 输入行数
我们将在下面解释如何编写eval函数。
前置知识 ¶
在编写eval函数之前，你可能需要以下这些知识。
数据类型
在 Databend 中数据类型分为：逻辑类型和物理类型两个形态。
逻辑数据类型是我们在Databend中使用的数据类型，物理数据类型是我们在执行/计算引擎中使用的数据类型。例如 Date32，它是一种逻辑数据类型，但是它的物理类型是 Int32，所以它的列由 DFInt32Array 表示。
内部有几种方式返回上述的两种数据类型：
- 通过
DataField 的 data_type() 获取逻辑数据类型
- 通过
DataColumn 的data_type() 获取物理数据类型
DataColumnsWithField 有 data_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),
}
一些指导方针 ¶
- 列转换
为了执行标量函数，我们肯定需要对参数的输入列进行遍历。由于我们已经在 return_type 函数中校验了数据类型，所以我们可以使用 i32 函数将输入列强制转换为特定类型的列，比如DFInt32Array 。
- 常量列
我们前面提到：应该注意 常量列匹配情况，以提高内存使用。
- 列迭代和有效位图的结合
当我们需要迭代列，我们可以使用 column.iter() 来生成一个迭代器，迭代项为 Option<T>，None 说明 null。不过这种办法比较低效，因为我们每次在循环内迭代列时都需要检查空值，这会污染CPU缓存。
根据 Arrow 的内存布局，我们可以直接利用 原始列的有效性位图 来表示空值。所以我们有ArrayApply trait 来协助你迭代列。如果有两个zip迭代器，我们可以使用 binary 函数来合并两列的有效性位图：
let array: DFUInt8Array = self.apply_cast_numeric(|a| {
AsPrimitive::<u8>::as_(a - (a / rhs) * rhs)
});
binary(x_series.f64()?, y_series.f64()?, |x, y| x.pow(y))
- Nullable 检查
有时 Nullable 很令人讨厌，但在大多数情况下我们可以接受 DataType::Null 参数。
- 隐式转换
Databend 可以接受隐式转换，例如：pow('3', 2), sign('1232') 我们可以使用cast_with_type 将参数转换到特定的列。
参考 ¶
正如你在上面所看到的, 在 Databend 中添加一个新的标量函数并不像你想象的那么难。不过在你开始添加之前, 也可以参考其他标量函数的例子, 如 sign, expr, tan, atan 。
测试 ¶
为了成为一名优秀的工程师，不要忘记测试你的代码，请在你完成新的标量函数后添加单元测试和无状态测试。
总结 ¶
我们欢迎所有社区用户为 Databend 贡献更强大的功能。如果您发现任何问题，也请随时在GitHub上 给我们提一个issue，我们会尽最大努力帮助你。
本文由博客一文多发平台 OpenWrite 发布！
