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

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

仅在聚合函数中使用cast（0作为“数据类型”）时，您不会得到错误。

Answer 2

如果您不知道数组的长度（如您的示例）：

import pyspark.sql.functions as F
psaudo_counts = df.select('name').distinct().withColumn('score', F.lit(0))
df = df.select('name', F.explode('scores').alias('score')).unionByName(psaudo_counts)
df = df.groupby('name').agg(F.sum('name').alias('Total'))

如果您知道数组的长度：

import pyspark.sql.functions as F
length_of_array = 3
df = df.select('name', sum([F.col('scores').getItem(i) for i in range(length_of_array)]).alias('Total'))

感谢cricket_007的提示，感谢old mail的固定长度建议

Answer 3

您可以使用高阶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）。

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

Answer 4

我建议使用withColumn方法进行更简单的表述/实现：

width < and = 400px

Answer 5

from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()
data = spark.createDataFrame(
    [
        ('Dan', [10, 5, 2, 12]), 
        ('Ann', [12, 3, 5]),
        ('Jon', [])
    ],
    ['name', 'scores'])

data.show()

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

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf

udf_function = udf(lambda x: int(np.array(x, dtype=int).sum()), IntegerType())
data = data.withColumn('sum', udf_function('scores'))
data.show()

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

pyspark - 在Array（Integer（））类型的列中对值进行求和的最佳方法

5 个答案: