注册用户定义函数 — register_scalar_function • Arrow R 包

这些函数支持从查询引擎执行中调用 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)

参数

名称: 要在 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）的函数，其中 types 是一个 DataType 的 list()。如果是一个函数，它必须返回一个 DataType。
auto_convert: 使用 TRUE 在传递给 fun 之前转换输入，并从输出构造正确类型的 Array。使用此选项编写 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()
}