平均元素明智 pyspark 数据帧每行中的密集向量列表

Question

我在pyspark数据帧中有一列，其中包含密集向量列表。不同的行可能具有不同大小的列表，但列表中的每个向量大小相同。我想计算每个列表的元素平均值。

更具体地说，假设我有以下 df：

|ID        | Column                                      |
| -------- | ------------------------------------------- |
| 0        | List(DenseVector(1,2,3), DenseVector(2,4,5))|
| 1        | List(DenseVector(1,2,3))                    |
| 2        | List(DenseVector(2,2,3), DenseVector(2,4,5))|

我想得到的是

|ID        | Column              |
| -------- | --------------------|
| 0        | DenseVector(1.5,3,4)|
| 1        | DenseVector(2,4,5)  |
| 2        | DenseVector(2,3,4)  |

非常感谢！

Answer 1

我认为没有直接的 pyspark 函数来做到这一点。这里有一个ElementwiseProduct(与此处的预期不同)和其他。因此，您可以尝试通过udf来实现这一点。

from pyspark.sql import functions as F
from pyspark.ml.linalg import Vectors, VectorUDT
def elementwise_avg(vector_list):
x = y = z = 0
no_of_v = len(vector_list)
for i, elem in enumerate(vector_list):
x += elem[i][0]
y += elem[i][1]
z += elem[i][2]
return Vectors.dense(x/no_of_v,y/no_of_v,z/no_of_v)
elementwise_avg_udf = F.udf(elementwise_avg, VectorUDT())
df = df.withColumn("Elementwise Avg", elementwise_avg_udf("Column"))