我有一个pyspark数据框，我正在尝试为我的数据中的每个组拟合一个模型。到目前为止，我使用循环来遍历每个组，并且代码可以正常工作，但是因为我有600个组，所以要花费很多时间（3个组需要10分钟），并且最终需要OOM。 是否有人通过分组保持工作分配来创建模型的有效方法？ （我想到的问题的解决方案是创建一个对角矩阵并为整个矩阵训练一个回归，但我宁愿有另一个解决方案）

unique_groups = regression_data.select("id")
unique_groups_list = list(set([str(row.id) for row in unique_groups.collect()]))[:3]

for i in unique_contracts_list:
    df = regression_data.filter(F.col("id") == i)
    lr = LinearRegression(featuresCol = "features",labelCol="output",maxIter=7,regParam=0.0,elasticNetParam=0.0)
    lr_model = lr.fit(df)
    df_output = first_df \
        .withColumn("id",F.lit(i)) \
        .withColumn("coefficients",F.lit(str(lr_model.coefficients))) \
        .withColumn("intercept",F.lit(str(lr_model.intercept)))
    prediction_df = prediction_df.union(df_output)