Calculate cumulative sum of pyspark array column

user_8844585 0
Viewed 476

I have a spark dataframe with an array column that looks like this:

+--------------+
|            x |
+--------------+
| [1, 1, 0, 1] |
| [0, 0, 0, 0] |
| [0, 0, 1, 1] |
| [0, 0, 0, 1] |
|    [1, 0, 1] |
+--------------+

I want to add a new column with another array that contains the cumulative sum of x at each index. The result should look like this:

+--------------+---------------+
|            x | x_running_sum |
+--------------+---------------+
| [1, 1, 0, 1] |  [1, 2, 2, 3] |
| [0, 0, 0, 0] |  [0, 0, 0, 0] |
| [0, 0, 1, 1] |  [0, 0, 1, 2] |
| [0, 0, 0, 1] |  [0, 0, 0, 1] |
|    [1, 0, 1] |     [1, 1, 2] |
+--------------+---------------+

How can I create the x_running_sum column? I've tried using some of the higher order functions like transform, aggregate, and zip_with, but I haven't found a solution yet.

pyspark apache-spark-sql apache-spark
1 Answers

To perform a cumulative sum I sliced the array by index position and reduce the values from it:

from pyspark.sql import Row


df = spark.createDataFrame([
  Row(x=[1, 1, 0, 1]),
  Row(x=[0, 0, 0, 0]),
  Row(x=[0, 0, 1, 1]),
  Row(x=[0, 0, 0, 1]),
  Row(x=[1, 0, 1])
])

(df
 .selectExpr('x', "TRANSFORM(sequence(1, size(x)), index -> REDUCE(slice(x, 1, index), CAST(0 as BIGINT), (acc, el) -> acc + el)) AS x_running_sum")
 .show(truncate=False))

Output

+------------+-------------+
|x           |x_running_sum|
+------------+-------------+
|[1, 1, 0, 1]|[1, 2, 2, 3] |
|[0, 0, 0, 0]|[0, 0, 0, 0] |
|[0, 0, 1, 1]|[0, 0, 1, 2] |
|[0, 0, 0, 1]|[0, 0, 0, 1] |
|[1, 0, 1]   |[1, 1, 2]    |
+------------+-------------+
user_60802760
Related