Spark DataFrame数组列行求和及posexplode逆操作可靠实现

问题场景 #

现有如下Spark DataFrame，其中values为double类型数组：

df = spark.createDataFrame([(1,1,1,  [1,2,3,4,5]),(1,1,2,  [1,2,3,4,5]), 
(1,1,3,  [1,2,3,4,5]), 
(1,1,4,  [1,2,3,4,5]), 
(1,2,5,  [1,2,3,4,5]), 
(1,2,6,  [1,2,3,4,5]), 
(1,2,7,  [1,2,3,4,5]), 
(1,3,8,  [1,2,3,4,5]), 
(1,3,9, [1,2,3,4,5]), 
(1,3,10, [1,2,3,4,5])], ["acc", "sub-acc", "item", "values"])

需求是按acc和sub-acc分组，对values数组的对应位置元素分别求和，最终得到同长度的数组小计（伪代码逻辑：df.groupby('acc','sub-acc').sum('values')）。

现有方案的问题 #

原方案通过posexplode_outer展开数组，分组求和后再collect_list合并，但在百万行数据+万级长度数组的场景下，并行处理时的排序操作无法保证位置顺序的绝对可靠，导致最终数组元素位置错位。

可靠优化方案 #

方案1：利用Spark高阶函数直接操作数组（推荐） #

无需展开数组，直接在分组后对数组列表进行位置累加，完全避免排序和数据膨胀问题，顺序绝对可靠：

from pyspark.sql import functions as F

result_df = df.groupBy("acc", "sub-acc") \
    .agg(F.collect_list("values").alias("values_list")) \
    .withColumn(
        "values",
        F.expr("""
            aggregate(
                values_list,
                cast(array_repeat(0.0, size(values_list[0])) as array<double>),
                (acc_arr, curr_arr) -> transform(arrays_zip(acc_arr, curr_arr), x -> x.acc_arr + x.curr_arr)
            )
        """)
    ) \
    .drop("values_list") \
    .orderBy("acc", "sub-acc")

# 查看结果
result_df.show(truncate=False)

逻辑说明：

1. 按acc和sub-acc分组，收集每组内的所有values数组到values_list；
2. 用aggregate初始化一个与原数组长度一致的全0数组作为累加器；
3. 遍历values_list中的每个数组，通过arrays_zip将累加器数组与当前数组按位置配对，再用transform对每对元素求和，更新累加器；
4. 最终累加器即为每组对应位置求和后的结果数组。

方案2：改进posexplode方案（确保位置顺序） #

若必须使用展开方式，可通过显式按位置索引提取结果，避免排序依赖：

from pyspark.sql import functions as F

# 展开数组并按位置求和
df_exploded = df.select("acc", "sub-acc", F.posexplode_outer(df.values).alias("pos", "val")) \
    .groupBy("acc", "sub-acc", "pos") \
    .agg(F.sum("val").alias("sum_val"))

# 获取原数组长度（确保覆盖所有位置）
arr_length = df.select(F.size("values")).first()[0]

# 按位置索引提取结果，生成有序数组
result_df = df_exploded.groupBy("acc", "sub-acc") \
    .agg(F.collect_list(F.struct("pos", "sum_val")).alias("pos_sum_list")) \
    .withColumn(
        "values",
        F.expr(f"transform(array({','.join(map(str, range(arr_length)))}), idx -> coalesce(element_at(filter(pos_sum_list, x -> x.pos = idx), 1).sum_val, 0.0))")
    ) \
    .drop("pos_sum_list") \
    .orderBy("acc", "sub-acc")

# 查看结果
result_df.show(truncate=False)

逻辑说明：

1. 展开数组后按位置分组求和，记录每个位置的总和；
2. 生成原数组长度的索引序列，通过filter和element_at按索引提取对应位置的求和值，确保顺序完全匹配原数组位置；
3. 用coalesce处理可能的空值（如原数组存在缺失位置的情况）。

结果验证 #

两种方案均可得到预期结果：

+---+-------+---------------------+
|acc|sub-acc|values               |
+---+-------+---------------------+
|1  |1      |[4.0, 8.0, 12.0, 16.0, 20.0]|
|1  |2      |[3.0, 6.0, 9.0, 12.0, 15.0]|
|1  |3      |[3.0, 6.0, 9.0, 12.0, 15.0]|
+---+-------+---------------------+

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