PySpark根据列键将DataFrame分成多个帧，并在每个帧上训练ML lib模型

Question

我有一个带有“组”列的PySpark数据框。我也有功能列和标签列。我想为每个组拆分数据框，然后训练模型，最后得到一个词典，其中的键是“组”名称，值是训练后的模型。

This question essentially give an answer to this problem。这种方法效率低下。

  

这里明显的问题是，它需要对每个级别进行完整的数据扫描，因此这是一项昂贵的操作。

答案是旧的，我希望从那以后PySpark会有改进。对于我的用例，我有1万个组，数据大小严重偏斜。最大的组可以有10亿条记录，最小的组可以有1条记录。

编辑：如此处所建议，这是一个可复制的小示例。

df = sc.createDataFrame(
    [
        ('A', 1, 0, True),
        ('A', 3, 0, False),
        ('B', 2, 2, True),
        ('B', 3, 3, True),
        ('B', 5, 2, False)
    ],
    ('group', 'feature_1', 'feature_2', 'label')
)

我可以按照以上链接中的建议拆分数据：

from itertools import chain
from pyspark.sql.functions import col

groups = chain(*df.select("group").distinct().collect())
df_by_group = {group: 
  train_model(df.where(col("group").eqNullSafe(group))) for group in groups}

其中train_model是一个函数，该函数获取具有column = {[feature_1, feature_2, label]的数据帧并在该数据帧上返回经过训练的模型。