Spark SQL running sum with over() function - summing an indicator - sql

Question

I have the following data:
df = [{"Category": 'A', "date": '01/01/2022', "Indictor": 1},
{"Category": 'A', "date": '02/01/2022', "Indictor": 0},
{"Category": 'A', "date": '03/01/2022', "Indictor": 1},
{"Category": 'A', "date": '04/01/2022', "Indictor": 1},
{"Category": 'A', "date": '05/01/2022', "Indictor": 1},
{"Category": 'B', "date": '01/01/2022', "Indictor": 0},
{"Category": 'B', "date": '02/01/2022', "Indictor": 1},
{"Category": 'B', "date": '03/01/2022', "Indictor": 1},
{"Category": 'B', "date": '04/01/2022', "Indictor": 0},
{"Category": 'B', "date": '05/01/2022', "Indictor": 0},
{"Category": 'B', "date": '06/01/2022', "Indictor": 1}]
df = spark.createDataFrame(df)
I want to use a LEAD() function to group by 'Category' and then order by 'date' ascending. Then create a new field called 'consec_ind' which is a counter of the number of consecutive days that the indicator has been 1.
This is the code I have tried but it doesn't quite work.
df.createOrReplaceTempView('df')
%sql
select date, Indictor,
case when Indictor > 0 THEN
(sum(count(Indictor)) over (order by date)) else 0 end as running_total
from df
WHERE Category = 'A'
group by date, Indictor
order by date, Indictor;
This is what I would like the data to look like:
{"Category": 'A', "date": '02/01/2022', "Indictor": 0,"consec_ind":0},
{"Category": 'A', "date": '03/01/2022', "Indictor": 1,"consec_ind":1},
{"Category": 'A', "date": '04/01/2022', "Indictor": 1,"consec_ind":2},
{"Category": 'A', "date": '05/01/2022', "Indictor": 1,"consec_ind":3},
{"Category": 'B', "date": '01/01/2022', "Indictor": 0,"consec_ind":0},
{"Category": 'B', "date": '02/01/2022', "Indictor": 1,"consec_ind":1},
{"Category": 'B', "date": '03/01/2022', "Indictor": 1,"consec_ind":2},
{"Category": 'B', "date": '04/01/2022', "Indictor": 0,"consec_ind":0},
{"Category": 'B', "date": '05/01/2022', "Indictor": 0,"consec_ind":0},
{"Category": 'B', "date": '06/01/2022', "Indictor": 1,"consec_ind":1}]

Answer 1

Here is my solution
First step: Partition by Category
Second step: Try to partition by your custom condition within first partitioning - i am doing it by incrementing counter every time i meet 0 while iterating over first window
Third step: Calcuate sum within final partitions
Here is version in PySpark:
import pyspark.sql.functions as F
from pyspark.sql import Window
df = [
{"Category": "A", "date": "01/01/2022", "Indictor": 1},
{"Category": "A", "date": "02/01/2022", "Indictor": 0},
{"Category": "A", "date": "03/01/2022", "Indictor": 1},
{"Category": "A", "date": "04/01/2022", "Indictor": 1},
{"Category": "A", "date": "05/01/2022", "Indictor": 1},
{"Category": "B", "date": "01/01/2022", "Indictor": 0},
{"Category": "B", "date": "02/01/2022", "Indictor": 1},
{"Category": "B", "date": "03/01/2022", "Indictor": 1},
{"Category": "B", "date": "04/01/2022", "Indictor": 0},
{"Category": "B", "date": "05/01/2022", "Indictor": 0},
{"Category": "B", "date": "06/01/2022", "Indictor": 1},
]
df = spark.createDataFrame(df)
windowSpec = Window.partitionBy("Category").orderBy("date")
df.withColumn(
"partition_number", F.sum((F.col("Indictor") == 0).cast("int")).over(windowSpec)
).withColumn(
"part_sum",
F.sum(F.col("Indictor")).over(
Window.partitionBy("Category", "partition_number").orderBy("date")
),
).drop(
"partition_number"
).show()
and here is sql:
select
Category,
Indictor,
date,
sum(Indictor) over (
PARTITION BY Category,
partition_number
ORDER BY
date
) as part_sum
from
(
select
*,
sum(
case when Indictor == 0 then 1 else 0 end
) over (
PARTITION BY Category
ORDER BY
date
) as partition_number
from
df
)
Output is:
+--------+--------+----------+--------+
|Category|Indictor| date|part_sum|
+--------+--------+----------+--------+
| A| 1|01/01/2022| 1|
| A| 0|02/01/2022| 0|
| A| 1|03/01/2022| 1|
| A| 1|04/01/2022| 2|
| A| 1|05/01/2022| 3|
| B| 0|01/01/2022| 0|
| B| 1|02/01/2022| 1|
| B| 1|03/01/2022| 2|
| B| 0|04/01/2022| 0|
| B| 0|05/01/2022| 0|
| B| 1|06/01/2022| 1|
+--------+--------+----------+--------+