Merge overlapping date ranges in pyspark

user_9325760 0
Viewed 532

I have a dataframe with the schema (ID, start_date, end_date)

Sample dataframe:

id start_date end_date
1 03/05/2020 07/05/2020
1 05/05/2020 07/05/2020
1 01/05/2020 07/05/2020
1 01/05/2020 13/05/2020
2 06/05/2020 09/05/2020
1 18/05/2020 20/05/2020
2 09/05/2020 17/05/2020

I want the result in below format

id start_dates days
1 [03/05/2020, 05/05/2020, 01/05/2020, 01/05/2020, 18/05/2020] 16
2 [06/05/2020, 09/05/2020] 12

days is the sum of non overlapping date intervals (start_date - end_date).

The approach that i tried is to group by id, collect start_date and end_date as lists and then write a UDF to get the days column. I would like to know if there is a better way to achieve this.

pyspark apache-spark python
1 Answers

Here is a possible solution which generates sequence and finds the array_overlaps with the help of the lag function. Then separates each group within the id column and sums the values. Then sums again for the final result:

from pyspark.sql import functions as F, Window

w = Window.partitionBy("id").orderBy(F.monotonically_increasing_id())

out = (df.withColumn("start_date",F.to_date("start_date","dd/MM/yyyy"))
       .withColumn("end_date",F.to_date("end_date","dd/MM/yyyy"))
       .withColumn("Seq",F.sequence("start_date","end_date"))
       .withColumn("Lag",F.arrays_overlap("Seq",F.lag("Seq").over(w)))
       .withColumn("Flag",F.when(F.col("Lag")|F.col("Lag").isNull(),1).otherwise(0))
        
        .groupBy("Id","Flag").agg(F.collect_list("start_date").alias("start_date"),
                (F.datediff(F.max("end_date"),F.min("start_date"))+1).alias("Days"))
        .groupBy("id").agg(F.flatten(F.collect_list("start_date")).alias("start_date"),
                           F.sum("Days").alias("Days"))
       )

out.show(truncate=False)

+---+------------------------------------------------------------+----+
|id |start_date                                                  |Days|
+---+------------------------------------------------------------+----+
|1  |[2020-05-03, 2020-05-05, 2020-05-01, 2020-05-01, 2020-05-18]|16  |
|2  |[2020-05-06, 2020-05-09]                                    |12  |
+---+------------------------------------------------------------+----+
user_98406370
Related