使用PySpark中的pandas\u udf平均分配组任务

我有一个Spark数据框，其中包含多组训练数据。每组由“组”列标识

我想为每个组并行地训练一个python ML模型（在我的例子中是lightgbm）

因此，我有以下工作代码：

schema = T.StructType([T.StructField("group_id", T.IntegerType(), True),
                       T.StructField("model", T.BinaryType(), True)])

@F.pandas_udf(schema, F.PandasUDFType.GROUPED_MAP)
def _fit(pdf):
    group_id = pdf.loc[0, "group"]
    X = df.loc[: X_col]
    y = df.loc[:, y_col].values

    # train model
    model = ...

    out_df = pd.DataFrame(
        [[group_id, pickle.dumps(model)],
         columns=["group_id", "model"]]
    )

    return out_df

df.groupby("group").apply(_fit)

我在数据集中有10个组和10个工作节点

大多数情况下，每个组都分配给一个执行者，处理速度非常快

但是，有时会将多个组分配给一个执行者，而其他一些执行者则可以自由活动

这导致处理变得非常缓慢，因为执行者必须同时训练多个模型

问题：我如何安排每个小组在单独的执行人身上进行培训以避免此问题？

schema = T.StructType([T.StructField("group_id", T.IntegerType(), True),
                       T.StructField("model", T.BinaryType(), True)])

@F.pandas_udf(schema, F.PandasUDFType.GROUPED_MAP)
def _fit(pdf):
    group_id = pdf.loc[0, "group"]
    X = df.loc[: X_col]
    y = df.loc[:, y_col].values

    # train model
    model = ...

    out_df = pd.DataFrame(
        [[group_id, pickle.dumps(model)],
         columns=["group_id", "model"]]
    )

    return out_df

df.groupby("group").apply(_fit)