I would like help with collapsing time intervals that overlap each other within each group.
Specifically, this is what I have:
| id | time_start | time_end |
|---|---|---|
| 1 | 8:00 | 9:00 |
| 1 | 8:30 | 9:30 |
| 1 | 9:45 | 10:00 |
| 2 | 8:00 | 9:00 |
| 2 | 8:30 | 8:40 |
This is what I desire:
| id | time_start | time_end |
|---|---|---|
| 1 | 8:00 | 9:30 |
| 1 | 9:45 | 10:00 |
| 2 | 8:00 | 9:00 |
The data is very large and needs to be handled as a Spark dataframe.
Appreciate any help!
You can add a grouping column as below:
from pyspark.sql import functions as F, Window
df2 = df.withColumn(
'time_start', F.lpad('time_start', 5, '0')
).withColumn(
'time_end', F.lpad('time_end', 5, '0')
).withColumn(
'overlap',
F.when(
F.max('time_end').over(
Window.partitionBy('id')
.orderBy('time_start')
.rowsBetween(Window.unboundedPreceding, -1)
) >= F.col('time_start'),
0
).otherwise(1)
).withColumn(
'group',
F.sum('overlap').over(Window.partitionBy('id').orderBy('time_start'))
).groupBy('id', 'group').agg(
F.min('time_start').alias('time_start'),
F.max('time_end').alias('time_end')
).drop('group')
df2.show()
+---+----------+--------+
| id|time_start|time_end|
+---+----------+--------+
| 1| 08:00| 09:30|
| 1| 09:45| 10:00|
| 2| 08:00| 09:00|
+---+----------+--------+
Behind the scenes before grouping:
+---+----------+--------+-------+-----+
| id|time_start|time_end|overlap|group|
+---+----------+--------+-------+-----+
| 1| 08:00| 09:00| 1| 1|
| 1| 08:30| 09:30| 0| 1|
| 1| 09:45| 10:00| 1| 2|
| 2| 08:00| 09:00| 1| 1|
| 2| 08:30| 08:40| 0| 1|
+---+----------+--------+-------+-----+