跳到内容

这些函数支持从查询引擎执行中调用 R 代码 (例如,在 dplyr::mutate()dplyr::filter() 中对 TableDataset 进行操作)。使用 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_convertTRUE,则后续参数在传递给 fun 之前转换为 R 向量,并且输出通过 as_arrow_array() 自动构建为预期的输出类型。

in_type

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

out_type

输出类型的 DataType 或接受单个参数 (types) 的函数,该参数是一个 DataTypelist()。 如果是一个函数,它必须返回一个 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()
}