PySpark DataFrame每行DenseVector列表的元素级均值计算

要实现每行DenseVector列表的元素级平均，你可以通过自定义UDF（用户定义函数）来完成，具体步骤如下：

1. 导入依赖 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.ml.linalg import DenseVector, VectorType
import numpy as np

2. 定义平均计算函数 #

这个函数会接收DenseVector列表，将其转换为numpy数组后按列求平均，最后转回DenseVector：

def avg_dense_vectors(vec_list):
    # 处理空列表场景（可根据业务需求调整逻辑）
    if not vec_list:
        return DenseVector([])
    # 把所有DenseVector转为numpy数组并堆叠
    vec_array = np.array([vec.toArray() for vec in vec_list])
    # 按列计算平均值
    avg_array = np.mean(vec_array, axis=0)
    # 转回DenseVector类型
    return DenseVector(avg_array)

3. 注册UDF #

将自定义函数注册为Spark UDF，指定返回类型为VectorType：

avg_vec_udf = udf(avg_dense_vectors, VectorType())

4. 应用到DataFrame #

假设你的原始DataFrame名为df，执行以下代码替换原列：

result_df = df.withColumn("Column", avg_vec_udf(df["Column"]))

测试示例 #

你可以用以下代码创建示例DataFrame验证效果：

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

# 构造测试数据
data = [
    (0, [DenseVector([1,2,3]), DenseVector([2,4,5])]),
    (1, [DenseVector([1,2,3])]),
    (2, [DenseVector([2,2,3]), DenseVector([2,4,5])])
]

df = spark.createDataFrame(data, ["ID", "Column"])

# 应用UDF得到结果
result_df = df.withColumn("Column", avg_vec_udf(df["Column"]))
result_df.show(truncate=False)

执行后会输出你期望的结果：

+---+----------------+
|ID |Column          |
+---+----------------+
|0  |[1.5,3.0,4.0]   |
|1  |[1.0,2.0,3.0]   |
|2  |[2.0,3.0,4.0]   |
+---+----------------+

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