如何在Apache Arrow收集数据前使用R语言泛型函数？

问题背景 #

我在R中有一个泛型函数boop，会根据输入参数的类执行不同逻辑：

library(dplyr)

df <- data.frame(a = c("these", "are", "some", "strings"),
                 b = 1:4)

boop <- function(x, ...) UseMethod("boop", x)

boop.numeric <- function(x) mean(x, na.rm = TRUE)

boop.character <- function(x) mean(nchar(x), na.rm = TRUE)

df %>% summarise(across(everything(), boop))
#      a   b
# 1 4.75 2.5

现在我想在使用Apache Arrow读取Parquet文件后，在collect()之前就用boop做汇总，但Arrow的register_scalar_function需要固定的输入schema，不知道如何把R的泛型boop适配到Arrow环境中。

解决方案 #

Arrow的标量函数支持按类型重载，不需要显式定义"泛型"，只需要为每个目标Arrow类型（对应R的类）注册同名的boop实现即可，Arrow会自动根据列的类型匹配对应的函数逻辑。

1. 注册不同类型的boop函数 #

针对数值类型、字符串类型（以及可选的整数类型）分别注册boop：

library(arrow)
library(dplyr)

# 处理Arrow float64类型（对应R的numeric）
register_scalar_function(
  "boop",
  function(context, x) {
    mean(x, na.rm = TRUE)
  },
  in_type = schema(x = float64()),
  out_type = float64(),
  auto_convert = TRUE
)

# 处理Arrow string类型（对应R的character）
register_scalar_function(
  "boop",
  function(context, x) {
    mean(nchar(x), na.rm = TRUE)
  },
  in_type = schema(x = string()),
  out_type = float64(),
  auto_convert = TRUE
)

# 可选：处理Arrow int64类型（对应R的integer）
register_scalar_function(
  "boop",
  function(context, x) {
    mean(x, na.rm = TRUE)
  },
  in_type = schema(x = int64()),
  out_type = float64(),
  auto_convert = TRUE
)

2. 在Arrow数据集上使用boop #

注册完成后，即可像操作本地数据框一样使用boop，Arrow会自动根据列类型匹配对应实现，且逻辑在Arrow引擎中执行（无需提前collect()）：

# 读取Parquet数据集
pq_dataset <- open_dataset("path/to/your/parquet/files")

# 执行汇总操作，逻辑在Arrow引擎中运行
pq_dataset %>%
  summarise(across(everything(), boop)) %>%
  collect()

原理说明 #

Arrow的函数注册系统支持多态重载：当你注册同名但输入类型不同的标量函数时，Arrow执行查询时会根据列的实际类型自动匹配对应的函数实现，这和R中UseMethod的泛型函数逻辑完全等价。无需单独定义"泛型函数"，只需覆盖所有需要处理的类型即可。

内容的提问来源于stack exchange，提问作者Dan