如何在 PySpark 中按组执行线性回归？

【问题标题】：How to perform a Linear Regression by group in PySpark?如何在 PySpark 中按组执行线性回归？
【发布时间】：2022-06-30 08:25:13
【问题描述】：

目标是在 PySpark 中以可扩展的方式为每个用户执行线性回归。特点：x1 和 x2。输出：y

回归方程（零截距）：y = m(x1) + n(x2)

例子：

pdf = pd.DataFrame(
    {
        "user": [1, 1, 1, 2, 2, 2],
        "x1": [1, 2, 3, 1, 2, 3],
        "x2": [2, 3, 4, 5, 6, 7],
        "y": [2, 4, 6, 3, 6, 9],
    }
)
df = sc.createDataFrame(pdf)
df.show()

数据如下：

+----+---+---+---+
|user| x1| x2|  y|
+----+---+---+---+
|   1|  1|  2|  2|
|   1|  2|  3|  4|
|   1|  3|  4|  6|
|   2|  1|  5|  3|
|   2|  2|  6|  6|
|   2|  3|  7|  9|
+----+---+---+---+

【问题讨论】：

x1 和 x2 在哪里？
@Steven 已更新。

标签： machine-learning pyspark linear-regression apache-spark-ml

【解决方案1】：

我使用了 PandasUDF，它适用于我的用例。 Ben Webber 在他的 post

中解释了这一点

从他的帖子中我们可以采取以下方法：(BDR:6.4;Spark:2.4.5)

from pyspark.sql.functions import pandas_udf, PandasUDFType
from pyspark.sql.types import *
import statsmodels.api as sm
import pandas as pd

pdf = pd.DataFrame(
    {
        "user": [1, 1, 1, 2, 2, 2],
        "x1": [1, 2, 3, 1, 2, 3],
        "x2": [2, 3, 4, 5, 6, 7],
        "y": [2, 4, 6, 3, 6, 9],
    }
)
df = spark.createDataFrame(pdf)
   
schema = StructType([StructField('user', DoubleType(), True),
                     StructField('r_squared', DoubleType(), True)])  

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def train_LR(input_pd):
    usr = input_pd.iloc[0]['user']
    # Implement lineal regression, as per your needs
    model = sm.OLS(input_pd['y'], input_pd[['x1','x2']]).fit()
    R_sq = model.rsquared
    return pd.DataFrame({'user': usr, 'r_squared': R_sq }, index=[0])

results = df.groupby('user').apply(train_LR)
display(results)

请注意，UDF 正在获取仅包含 groupby 数据的 pandas 数据帧

【讨论】：