如何在Pyspark数据框中进行分组后收集地图？

2024-05-10 • 问答

我有一个像这样的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。但是还有其他方法（更短或更有效）吗？

您可以从pyspark.sql.functions.map_from_entries使用此功能如果我们认为您的数据框为df，则应执行以下操作：

if (bill.has(keyInt))
  bill.delete(keyInt)
else
  bill.add(keyInt)

这是我的方法。

代码

import pyspark.sql.functions as F
from pyspark.sql.types import StringType

df = spark.createDataFrame([(1,'t1','a'),(1,'t2','b'),(2,'t3','t4','c'),'t5',\
(3,'t6',(3,'t7','t8','a')],\
('id','time','cat'))

(df.groupBy(['id','cat'])
   .agg(F.count(F.col('cat')).cast(StringType()).alias('counted'))
   .select(['id',F.concat_ws('->',F.col('cat'),F.col('counted')).alias('arrowed')])
   .groupBy('id')
   .agg(F.collect_list('arrowed'))
   .show()
)

输出

+-------+---------------------+
|     id|collect_list(arrowed)|
+-------+---------------------+
|   1   |    [a -> 1,b -> 1] |
|   3   |            [a -> 3] |
|   2   |    [b -> 2,c -> 1] |
+-------+---------------------+

编辑

(df.groupBy(['id','cat'])
   .count()
   .select(['id',F.create_map('cat','count').alias('map')])
   .groupBy('id')
   .agg(F.collect_list('map').alias('cat'))
   .show()
)

#+---+--------------------+
#| id|                 cat|
#+---+--------------------+
#|  1|[[a -> 1],[b -> 1]]|
#|  3|          [[a -> 3]]|
#|  2|[[b -> 2],[c -> 1]]|
#+---+--------------------+

类似于yasi的答案

import pyspark.sql.functions as F

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

如何在Pyspark数据框中进行分组后收集地图？

nanyun1 回答：如何在Pyspark数据框中进行分组后收集地图？

大家都在问