How to fill null with 0 and accumulate `count` with spark `pivot`?

user_5286990 0
Viewed 208

I have some data of product sales like below:

product date money
A 2020-01 60
B 2020-03 80
A 2020-05 41
B 2020-08 50
B 2020-12 76
A 2020-11 76

And I want to group the data by date and pivot by product My code is below

     df.groupBy("date").pivot("product").agg(
      sum("money").as("month-sum"),
      sum(sum("money")).over(Window.orderBy("date").partitionBy("product")).as("cur-cumulative")
    ).orderBy("date").show()

And result is

|   date|A_month-sum|A_cur-cumulative|B_month-sum|B_cur-cumulative|
+-------+-----------+----------------+-----------+----------------+
|2020-01|         60|              60|       null|            null|
|2020-03|       null|            null|         80|             140|
|2020-05|         41|             181|       null|            null|
|2020-08|       null|            null|         50|             231|
|2020-11|         76|             307|       null|            null|
|2020-12|       null|            null|         76|             383|

My expectation is that null of month-sum can be filled with 0, and null of cur-cumulative can be filled with value of last row, just like this:

|   date|A_month-sum|A_cur-cumulative|B_month-sum|B_cur-cumulative|
+-------+-----------+----------------+-----------+----------------+
|2020-01|         60|              60|          0|               0|
|2020-03|          0|              60|         80|              80|
|2020-05|         41|             101|          0|              80|
|2020-08|          0|             101|         50|             130|
|2020-11|         76|             177|          0|             130|
|2020-12|          0|             177|         76|             206|
+-------+-----------+----------------+-----------+----------------+

Is there any suggestion? Thanks in advance!

apache-spark-sql apache-spark pivot scala sql
1 Answers

You can do a .na.fill(0) before doing the cumulative sums:

import org.apache.spark.sql.expressions.Window

val df2 = df
    .groupBy("date")
    .pivot("product")
    .agg(sum("money"))

val df3 = df2
    .na.fill(0)
    .select(
        col("date") +: 
        df2.columns.tail.flatMap(x => 
            Seq(
                col(x).as(x + "_month-sum"),
                sum(x).over(Window.orderBy("date")).as(x + "_cur-cumulative")
            )
        ): _*
    )
    .orderBy("date")

df3.show
+-------+-----------+----------------+-----------+----------------+
|   date|A_month-sum|A_cur-cumulative|B_month-sum|B_cur-cumulative|
+-------+-----------+----------------+-----------+----------------+
|2020-01|       60.0|            60.0|        0.0|             0.0|
|2020-03|        0.0|            60.0|       80.0|            80.0|
|2020-05|       41.0|           101.0|        0.0|            80.0|
|2020-08|        0.0|           101.0|       50.0|           130.0|
|2020-11|       76.0|           177.0|        0.0|           130.0|
|2020-12|        0.0|           177.0|       76.0|           206.0|
+-------+-----------+----------------+-----------+----------------+
user_141657300
Related