如何在PySpark中基于指定列列表添加行均值列？

要实现这个需求，核心是针对指定列生成行级数组，再计算数组的平均值，下面分Spark版本给出具体实现：

方法一：Spark 3.0+ （推荐，更简洁） #

Spark 3.0及以上提供了array_avg()函数，可以直接对数组求均值，代码如下：

from pyspark.sql import SparkSession
from pyspark.sql.functions import array, array_avg

# 初始化SparkSession
spark = SparkSession.builder.appName("RowMeanCalculation").getOrCreate()

# 构造示例DataFrame
data = [
    (1, 2, 3, 4),
    (5, 6, 7, 8),
    (9, 10, 11, 12)
]
df = spark.createDataFrame(data, ["A", "B", "C", "D"])

# 定义可变列列表（可根据需求修改）
l = ["A", "C"]

# 添加行均值列
df_with_mean = df.withColumn("row_mean", array_avg(array(*l)))

# 查看结果
df_with_mean.show()

解释：

• array(*l)：把列表l中的列名转为Spark数组列
• array_avg()：计算数组中元素的平均值，自动处理数值类型的列

方法二：Spark 2.x 兼容版本 #

如果你的Spark版本低于3.0，可以用aggregate()函数手动计算总和再除以列数：

from pyspark.sql import SparkSession
from pyspark.sql.functions import array, aggregate, lit, col

spark = SparkSession.builder.appName("RowMeanCalculation").getOrCreate()

# 构造示例DataFrame（同上）
data = [
    (1, 2, 3, 4),
    (5, 6, 7, 8),
    (9, 10, 11, 12)
]
df = spark.createDataFrame(data, ["A", "B", "C", "D"])

# 定义可变列列表
l = ["A", "B", "C"]

# 计算行均值：总和 / 列数量
df_with_mean = df.withColumn(
    "row_mean",
    aggregate(
        array(*l),
        lit(0.0),  # 初始值设为浮点型，避免整数除法
        lambda acc, x: acc + x,
        lambda acc: acc / len(l)
    )
)

df_with_mean.show()

解释：

• aggregate()的三个参数：待处理的数组、累加器初始值、累加逻辑、最终计算逻辑
• 初始值用lit(0.0)确保结果是浮点型，避免整数除法导致的精度丢失

注意事项 #

• 确保列表l中的列名全部存在于DataFrame中，否则会抛出列不存在的错误
• 如果列包含非数值类型数据，需要先做类型转换，否则计算会失败

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