pyspark-按元素聚合(和)矢量

时间:2019-01-24 20:34:26

标签: apache-spark pyspark

我有一个看似简单的问题,但我一直不停地往墙上撞。我实质上是在尝试与this post做相同的事情,除了我不在乎那个帖子的“分组依据”方面,我只想对所有行进行汇总。 / p>

为了解释链接的帖子,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.

2 个答案:

答案 0 :(得分:0)

我认为您必须先将向量列转换为数组,然后才能进行汇总。

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)

答案 1 :(得分:0)

我最终想通了(我在撒谎,我的一位同事为我想通了),所以我将答案贴在这里,以防有人遇到相同的问题。

您可以使用fold,类似于原始问题中链接的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_cols,它只是一个0的数组。第二个函数是应用于数组并用于遍历数据帧的每一行的函数。因此,对于每一行,它都会执行lambda a,b: [x + y for x, y in zip(a, b)],它只会将该行逐元素添加到到目前为止已计算的内容中。

您可以在原始问题中使用我的代码来生成玩具数据框以对其进行测试。希望对某人有帮助。