Rust UDF for Apache Arrow

用法

将以下行添加到您的 Cargo.toml

[dependencies]
arrow-udf = "0.2"

使用#[function] 宏定义您的函数

use arrow_udf::function;

#[function("gcd(int, int) -> int", output = "eval_gcd")]
fn gcd(mut a: i32, mut b: i32) -> i32 {
    while b != 0 {
        (a, b) = (b, a % b);
    }
    a
}

该宏将生成一个接受 RecordBatch 作为输入并返回 RecordBatch 作为输出的函数。该函数可以由可选的 output 参数命名。如果未指定，它将被任意命名，例如 gcd_int4_int4_int4_eval。

然后您可以在 RecordBatch 上调用生成的函数

let input: RecordBatch = ...;
let output: RecordBatch = eval_gcd(&input).unwrap();

如果您打印输入和输出批次，它将如下所示

 input     output
+----+----+-----+
| a  | b  | gcd |
+----+----+-----+
| 15 | 25 | 5   |
|    | 1  |     |
+----+----+-----+

可能失败的函数

如果您的函数返回一个 Result

use arrow_udf::function;

#[function("div(int, int) -> int", output = "eval_div")]
fn div(x: i32, y: i32) -> Result<i32, &'static str> {
    x.checked_div(y).ok_or("division by zero")
}

输出批次将包含一个错误列。错误行将在输出列中填充NULL，错误消息将存储在错误列中。

 input     output
+----+----+-----+------------------+
| a  | b  | div | error            |
+----+----+-----+------------------+
| 15 | 25 | 0   |                  |
| 5  | 0  |     | division by zero |
+----+----+-----+------------------+

结构类型

您可以使用 StructType trait 定义一个结构类型

use arrow_udf::types::StructType;

#[derive(StructType)]
struct Point {
    x: f64,
    y: f64,
}

然后您可以在函数签名中使用该结构类型

use arrow_udf::function;

#[function("point(float8, float8) -> struct Point", output = "eval_point")]
fn point(x: f64, y: f64) -> Point {
    Point { x, y }
}

目前结构类型仅支持作为返回类型。

函数注册表

如果您想按签名查找函数，可以启用 global_registry 功能

[dependencies]
arrow-udf = { version = "0.2", features = ["global_registry"] }

在定义时，每个函数都将注册到全局注册表中。然后您可以从 REGISTRY 中查找函数

use arrow_schema::DataType::Int32;
use arrow_udf::sig::REGISTRY;

let sig = REGISTRY.get("gcd", &[Int32, Int32], &Int32).expect("gcd function");
let output = sig.function.as_scalar().unwrap()(&input).unwrap();

有关更多详细信息，请参阅示例