PySpark与numpy互操作：如何高效对ArrayType列执行numpy数学运算

完全可以通过Pandas UDF简化你的实现逻辑，Pandas UDF基于Apache Arrow实现Spark和Python之间的类型自动转换，不需要手动做numpy类型到Python原生类型的适配，同时性能比普通Python UDF高一个数量级以上，完全适配你的业务场景。

通用UDF封装实现 #

单返回值场景（标量/数组/复数结果） #

你只需要通过装饰器声明返回的Spark类型，内部直接调用numpy函数即可，不同返回类型的封装逻辑高度一致：

import pyspark.sql.functions as F
from pyspark.sql.types import DoubleType, ArrayType, StructType, StructField
import numpy as np
import pandas as pd

# 标量返回：中位数
@F.pandas_udf(returnType=DoubleType())
def np_median(s: pd.Series) -> pd.Series:
    return s.apply(lambda x: np.median(np.array(x)))

# 数组返回：绝对值
@F.pandas_udf(returnType=ArrayType(DoubleType()))
def np_abs(s: pd.Series) -> pd.Series:
    return s.apply(lambda x: np.abs(np.array(x)).tolist())

# 复数数组返回：FFT结果，每个复数拆分为[实部, 虚部]存储
@F.pandas_udf(returnType=ArrayType(ArrayType(DoubleType())))
def np_fft(s: pd.Series) -> pd.Series:
    def convert_fft(arr):
        fft_res = np.fft.fft(np.array(arr))
        return [[float(v.real), float(v.imag)] for v in fft_res]
    return s.apply(convert_fft)

多返回值场景 #

如果需要一个函数同时返回多个计算结果，直接声明Struct类型作为返回值即可，不需要单独封装多个UDF：

# 声明多返回值的结构
stats_return_type = StructType([
    StructField("mean", DoubleType()),
    StructField("median", DoubleType()),
    StructField("quantile_90", DoubleType())
])

@F.pandas_udf(returnType=stats_return_type)
def array_stats(s: pd.Series) -> pd.Series:
    def calc_stats(arr):
        np_arr = np.array(arr)
        return (
            float(np.mean(np_arr)),
            float(np.median(np_arr)),
            float(np.quantile(np_arr, 0.9))
        )
    return s.apply(calc_stats)

批量处理上百列的简化写法 #

针对你有大量同类列需要处理的场景，可以通过循环批量生成计算逻辑，不需要逐行写withColumn：

# 待处理的列列表，可根据实际业务动态生成
process_cols = ["P1", "P2"]
calc_exprs = []

for col_name in process_cols:
    # 批量添加每列的计算逻辑
    calc_exprs.append(np_median(F.col(col_name)).alias(f"{col_name}_median"))
    calc_exprs.append(np_abs(F.col(col_name)).alias(f"{col_name}_abs"))
    calc_exprs.append(np_fft(F.col(col_name)).alias(f"{col_name}_fft"))
    calc_exprs.append(array_stats(F.col(col_name)).alias(f"{col_name}_stats"))

# 一次性完成所有列的计算
df_processed = df.select("*", *calc_exprs)

调用后的使用逻辑和你原来的普通UDF完全一致，不需要修改其他业务代码。

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