如何在PySpark中对DataFrame数组列执行按元素求和？

问题描述 #

我有一个PySpark DataFrame，其中包含一列名为“c1”的数组列，每行均为整数数组，示例数据如下：

需要对该列执行按元素求和（即常规向量加法），最终归约为单个数组[12, 15, 18]。希望实现类似df.select(sum("c1"))的操作，但要针对数组列执行向量/数组加法。

搜索时大多找到基于zip_with和高阶函数的解决方案，但这些方案都是假设存在多列并对行进行求和，可能是表述有误导致的。

补充说明：“c1”列仅为示例，实际场景中该列包含数千行，每个数组有数百个元素。

解决方案 #

在notNull的帮助下，得出以下解决方案：

# 先获取数组的最大元素个数，超出数组长度的位置会被设为null
max_elements = df.select(max(size("c1"))).first()[0]

# 创建用于select语句的列表达式列表，对每个位置的元素求和
selected_cols = [sum(col("c1")[i]).alias(f"elem_{i}") for i in range(max_elements)]

# 将求和后的单个元素组合成新数组
df_new = df.select(array(*selected_cols).alias("sum_result"))

注：原代码中df_mv_pr和movement_profile应为笔误，已修正为与问题描述一致的df和c1；sum作为别名存在语法冲突，改为sum_result避免问题。

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