从 pyspark 数据框创建 multidict

Question

我是 pyspark 的新手，想从 pyspark 数据框创建字典。我确实有一个可以工作的熊猫代码，但我需要 pyspark 中的等效命令，但不知何故我无法弄清楚如何去做。

df = spark.createDataFrame([
(11, 101, 5.9),
(11, 102, 5.4),
(22, 111, 5.2),
(22, 112, 5.9),
(22, 101, 5.7),
(33, 101, 5.2),
(44, 102, 5.3),
], ['user_id', 'team_id', 'height'])
df = df.select(['user_id', 'team_id'])
df.show()

-------+-------+
|user_id|team_id|
+-------+-------+
|     11|    101|
|     11|    102|
|     22|    111|
|     22|    112|
|     22|    101|
|     33|    101|
|     44|    102|
+-------+-------+


df.toPandas().groupby('user_id')[
        'team_id'].apply(list).to_dict()


Result: 
{11: [101, 102], 22: [111, 112, 101], 33: [101], 44: [102]}

在 pyspark 中寻找有效的方法来创建上述 multidict。

Answer 1

您可以将team_id 列聚合为列表，然后使用collectAsMap 方法将rdd 收集为字典：

mport pyspark.sql.functions as F

df.groupBy("user_id").agg(F.collect_list("team_id")).rdd.collectAsMap()
# {33: [101], 11: [101, 102], 44: [102], 22: [111, 112, 101]}