PySpark对含Map类型列的DataFrame实现各列按规则求和

问题说明 #

待处理的PySpark DataFrame包含Map<String, Integer>类型的Item列，样例数据如下：

Date              Item (Map<Str,int>)                           Total Items      ColA
2021-02-01    Item_A -> 3, Item_B -> 10, Item_C -> 2                 15            10
2021-02-02    Item_A -> 1, Item_D -> 5, Item_E ->  7                 13            20
2021-02-03    Item_A -> 8, Item_E -> 3, Item_C ->  1                 12            30

聚合规则：

• Map类型的Item列：相同Key对应的Value累加
• 普通数值列（Total Items、ColA）：直接累加求和
  期望输出参考：

[Item_A -> 12, Item_B -> 10, Item_C -> 3, Item_D -> 5, Item_E -> 10, 40, 60]

原有实现的问题：直接对Python字典对象做+运算会触发类型错误，且未覆盖其他数值列的聚合逻辑。

实现方案 #

方案1：RDD自定义聚合（逻辑直观） #

通过自定义Map累加函数，在reduce阶段同时处理Map列和数值列的聚合：

def merge_map(map_a, map_b):
    # 合并两个Map，相同Key的Value累加
    merged = map_a.copy()
    for key, val in map_b.items():
        merged[key] = merged.get(key, 0) + val
    return merged

# 全表聚合
res = df.rdd.map(
    lambda row: (1, (row["Item"], row["Total Items"], row["ColA"]))
).reduceByKey(
    lambda x, y: (
        merge_map(x[0], y[0]),
        x[1] + y[1],
        x[2] + y[2]
    )
).collect()[0][1]

运行输出：

({'Item_A': 12, 'Item_B': 10, 'Item_C': 3, 'Item_D': 5, 'Item_E': 10}, 40, 60)

方案2：Spark SQL内置函数实现（性能更优） #

避免RDD序列化开销，直接使用DataFrame原生API聚合，适合大数据量场景：

from pyspark.sql import functions as F

# 聚合普通数值列
num_col_agg = df.agg(
    F.sum("Total Items").alias("total_items_sum"),
    F.sum("ColA").alias("col_a_sum")
).first()

# 聚合Map列：先炸开Map为KV对，按Key分组求和后再重组为Map
map_col_agg = df.select(
    F.explode("Item").alias("item_key", "item_val")
).groupBy("item_key").agg(
    F.sum("item_val").alias("item_val")
).select(
    F.map_from_entries(F.collect_list(F.struct("item_key", "item_val"))).alias("item_sum")
).first()["item_sum"]

final_res = [map_col_agg, num_col_agg["total_items_sum"], num_col_agg["col_a_sum"]]

输出结果和方案1完全一致，执行效率更高。

内容的提问来源于stack exchange，提问作者Abhishek Patil