pyspark -- best way to sum values in column of type Array(Integer())

Question

Lets say this is my dataframe ...

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

Desired output is something like

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

I made a UDF along the lines of ....

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

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

However, I have learned that UDFs are relatively slow to pure pySpark functions.

Any way to do code above in pySpark without a UDF ?

Answer 1

for Spark 3.1+, you could simply call pyspark.sql.functions.aggregate:

import pyspark.sql.functions as F
df = df.withColumn(
    "Total",
    F.aggregate("scores", F.lit(0), lambda acc, x: acc + x)
)

Notice that you should use F.lit(0.0) if the column type is not integer.