PySpark实现字典/Map类型列按行求和并新增总和列

该需求可直接在PySpark中实现，无需强制转换为pandas DataFrame。你之前的代码逻辑错误点在于：炸开Map列后未保留原行的关联关系，直接按key全局分组聚合，得到的是所有行同key的总和，而非单行内所有value的总和。

方案1：原生Spark高阶函数实现（推荐，性能最优） #

Spark 2.4及以上版本支持数组高阶聚合函数，无需炸列、无shuffle开销，是大数据量场景下的最优选择：

• 先用map_values提取Map类型列中所有value，返回整型数组
• 用aggregate高阶函数遍历数组完成逐行累加

from pyspark.sql import functions as F

df_with_total = df.withColumn(
    "total_value",
    F.aggregate(
        F.map_values("count"),
        F.lit(0),
        lambda acc, current_val: acc + current_val
    )
)

执行后返回的结果与预期完全一致：第一行2+4+5+1+9=21，第二行3+8+3+8+1=23。

方案2：低版本Spark兼容写法 #

如果你的集群Spark版本低于2.4，不支持高阶函数，可以通过炸开Map后按原行ID分组求和，再关联回原表的方式实现：

from pyspark.sql import functions as F

# 按行炸开Map，以原表ID为分组维度计算单行总和
row_sum_df = df.select(
    "ID",
    F.explode("count").alias("map_key", "map_val")
).groupBy("ID").agg(F.sum("map_val").alias("total_value"))

# 关联回原表得到完整结果
df_with_total = df.join(row_sum_df, on="ID", how="left")

该方案存在explode和join操作，会产生额外shuffle，性能低于方案1，仅作低版本兼容使用。

方案3：转换为pandas实现（不推荐） #

仅当数据集规模极小、能完全加载进Driver内存时可使用该方案，大数据量下极易触发Driver内存溢出：

import pandas as pd

pdf = df.toPandas()
pdf["total_value"] = pdf["count"].apply(lambda map_obj: sum(map_obj.values()))
df_with_total = spark.createDataFrame(pdf)

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