apache-spark - pyspark - 按元素聚合(求和)向量

Question

我有一个看起来很简单的问题，但我一直用头撞墙，但没有成功。我基本上是在尝试做与 this post 相同的事情除了我不关心那个帖子的“分组依据”方面，我只想对所有行进行总结。

解释一下链接的帖子，DataFrame 看起来像:

ID,Vec
1,[0,0,5]
2,[3,3,4]
3,[0,8,1]
....

我想按元素对向量求和。

上面例子的期望输出是一行:

SumOfVectors
[3,11,10]

另一个很大的区别是 我正在使用 pyspark ，不是 Scala。我试着让 rdd.fold()工作，但要么工作不一样，要么我无法弄清楚 pyspark 中的语法。

最后一个警告是，我在 ~1MM 行的数据帧和长度 ~10k 的向量上执行此操作，因此这必须相当有效。

谢谢你的帮助!根据评论，下面是一个可重复的玩具数据框。

import numpy as np
from pyspark.ml.linalg import Vectors

n_rows = 100

pdf = np.concatenate([np.array(range(n_rows)), np.random.randn(n_rows), 3*np.random.randn(n_rows)+2, 6*np.random.randn(n_rows)-2]).reshape(n_rows,-1)
dff = map(lambda x: (int(x[0]), Vectors.dense(x[1:])), pdf)

df = spark.createDataFrame(dff,schema=["ID", "Vec"])

df.schema应该看起来像 StructType(List(StructField(ID,LongType,true),StructField(Vec,VectorUDT,true)))
只是打印 df给我 DataFrame[ID: bigint, Vec: vector]
同样重要的是， 我在 Spark 2.4

$ spark-submit --version
Welcome to
      ____              __
     / __/__  ___ _____/ /__
    _\ \/ _ \/ _ `/ __/  '_/
   /___/ .__/\_,_/_/ /_/\_\   version 2.4.0
      /_/

Using Scala version 2.11.12, OpenJDK 64-Bit Server VM, 1.8.0_191
Branch HEAD
Compiled by user ec2-user on 2018-12-07T19:51:27Z
Revision bab859f34a291cb7b3f4e724b59e1b48af69016b
Url git@aws157git.com:/pkg/Aws157BigTop
Type --help for more information.

Answer 1

我认为您必须先将向量列转换为数组，然后才能对其进行聚合。

from pyspark.ml.linalg import Vectors, VectorUDT
from pyspark.sql import functions as F
from pyspark.sql import types as T

def vec2array(v):
  v = Vectors.dense(v)
  array = list([float(x) for x in v])
  return array

vec2array_udf = F.udf(vec2array, T.ArrayType(T.FloatType()))

df = df.withColumn('Vec', vec2array_udf('Vec'))

n = len(df.select('Vec').first()[0])
bla = df.agg(F.array(*[F.sum(F.col("Vec")[i]) for i in range(n)]).alias("sum"))
bla.show(truncate=False)