如何在PySpark中计算数组列的按索引平均值

问题描述 #

现有Spark DataFrame包含fftAbs列，类型为ArrayType(DoubleType())，存储FFT计算后的绝对值数组。需要计算该列所有数组按索引位置的平均值：

期望结果为[1, 4, 6]，计算逻辑：

• 索引0：(0+2+1)/3 = 1
• 索引1：(1+3+8)/3 = 4
• 索引2：(2+12+4)/3 = 6

错误原因分析 #

1. 直接使用F.avg('fftAbs')报错：Spark内置的avg函数仅支持数值或interval类型，不支持数组类型，因此会触发AnalysisException类型不匹配错误。
2. 带Numpy的UDF报错：Spark的UDF序列化/反序列化机制无法处理Numpy的dtype对象，导致PickleException。

解决方案 #

方案一：使用Spark内置函数（推荐，性能更优） #

利用Spark内置函数拆分数组、计算单索引平均值，再合并为结果数组。该方案无需UDF，避免序列化问题且性能更好。

import pyspark.sql.functions as F

# 假设所有行的数组长度一致，先获取数组长度
array_length = fftDf.select(F.size('fftAbs')).first()[0]

# 生成每个索引位置的平均值计算列
avg_cols = [
    F.avg(F.col('fftAbs')[i]).alias(f'avg_{i}') 
    for i in range(array_length)
]

# 聚合后将所有平均值列合并为数组
avgDf = fftDf.agg(*avg_cols).select(
    F.array(*[f'avg_{i}' for i in range(array_length)]).alias('index_avg')
)

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

注意：该方案要求所有行的数组长度一致，若存在长度不一致的行，需先过滤或填充至统一长度。

方案二：修复自定义UDF（仅当必须使用UDF时） #

避免使用Numpy，改用纯Python逻辑实现按索引求平均，避免序列化问题：

import pyspark.sql.functions as F
import pyspark.sql.types as T

def _index_avg(two_d_list):
    if not two_d_list:
        return []
    # 转置二维列表，将同索引元素归为一组
    transposed = list(zip(*two_d_list))
    # 计算每组平均值，返回原生Python float列表
    return [sum(group)/len(group) for group in transposed]

# 注册UDF
spark_index_avg = F.udf(_index_avg, T.ArrayType(T.DoubleType(), False))

# 聚合计算
avgDf = fftDf.agg(
    spark_index_avg(F.collect_list('fftAbs')).alias('index_avg')
)

注意：UDF性能低于Spark内置函数，数据量较大时优先选择方案一。

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