PySpark中按列分组后对数组列逐元素求和的高效实现

针对你的需求，我们可以利用PySpark的高阶函数直接在数组层面完成分组聚合，避免展开数组为多列的高开销操作，非常适合处理20亿条记录、1500元素数组的大规模场景。

核心思路 #

按CATEGORY分组后，收集同组的所有数组，然后通过aggregate和zip_with函数对数组进行对应位置元素累加，全程无需展开数组。

代码实现 #

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

# 初始化Spark会话
spark = SparkSession.builder.appName("array_elementwise_sum").getOrCreate()

# 构造示例数据
data = [
    ("A", [4, 5, 6]),
    ("A", [1, 2, 3]),
    ("B", [7, 8, 9])
]
df = spark.createDataFrame(data, ["CATEGORY", "VALUE"])

# 获取数组长度（假设所有数组长度一致）
array_length = df.select(F.size("VALUE")).first()[0]

# 分组聚合：对应位置求和
result_df = df.groupBy("CATEGORY").agg(
    F.aggregate(
        F.collect_list("VALUE"),  # 收集同组的所有数组
        F.array([F.lit(0)] * array_length),  # 初始化累加器为全0数组
        lambda acc, arr: F.zip_with(acc, arr, lambda x, y: x + y)  # 对应位置相加
    ).alias("VALUE")
)

result_df.show(truncate=False)

输出结果 #

+--------+---------+
|CATEGORY|VALUE    |
+--------+---------+
|A       |[5, 7, 9]|
|B       |[7, 8, 9]|
+--------+---------+

方案优势 #

• 低开销：无需将数组展开为1500列，避免了数据膨胀和大量列的聚合计算，大幅降低IO和内存消耗
• 高效执行：利用PySpark内置高阶函数的向量式操作，执行计划经过优化，适合处理超大规模数据
• 通用性：只要数组长度固定，即可直接适配1500元素的场景，无需修改核心逻辑

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