从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中寻找有效的方法来创建上述多指数。

Answer 1

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

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]}