基于组的高效执行udf的方法

Question

我是pyspark的新手，并且遇到性能问题。给定数据框

| Group | Id   |  Selected |
| ----- | ---- | --------- |
|   A   |  id1 |      0    |
|   A   |  id2 |      0    |
|   A   |  id3 |      0    |
|   B   |  id4 |      0    |
|   B   |  id5 |      0    |

还有一个采样字典

sample_dict = {'A': 2, 'B': 1}

我想根据抽样字典随机选择（更新“选定的列”）数据框。例如

| Group |  Id  |  Selected |
| ----- | ---- | --------- |
|   A   |  id1 |     *1*   |
|   A   |  id2 |      0    |
|   A   |  id3 |     *1*   |
|   B   |  id4 |      0    |
|   B   |  id5 |     *1*   |

当前，我只能遍历每个组，获取子组，随机选择ID并根据ID更新原始数据框

for group in sample_dict:
  sub_df = df.filter(col('Group') == group)
  # id_list = a_udf_randomly_sample_n_Id_from_sub_df
  # update df['Selected'] if df['Id'].isin(id_list)

此方法的问题是顺序处理（每个组一个接一个地执行）。如果组和总行的数量增加，则pyspark代码的运行速度比简单的pandas版本（在数据块上）慢。

对于pyspark中的此问题，您是否可以共享更好的方法（例如，并行处理每个组中的事件选择）？

非常感谢您