Apache spark pyspark-聚合（和）向量元素

我有一个看似简单的问题，但我一直把头撞在墙上，没有成功。我基本上是在做与相同的事情，只是我不关心那篇文章的“分组方式”，我只想对所有行进行汇总

为了解释链接的帖子，数据框如下所示：

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（列表（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.

---

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

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)

---

我最终发现了这一点（我在撒谎，我的一个同事帮我发现了这一点），所以我将在这里发布答案，以防有人有同样的问题

您可以使用与原始问题中链接的scala示例类似的

折叠方法。pyspark中的语法如下所示：

# find out how many Xs we're iterating over to establish the range below
vec_df = df.select('Vec')
num_cols = len(vec_df.first().Vec)

# iterate over vector to sum each "column"    
vec_sums = vec_df.rdd.fold([0]*num_cols, lambda a,b: [x + y for x, y in zip(a, b)])

简要说明：
rdd.fold（）
接受两个参数。第一个是初始化数组，在本例中，
[0]*num\u cols
，它只是一个0的数组。第二个是应用于数组并用于迭代数据帧的每一行的函数。因此，对于每一行，它都会执行
lambda，b:[x+y代表x，y在zip（a，b）中）]
，它只是将这一行元素添加到迄今为止的计算中

您可以使用我在原始问题中的代码生成一个玩具数据框来测试这一点。希望这对其他人有所帮助。
您可以使用PandasUDF实现用户定义的聚合函数。看一看。对不起，我没有尝试这个，因为我自己发现了（答案如下）。不过，这也可能奏效。