如何在R语言的sparklyR中为每个Spark分区进行并行模型训练？

sparklyr中的group_by() + ml_linear_regression()并不会按分组拟合独立模型，而是基于全量数据训练全局模型，仅保留分组元数据。要实现每个分组单独训练模型，确实需要用到spark_apply()——它是sparklyr中唯一支持对Spark数据的分组执行自定义R逻辑的工具，完全适配你的需求。

实现步骤与代码示例 #

以下是完整的分组拟合线性回归模型的代码：

library(sparklyr)
library(stats)

# 生成测试数据
set.seed(134)
n <- 10
x <- 1:n
y <- -x + rep(c(10, 20), each = 5) + rnorm(n)
df <- data.frame(
  group = rep(c("A", "B"), each = n/2),
  y = y,
  x = x
)

# 连接本地Spark集群并复制数据
sc <- spark_connect(master = "local")
sdf <- copy_to(sc, df, name = "sdf")

# 按group分组拟合独立模型
group_models <- sdf %>%
  spark_apply(
    function(group_df) {
      # 对当前分组的本地数据框拟合线性回归
      model <- lm(y ~ x, data = group_df)
      # 提取模型关键参数并返回结构化数据
      data.frame(
        group = unique(group_df$group),
        intercept = coef(model)[[1]],
        x_coefficient = coef(model)[[2]],
        r_squared = summary(model)$r.squared
      )
    },
    # 指定分组列
    group_by = "group",
    # 显式定义返回数据的Schema（提升稳定性）
    schema = struct_type(
      struct_field("group", "string"),
      struct_field("intercept", "double"),
      struct_field("x_coefficient", "double"),
      struct_field("r_squared", "double")
    )
  )

# 将结果拉回本地查看
collect(group_models)

关键说明 #

• group_by = "group"：告诉Spark按group列拆分数据，每个分组的数据集会被传递给自定义函数作为本地R数据框。
• 自定义函数内使用base R的lm()：此时处理的是本地数据，无需依赖Spark的ML库，灵活性更高。
• 指定schema：明确返回数据的字段类型，避免Spark自动推断结构时出现错误，尤其在复杂返回场景下更重要。
• 结果处理：group_models是Spark数据框，用collect()可以将结果拉回本地R环境查看。

注意事项 #

如果你的分组数据量极大，本地拟合模型可能会遇到内存瓶颈，这种情况下建议评估分组的合理性，或者考虑Spark MLlib的其他分布式建模方案，但sparklyr中针对"每个分组独立训练模型"的场景，spark_apply()是最直接的实现方式。

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