PySpark实现按元素对数组列求和：仅使用Spark内置函数

要解决大规模数据下的OOM问题，核心是避免将所有数组集中到单节点内存，改用Spark分布式内置函数处理。以下是两种可行方案：

方案一：固定长度数组（通用高效） #

适用于数组长度固定的场景，通过拆分元素位置实现分布式求和：

import pyspark.sql.functions as F

data = [
    ([1, 2, 3],),
    ([1, 2, 3],),
    ([1, 2, 3],),
]

df = spark.createDataFrame(data, ["array_column"])

# 1. 拆分数组元素与对应位置索引
exploded_df = df.select(F.posexplode("array_column").alias("pos", "val"))

# 2. 按位置分组，对同位置元素求和
summed_df = exploded_df.groupBy("pos").agg(F.sum("val").alias("sum_val"))

# 3. 按位置排序，聚合为最终数组
result_df = summed_df.orderBy("pos").agg(F.collect_list("sum_val").alias("summed"))

result_df.display()

原理说明： #

• posexplode将每个数组的元素拆分为独立行，同时保留元素在原数组中的位置索引，数据分布式存储在集群节点中。
• 按位置分组求和的计算逻辑由Spark分布式执行，不会将所有数据集中到单节点。
• 最后按位置排序并聚合数组，保证结果顺序与原数组元素位置一致。

方案二：可变长度数组 #

如果数组长度不固定，可先获取最大数组长度，再对每个位置的元素单独求和：

import pyspark.sql.functions as F

data = [
    ([1, 2, 3],),
    ([1, 2],),
    ([1],),
]

df = spark.createDataFrame(data, ["array_column"])

# 获取数组的最大长度
max_arr_len = df.select(F.max(F.size("array_column"))).first()[0]

# 生成每个位置的求和列
sum_columns = [
    F.sum(F.element_at("array_column", pos + 1)).alias(f"sum_pos_{pos}")
    for pos in range(max_arr_len)
]

# 聚合求和列并转为数组
result_df = df.agg(*sum_columns).select(
    F.array(*[f"sum_pos_{pos}" for pos in range(max_arr_len)]).alias("summed")
)

result_df.display()

对比原UDF方案的优势： #

原方案使用collect_list将所有数组收集到单节点内存，再通过UDF处理，数据量增大时极易触发OOM。而上述内置函数方案全程分布式计算，数据分散在集群节点中，从根源上避免了单节点内存过载问题。

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