group-by - 如何在 Pyspark 数据框中分组后收集 map ？

Question

我有一个这样的 pyspark 数据框:

|  id  |  time  |  cat  |
-------------------------
   1       t1       a
   1       t2       b
   2       t3       b
   2       t4       c
   2       t5       b
   3       t6       a
   3       t7       a
   3       t8       a

现在，我想按“id”对它们进行分组，并将它们聚合成这样的 map :

|   id  |       cat       |
---------------------------
|   1   |  a -> 1, b -> 1 |
|   2   |  b -> 2, c -> 1 |
|   3   |  a -> 3         |

我想我们可以使用 pyspark sql 函数的 collect_list 将它们收集为列表，然后我可以应用一些 UDF 函数将列表转换为 dict。但是有没有其他(更短或更有效)的方法来做到这一点？

Answer 1

您可以从 pyspark.sql.functions.map_from_entries 使用此功能如果我们认为您的数据框是 df 您应该这样做:

import pyspark.sql.functions as F

df1 = df.groupby("id", "cat").count()
df2 = df1.groupby("id")\
         .agg(F.map_from_entries(F.collect_list(F.struct("cat","count"))).alias("cat"))