继续浏览精彩内容
慕课网APP
程序员的梦工厂
打开
继续
感谢您的支持,我会继续努力的
赞赏金额会直接到老师账户
将二维码发送给自己后长按识别
微信支付
支付宝支付

Databend 内置标量函数开发指南

Databend
关注TA
已关注
手记 43
粉丝 3
获赞 5

file

什么是标量函数?¶

标量函数(有时被称为用户自定义函数/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 ⇒ 表示这个函数的名称,如 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,所以我们可以用一个常量列来表示它。这在计算期间对于节省内存非常有用。

因此,DatabendDataColumn 表示为:

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<T>,None 说明 null。不过这种办法比较低效,因为我们每次在循环内迭代列时都需要检查空值,这会污染CPU缓存。

根据 Arrow 的内存布局,我们可以直接利用 原始列的有效性位图 来表示空值。所以我们有ArrayApply trait 来协助你迭代列。如果有两个zip迭代器,我们可以使用 binary 函数来合并两列的有效性位图:

  • ArrayApply
let array: DFUInt8Array = self.apply_cast_numeric(|a| {
    AsPrimitive::<u8>::as_(a - (a / rhs) * rhs)
});
  • binary
binary(x_series.f64()?, y_series.f64()?, |x, y| x.pow(y))
  1. Nullable 检查

有时 Nullable 很令人讨厌,但在大多数情况下我们可以接受 DataType::Null 参数。

  1. 隐式转换

Databend 可以接受隐式转换,例如:pow('3', 2), sign('1232') 我们可以使用cast_with_type 将参数转换到特定的列。

参考

正如你在上面所看到的, 在 Databend 中添加一个新的标量函数并不像你想象的那么难。不过在你开始添加之前, 也可以参考其他标量函数的例子, 如 sign, expr, tan, atan

测试

为了成为一名优秀的工程师,不要忘记测试你的代码,请在你完成新的标量函数后添加单元测试和无状态测试。

总结

打开App,阅读手记
0人推荐
发表评论
随时随地看视频慕课网APP