PySpark - 在不使用 UDF 或连接的情况下创建多个聚合地图列答案

【问题标题】：PySpark - create multiple aggregative map columns without using UDF or joinPySpark - 在不使用 UDF 或连接的情况下创建多个聚合地图列
【发布时间】：2022-01-03 20:35:33
【问题描述】：

我有一个类似这样的巨大数据框：

+----+-------+-------+-----+
|name|level_A|level_B|hours|
+----+-------+-------+-----+
| Bob|     10|      3|    5|
| Bob|     10|      3|   15|
| Bob|     20|      3|   25|
| Sue|     30|      3|   35|
| Sue|     30|      7|   45|
+----+-------+-------+-----+

我想要的输出：

+----+--------------------+------------------+
|name|         map_level_A|       map_level_B|
+----+--------------------+------------------+
| Bob|{10 -> 20, 20 -> 25}|         {3 -> 45}|
| Sue|          {30 -> 80}|{7 -> 45, 3 -> 35}|
+----+--------------------+------------------+

含义，按name 分组，将映射level_A 和level_B 的2 个MapType 列添加到hours 的总和中。

我知道我可以使用 UDF 或连接操作获得该输出。

但是，实际上，数据非常大，不是2个map列，而是几十个，所以join/UDF成本太高了。

有没有更有效的方法来做到这一点？

【问题讨论】：

标签： apache-spark join pyspark apache-spark-sql user-defined-functions

【解决方案1】：

您可以考虑使用窗口函数。您需要为每个由name 和level_X 分区的level_X 提供一个windowspec 来计算hours 的总和。然后按name 分组并从结构数组创建映射：

from pyspark.sql import Window
import pyspark.sql.functions as F

df = spark.createDataFrame([("Bob", 10, 3, 5), ("Bob", 10, 3, 15), ("Bob", 20, 3, 25), 
                            ("Sue", 30, 3, 35),("Sue", 30, 7, 45), ], 
                           ["name", "level_A", "level_B", "hours"])

wla = Window.partitionBy("name", "level_A")
wlb = Window.partitionBy("name", "level_B")

result = df.withColumn("hours_A", F.sum("hours").over(wla)) \
    .withColumn("hours_B", F.sum("hours").over(wlb)) \
    .groupBy("name") \
    .agg(
        F.map_from_entries(
            F.collect_set(F.struct(F.col("level_A"), F.col("hours_A")))
        ).alias("map_level_A"),
        F.map_from_entries(
            F.collect_set(F.struct(F.col("level_B"), F.col("hours_B")))
        ).alias("map_level_B")
    )

result.show()

#+----+--------------------+------------------+
#|name|         map_level_A|       map_level_B|
#+----+--------------------+------------------+
#| Sue|          {30 -> 80}|{3 -> 35, 7 -> 45}|
#| Bob|{10 -> 20, 20 -> 25}|         {3 -> 45}|
#+----+--------------------+------------------+

【讨论】：