Apache Arrow 的 JavaScript UDF

使用方法

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

[dependencies]
arrow-udf-js = "0.3"

创建一个 Runtime 并以字符串形式定义您的 JS 函数。请注意，函数必须导出，并且其名称必须与传递给 add_function 的名称匹配。

use arrow_udf_js::{Runtime, CallMode};

let mut runtime = Runtime::new().unwrap();
runtime
    .add_function(
        "gcd",
        arrow_schema::DataType::Int32,
        CallMode::ReturnNullOnNullInput,
        r#"
        export function gcd(a, b) {
            while (b != 0) {
                let t = b;
                b = a % b;
                a = t;
            }
            return a;
        }
        "#,
    )
    .unwrap();

然后您可以在 RecordBatch 上调用 JS 函数

let input: RecordBatch = ...;
let output: RecordBatch = runtime.call("gcd", &input).unwrap();

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

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

对于返回集合的函数（或所谓的表函数），将函数定义为生成器

use arrow_udf_js::{Runtime, CallMode};

let mut runtime = Runtime::new().unwrap();
runtime
    .add_function(
        "range",
        arrow_schema::DataType::Int32,
        CallMode::ReturnNullOnNullInput,
        r#"
        export function* range(n) {
            for (let i = 0; i < n; i++) {
                yield i;
            }
        }
        "#,
    )
    .unwrap();

然后您可以通过 call_table_function 调用表函数

let chunk_size = 1024;
let input: RecordBatch = ...;
let outputs = runtime.call_table_function("range", &input, chunk_size).unwrap();
for result in outputs {
    let output: RecordBatch = result?;
    // do something with the output
}

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

+-----+-------+
| row | range |
+-----+-------+
| 0   | 0     |
| 2   | 0     |
| 2   | 1     |
| 2   | 2     |
+-----+-------+

JS 代码将在嵌入式 QuickJS 解释器中运行。

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

类型映射

以下表格显示了 Arrow 与 JavaScript 之间的类型映射

此 crate 还支持以下 Arrow 扩展类型