PySpark：对DataFrame每个单元格中的所有DenseVector进行逐元素求和

嘿，我来帮你搞定这个逐元素求和的需求！你要做的就是把每个单元格里的所有向量（不管是数组形式还是DenseVector）对应位置的元素加起来，得到一个合并后的向量对吧？咱们一步步来实现。

先处理数组形式的向量（你的示例数据） #

首先看你给出的示例，每个单元格是数组的数组，咱们先基于这个场景来解决。

1. 完整创建测试DataFrame

先把你给出的代码补全，创建出测试用的DataFrame：

from pyspark.sql import Row
from pyspark.sql.types import StructType, StructField, ArrayType, FloatType

schema = StructType([
  StructField("abc", ArrayType(ArrayType(FloatType())), True)
])

test_data = [
    Row([[1.0, 2.0, 3.0], [1.2, 2.2, 3.2]]),
    Row([[2.0, 3.0, 4.0], [3.1, 4.1, 5.1]]),
    Row([[3.0, 4.0, 5.0], [4.3, 5.3, 6.3]]),
    Row([[4.0, 5.0, 6.0], [5.2, 4.2, 7.2]])
]

df = spark.createDataFrame(test_data, schema)

2. 用内置函数实现逐元素求和

这里用PySpark的aggregate和zip_with函数来高效处理，不用写UDF（性能更好）：

from pyspark.sql import functions as F

# 通用版本：自动匹配向量长度，不用硬编码初始值
result_df = df.withColumn(
    "abc",
    F.aggregate(
        # 要处理的向量数组
        F.col("abc"),
        # 初始累加值：和第一个向量长度一致的全0数组
        F.transform(F.element_at(F.col("abc"), 1), lambda x: F.lit(0.0)),
        # 累加逻辑：把当前累加值和当前向量对应元素相加
        lambda acc, vec: F.zip_with(acc, vec, lambda x, y: x + y)
    )
)

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

运行后就能得到你想要的输出：

+-----------------+
|abc              |
+-----------------+
|[2.2, 4.2, 6.2]  |
|[5.1, 7.1, 9.1]  |
|[7.3, 9.3, 11.3] |
|[9.2, 9.2, 13.2] |
+-----------------+

如果是MLlib的DenseVector类型 #

如果你的单元格里存的是DenseVector（而不是数组），那需要先把DenseVector转换成数组，再用上面的方法：

from pyspark.ml.linalg import DenseVector
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, FloatType

# 定义UDF：把DenseVector转成Python列表
vec_to_array = udf(lambda vec: vec.toArray().tolist(), ArrayType(FloatType()))

# 先把每个DenseVector转成数组，再处理求和
df_with_arrays = df.withColumn("abc", F.transform(F.col("abc"), vec_to_array))

# 然后用上面的aggregate逻辑处理
result_df = df_with_arrays.withColumn(
    "abc",
    F.aggregate(
        F.col("abc"),
        F.transform(F.element_at(F.col("abc"), 1), lambda x: F.lit(0.0)),
        lambda acc, vec: F.zip_with(acc, vec, lambda x, y: x + y)
    )
)

代码解释 #

• aggregate：遍历每个单元格里的向量数组，从初始的全0数组开始，逐个累加每个向量的元素。
• zip_with：把两个数组（累加值和当前向量）对应位置的元素配对，用lambda函数执行相加操作，生成新的数组。
• transform + element_at：自动生成和向量长度一致的全0初始值，不用手动硬编码向量长度，更通用。

备注：内容来源于stack exchange，提问作者Shruti