Date-Shift to Roll-up in Pyspark? - apache-spark

Question

I have data that looks similar to this:
What I want to do is replace the Stopdate of the first record with the stopdate of the last record so that I can roll-up all of the records that have a 1 in both gap columns. I know that this is an F.when statement but everything I can think to construct in my head does not give me the result I want. How would I do this while making sure it only applies to records with this id?
Can anyone please help? Thanks!
Sample data in text
ID Startdate Stopdate gap_from_previous_in_days gap_to_next_in_days
1 1/1/2021 1/2/2021
1 1/3/2021 1/4/2021 1 1
1 1/5/2021 1/6/2021 1 1
1 1/7/2021 1/8/2021 1 1
1 1/9/2021 1/10/2021 1 1
1 1/11/2021 1/12/2021 1 1
1 1/13/2021 1/14/2021 1 1
1 1/15/2021 1/16/2021 1 1
1 1/17/2021 1/18/2021 1 1
1 1/19/2021 1/20/2021 1 2
My desired result:
ID Startdate Stopdate gap_from_previous_in_days gap_to_next_in_days
1 1/1/2021 1/20/2021
So basically I'm trying to create a table that instead of looking like this:
ID Startdate Stopdate gap_from_previous_in_days gap_to_next_in_days
1 1/1/2021 1/2/2021
1 1/3/2021 1/4/2021 1 1
1 1/5/2021 1/6/2021 1 1
1 1/7/2021 1/8/2021 1 1
1 1/9/2021 1/10/2021 1 1
1 1/11/2021 1/12/2021 1 1
1 1/13/2021 1/14/2021 1 1
1 1/15/2021 1/16/2021 1 1
1 1/17/2021 1/18/2021 1 1
1 1/19/2021 1/20/2021 1 3
1 1/23/2021 1/25/2021 3
Would look like this
ID Startdate Stopdate gap_from_previous_in_days gap_to_next_in_days
1 1/1/2021 1/2/2021 3
1 1/23/2021 1/25/2021 3
Hopefully that helps illustrate what I'm trying to do.
I'm basically trying to combine records that are only one day apart from each other.

Answer 1

The idea is to create a grouping column based on the count of rows with gap_from_previous_in_days != 1, group by that grouping column and the ID, get the earliest start date and latest stop date, and their associated values of the gap:
from pyspark.sql import functions as F, Window
result = df.withColumn(
'Startdate2', F.to_date('Startdate', 'M/d/yyyy')
).withColumn(
'Stopdate2', F.to_date('Stopdate', 'M/d/yyyy')
).withColumn(
'grp',
F.count(
F.when(F.col('gap_from_previous_in_days') != 1, 1)
).over(
Window.partitionBy('ID').orderBy('Startdate2')
)
).groupBy('ID', 'grp').agg(
F.min(F.struct('Startdate2', 'Startdate', 'gap_from_previous_in_days')).alias('start'),
F.max(F.struct('Stopdate2', 'Stopdate', 'gap_to_next_in_days')).alias('end')
).select(
'ID',
'start.Startdate', 'end.Stopdate',
'start.gap_from_previous_in_days', 'end.gap_to_next_in_days'
)
result.show()
+---+---------+---------+-------------------------+-------------------+
| ID|Startdate| Stopdate|gap_from_previous_in_days|gap_to_next_in_days|
+---+---------+---------+-------------------------+-------------------+
| 1| 1/1/2021|1/20/2021| null| 3|
| 1|1/23/2021|1/25/2021| 3| null|
+---+---------+---------+-------------------------+-------------------+

Answer 2

Another approach is to construct all dates from the given intervals.
Starting from this list of dates new intervals can be recalculated.
For this approach the gap columns are not used/needed. Therefore, at the end they are recalculated but if not needed this can be omitted.
# Create test dataframe
import pyspark.sql.functions as F
df = (spark.createDataFrame([[1, '2021-01-01', '2021-01-10'],
[1, '2021-01-11', '2021-01-12'],
[1, '2021-01-14', '2021-01-16'],
[1, '2021-01-17', '2021-01-20'],
[2, '2021-01-01', '2021-01-10'],
[2, '2021-01-12', '2021-01-14'],
[2, '2021-01-14', '2021-01-15'],
[2, '2021-01-19', '2021-01-20'],
], schema="ID int, From string, To string")
.selectExpr('ID',
'to_date(From, "yyyy-MM-dd") as StartDate',
'to_date(To, "yyyy-MM-dd") as StopDate')
)
# Do actual calculation
df_result = (df
# Get all included dates
.selectExpr("ID", "explode(sequence(StartDate, StopDate)) as dates")
# Get previous and next date
.withColumn("Previous", F.expr('LAG(dates) OVER (PARTITION BY ID ORDER BY dates ASC)'))
.withColumn("Next", F.expr('LAG(dates) OVER (PARTITION BY ID ORDER BY dates DESC)'))
# Flag beginnings and endings of intervals
.withColumn("Begin", F.expr("datediff(dates, previous)> 1 OR previous is NULL"))
.withColumn("End", F.expr("datediff( next, dates)> 1 OR next is NULL"))
# Only keep beginnings and endings
.filter("Begin OR End")
# Get end next to begin and only keep beginnings
.withColumn("IntervalEnd", F.expr('LAG(dates) OVER (PARTITION BY ID ORDER BY dates DESC)'))
.filter("Begin")
# Rename columns + calculate gaps
.selectExpr(
"ID",
"dates as StartDate",
"IntervalEnd as StopDate",
"datediff(dates, LAG(IntervalEnd) OVER (PARTITION BY ID ORDER BY dates ASC)) as gap_from_previous_in_days",
"datediff(LAG(dates) OVER (PARTITION BY ID ORDER BY dates DESC), IntervalEnd ) as gap_to_next_in_days"
)
)
df_result.show()
+---+----------+----------+-------------------------+-------------------+
| ID| StartDate| StopDate|gap_from_previous_in_days|gap_to_next_in_days|
+---+----------+----------+-------------------------+-------------------+
| 1|2021-01-01|2021-01-12| null| 2|
| 1|2021-01-14|2021-01-20| 2| null|
| 2|2021-01-01|2021-01-10| null| 2|
| 2|2021-01-12|2021-01-15| 2| 4|
| 2|2021-01-19|2021-01-20| 4| null|
+---+----------+----------+-------------------------+-------------------+