注册用户定义函数

这些函数支持从查询引擎执行中调用 R 代码（即，在 Table 或 Dataset 上进行 dplyr::mutate() 或 dplyr::filter()）。使用 register_scalar_function() 将 Arrow 输入和输出类型附加到 R 函数，并使其可用于 dplyr 接口和/或 call_function()。标量函数是目前唯一支持的用户定义函数类型。在 Arrow 中，标量函数必须是无状态的，并且返回与输入具有相同形状（即，相同行数）的输出。

用法

register_scalar_function(name, fun, in_type, out_type, auto_convert = FALSE)

参数

name

将在 dplyr 绑定中使用的函数名

fun

一个 R 函数或 rlang 风格的 lambda 表达式。该函数将以第一个参数 context 调用，它是一个 list()，其中包含 batch_size（预期输出长度）和 output_type（所需的 DataType 输出）元素，可用于确保输出具有正确的类型和长度。后续参数按 in_types 指定的位置传递。如果 auto_convert 为 TRUE，则在传递给 fun 之前将后续参数转换为 R 向量，并通过 as_arrow_array() 自动使用预期的输出类型构建输出。

in_type

输入类型的 DataType 或具有多个参数的函数的 schema()。此签名将用于确定此函数是否适合给定的参数集。如果此函数适合多个签名，则传递上述内容的 list()。

out_type

输出类型的 DataType 或接受单个参数 (types) 的函数，它是一个 DataType 的 list()。如果为函数，则必须返回 DataType。

auto_convert

使用 TRUE 在传递给 fun 之前转换输入，并从输出构建正确类型的数组。使用此选项编写 R 对象的函数，而不是 Arrow R6 对象的函数。

返回值

NULL，隐式

示例

if (FALSE) { # arrow_with_dataset() && identical(Sys.getenv("NOT_CRAN"), "true")
library(dplyr, warn.conflicts = FALSE)

some_model <- lm(mpg ~ disp + cyl, data = mtcars)
register_scalar_function(
  "mtcars_predict_mpg",
  function(context, disp, cyl) {
    predict(some_model, newdata = data.frame(disp, cyl))
  },
  in_type = schema(disp = float64(), cyl = float64()),
  out_type = float64(),
  auto_convert = TRUE
)

as_arrow_table(mtcars) %>%
  transmute(mpg, mpg_predicted = mtcars_predict_mpg(disp, cyl)) %>%
  collect() %>%
  head()
}