如何利用PySpark算力对分组数据集执行LinearRegression线性回归计算

问题根因 #

你当前代码性能差的核心原因是将所有SEQ_ID拉取到Driver端循环执行，每次过滤都要扫描全量数据集，且所有训练任务由Driver串行调度，无法利用集群分布式算力，才会出现节点利用率低、任务数爆炸、运行慢的问题。

最优实现方案（基于applyInPandas分布式分组计算） #

Spark 3.0及以上版本原生支持groupBy.applyInPandas分组自定义计算能力，完全符合Spark分布式设计逻辑，可以完美解决你的需求。

前置依赖 #

• Spark 3.0+ 运行环境
• 所有Executor节点安装pandas、scikit-learn依赖

代码实现 #

import pandas as pd
from sklearn.linear_model import LinearRegression
from pyspark.sql.types import StructType, StructField, LongType, DoubleType
from pyspark.sql.functions import col

# 定义输出结果的Schema，和你原有输出结构完全对齐
output_schema = StructType([
    StructField("SEQ_ID", LongType(), nullable=False),
    StructField("Norm_time", DoubleType(), nullable=False),
    StructField("Norm_Med", DoubleType(), nullable=False),
    StructField("prediction", DoubleType(), nullable=False)
])

# 定义单个SEQ_ID分组的训练预测逻辑，直接在Executor端运行
def train_per_group(pdf: pd.DataFrame) -> pd.DataFrame:
    X = pdf[["Norm_time"]].values
    y = pdf["Norm_Med"].values
    # 等价于你原有pyspark.LinearRegression的无正则化线性回归配置
    lr = LinearRegression(fit_intercept=True)
    lr.fit(X, y)
    pdf["prediction"] = lr.predict(X)
    return pdf[["SEQ_ID", "Norm_time", "Norm_Med", "prediction"]]

# 一行代码完成分布式分组训练+结果合并，无需循环、无需手动union
pred_final = df.groupBy("SEQ_ID").applyInPandas(train_per_group, schema=output_schema)

# 验证结果和原有逻辑完全一致
pred_final.show()

核心优势 #

• 算力充分利用：所有分组计算在Executor端分布式执行，20个节点可以全部跑满，没有Driver端调度瓶颈
• 计算效率提升：仅需扫描1次原数据集，避免原有逻辑N次全表扫描的开销，任务数从几万降到和分区数匹配的量级
• 结果完全对齐：普通最小二乘线性回归的计算结果和你原有pyspark.ml.LinearRegression的参数配置输出完全一致

可选优化 #

如果需要反复计算该结果，可以提前对数据集按SEQ_ID分区缓存，避免重复Shuffle：

df = df.repartition(col("SEQ_ID")).cache()

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