如何使用 foreachPartition 在 Spark 中为每个分区高效地构建一个 ML 模型？答案

【问题标题】：How can you efficiently build one ML model per partition in Spark with foreachPartition?如何使用 foreachPartition 在 Spark 中为每个分区高效地构建一个 ML 模型？
【发布时间】：2020-02-23 18:53:41
【问题描述】：

我正在尝试为我的数据集的每个分区拟合一个 ML 模型，但我不知道如何在 Spark 中做到这一点。

我的数据集基本上是这样的，按公司划分：

Company | Features | Target

A         xxx        0.9
A         xxx        0.8
A         xxx        1.0
B         xxx        1.2
B         xxx        1.0
B         xxx        0.9
C         xxx        0.7
C         xxx        0.9
C         xxx        0.9

我的目标是以并行方式为每家公司训练一个回归器（我有几亿条记录，有 10 万家公司）。我的直觉是我需要使用foreachPartition 来并行处理分区（即我的公司）并训练和保存每个公司模型。 我的主要问题是如何处理iterator 类型，该类型将在foreachPartition 调用的函数中使用。

这是它的样子：

dd.foreachPartition(

    iterator => {var company_df = operator.toDF()
                 var rg = RandomForestRegressor()
                                 .setLabelCol("target")
                                 .setFeaturesCol("features")
                                 .setNumTrees(10)
                 var model = rg.fit(company_df)
                 model.write.save(company_path)
                 }
)

据我了解，尝试将 iterator 转换为 dataframe 是不可能的，因为 RDD 的概念本身不能存在于 foreachPartition 语句中。

我知道这个问题很开放，但我真的被困住了。

【问题讨论】：

标签： scala apache-spark apache-spark-ml

【解决方案1】：

在 pyspark 中，您可以执行以下操作

import statsmodels.api as sm
# df has four columns: id, y, x1, x2

group_column = 'id'
y_column = 'y'
x_columns = ['x1', 'x2']
schema = df.select(group_column, *x_columns).schema

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
# Input/output are both a pandas.DataFrame
def ols(pdf):
    group_key = pdf[group_column].iloc[0]
    y = pdf[y_column]
    X = pdf[x_columns]
      X = sm.add_constant(X)
    model = sm.OLS(y, X).fit()

    return pd.DataFrame([[group_key] + [model.params[i] for i in   x_columns]], columns=[group_column] + x_columns)

beta = df.groupby(group_column).apply(ols)

【讨论】：

如果我切换到 PySpark 非常有用。 Scala Spark 中的任何等价物？
Pandas 是一个 python 库，您可以通过多种方式使用 JNI 之类的东西来使用 Scala 中的等价物，但实现起来会非常复杂。我想说坚持使用 PySpark 将是最干净的解决方案。