Fix date overlap in pyspark - pyspark

Question

I have a dataset like this below:
+------+------+-------+------------+------------+-----------------+
| key1 | key2 | price | date_start | date_end | sequence_number |
+------+------+-------+------------+------------+-----------------+
| a | b | 10 | 2022-01-03 | 2022-01-05 | 1 |
| a | b | 10 | 2022-01-02 | 2050-05-15 | 2 |
| a | b | 10 | 2022-02-02 | 2022-05-10 | 3 |
| a | b | 20 | 2024-02-01 | 2050-10-10 | 4 |
| a | b | 20 | 2024-04-01 | 2025-09-10 | 5 |
| a | b | 10 | 2022-04-02 | 2024-09-10 | 6 |
| a | b | 20 | 2024-09-11 | 2050-10-10 | 7 |
+------+------+-------+------------+------------+-----------------+
What I want to achieve is this:
+------+------+-------+------------+------------+
| key1 | key2 | price | date_start | date_end |
+------+------+-------+------------+------------+
| a | b | 10 | 2022-01-02 | 2024-09-10 |
| a | b | 20 | 2024-09-11 | 2050-10-10 |
+------+------+-------+------------+------------+
The sequence number is the order in which the data rows were received.
The resultant dataset should be able to fix the overlapping dates for each price but also consider the fact that when there is a new price for the same key columns the older record's date_end is updated to date_start-1
After the first three sequence numbers, the output looked like this:
+------+------+-------+------------+------------+
| key1 | key2 | price | date_start | date_end |
+------+------+-------+------------+------------+
| a | b | 10 | 2022-01-02 | 2050-05-15 |
+------+------+-------+------------+------------+
This covers the max range for the price.
After the 4th sequence number:
+------+------+-------+------------+------------+
| key1 | key2 | price | date_start | date_end |
+------+------+-------+------------+------------+
| a | b | 10 | 2022-01-02 | 2024-01-31 |
| a | b | 20 | 2024-02-01 | 2050-10-10 |
+------+------+-------+------------+------------+
After the 5th sequence number:
+------+------+-------+------------+------------+
| key1 | key2 | price | date_start | date_end |
+------+------+-------+------------+------------+
| a | b | 10 | 2022-01-02 | 2024-01-31 |
| a | b | 20 | 2024-02-01 | 2050-10-10 |
+------+------+-------+------------+------------+
No changes as the date overlaps
After the 6th sequence number:
+------+------+-------+------------+------------+
| key1 | key2 | price | date_start | date_end |
+------+------+-------+------------+------------+
| a | b | 10 | 2022-01-02 | 2024-09-10 |
| a | b | 20 | 2024-09-11 | 2050-10-10 |
+------+------+-------+------------+------------+
The date_start and date_end both are updated
After the 7th sequence number:
+------+------+-------+------------+------------+
| key1 | key2 | price | date_start | date_end |
+------+------+-------+------------+------------+
| a | b | 10 | 2022-01-02 | 2024-09-10 |
| a | b | 20 | 2024-09-11 | 2050-10-10 |
+------+------+-------+------------+------------+
No changes.

Answer 1

So here's an answer. First I expand the dates into rows. Then I use a group by with struct & collect_list, to capture the price and sequence_number together. I take the first item of the array returned (and sorted) from collect list to effectively give me the max of the sequence number to then pull back the price. From there a group the dates by min/max.
from pyspark.sql.functions import min, explode, col, expr, max
data = [
("a" ,"b" , 10 ,"2022-01-03" ,"2022-01-05" , 1 ),
("a" ,"b" , 10 ,"2022-01-02" ,"2050-05-15" , 2 ),
("a" ,"b" , 10 ,"2022-02-02" ,"2022-05-10" , 3 ),
("a" ,"b" , 20 ,"2024-02-01" ,"2050-10-10" , 4 ),
("a" ,"b" , 20 ,"2024-04-01" ,"2025-09-10" , 5 ),
("a" ,"b" , 10 ,"2022-04-02" ,"2024-09-10" , 6 ),
("a" ,"b" , 20 ,"2024-09-11" ,"2050-10-10" , 7 ),
]
columns = ["key1","key2","price","date_start","date_end","sequence_number"]
df = spark.createDataFrame(data).toDF(*columns)
#create all the dates as rows
df_sequence = df.select( "*" , explode( expr("sequence ( to_date(date_start),to_date(date_end), interval 1 day)").alias('day') ) )
df_date_range = df_sequence.groupby( \
col("key1"),
col("key2"),
col("col")
).agg( \# here's the magic
reverse( \#descending sort effectively on sequence number.
array_sort( \#ascending sort of array by first element in struct then second element in struct
collect_list( \# collect all elements of the group
struct( \# get all data we need
col("sequence_number").alias("sequence") ,
col("price" ).alias("price")
)
)
)
)[0].alias("mystruct")\ # pull the first element to get "max"
).select( \
col("key1"),
col("key2"),
col("col"),
col("mystruct.price") \#pull out price from struct
)
#without comments
#df_date_range = df_sequence.groupby( col("key1"),col("key2"), col("col") ).agg( reverse(array_sort(collect_list( struct(col("sequence_number").alias("sequence") ,col("price" ).alias("price") ) )))[0].alias("mystruct") ).select( col("key1"), col("key2"), col("col"), col("mystruct.price") )
df_date_range.groupby( col("key1"),col("key2"), col("price") ).agg( min("col").alias("date_start"), max("col").alias("date_end") ).show()
+----+----+-----+----------+----------+
|key1|key2|price| date_start | date_end|
+----+----+-----+----------+----------+
| a| b| 10|2022-01-02|2024-09-10|
| a| b| 20|2024-09-11|2050-10-10|
+----+----+-----+----------+----------+
This does assume you will only ever use the price once before changing it. If you needed to go to back to a price you'd have to use window logic to identify the price different prices ranges and add that to your collect_list as an extra factor to sort on.