apache-spark - pyspark-对Array(Integer())类型的列中的值求和的最佳方法

Question

可以说这是我的数据框...

name | scores
Dan  |  [10,5,2,12]
Ann  |  [ 12,3,5]
Jon  |  [ ] 

所需的输出类似于

name | scores         | Total
Dan  |  [10,5,2,12]   | 29
Ann  |   [ 12,3,5]    | 20
Jon  |  [ ]           | 0

我按照...的方式制作了UDF。

sum_cols = udf(lambda arr: if arr == [] then 0 else __builtins__.sum(arr),IntegerType())

df.withColumn('Total', sum_cols(col('scores'))).show()

但是，我了解到，UDF在使用纯pySpark函数方面相对较慢。

在没有UDF的情况下，可以在pySpark中执行以上代码吗？

Answer 1

您可以使用更高阶的SQL函数AGGREGATE(从函数式编程中减少)，如下所示:

import pyspark.sql.functions as F
df = df.select(
  'name',
  F.expr('AGGREGATE(scores, 0, (acc, x) -> acc + x)').alias('Total')
)

第一个参数是数组列，第二个参数是初始值(应与您求和的值具有相同的类型，因此，如果输入不是整数，则可能需要使用“0.0”或“DOUBLE(0)”等)，第三个参数是是lambda函数，它将数组的每个元素添加到一个累加器变量(在开始时，它将被设置为初始值0)。

转换将在单个投影运算符中运行，因此非常有效。另外，您不必事先知道数组的大小，并且数组在每一行上的长度可以不同。