PySpark：按元素对数据帧的每个单元格中的所有 DenseVector 求和

Question

我需要对每个句子进行element-wise求和ALLDenseVectors。每行的单元格中可以有多个 DenseVector（每个句子），即嵌入。 我将借助单元格中的多个列表进行解释。举个例子：

+----------------------------------+
|abc                               |
+----------------------------------+
|[[1.0, 2.0, 3.0], [1.2, 2.2, 3.2]]|
|[[2.0, 3.0, 4.0], [3.1, 4.1, 5.1]]|
|[[3.0, 4.0, 5.0], [4.3, 5.3, 6.3]]|
|[[4.0, 5.0, 6.0], [5.2, 4.2, 7.2]]|
+----------------------------------+

所需的输出是：

+-----------------+
|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]|
+-----------------+

我尝试了以下几种方法，但似乎没有任何效果。

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]])
    ]

test_df = spark.createDataFrame(data=test_data, schema=schema)

test_df.groupBy("abc").agg(array([sum(col("abc")[i].getItem(j)) for j in range(3) for i in range(len(test_df.select('abc').first()[0]))]).alias("sum")).show(truncate=False)

test_df.groupBy("abc").agg(array(*[array(*[sum(col("abc")[i].getItem(j)) for j in range(3) for i in range(len(test_df.select('abc').first()[0]))])]).alias("sum")).show(truncate=False)

test_df.groupBy("abc").agg(array(*[array(*[sum(col("abc")[i].getItem(j)) for i in range(len(test_df.select('abc').first()[0]))]) for j in range(3) ]).alias("sum")).show(truncate=False)

Answer 1

您可以使用

reduce

和

zip_with

在数组单元格内聚合。

df = (df.withColumn('abc', 
                    F.reduce('abc', 
                             F.array([F.lit(0.0), F.lit(0.0), F.lit(0.0)]), 
                             lambda acc, x: F.zip_with(acc, x, lambda iacc, ix: iacc + ix))))