Compute overlapping time intervals across dataframe

user_1612432 0
Viewed 202

I am trying to compute which rows overlaps in time with any other row on a DataFrame. Given this example:

            df = spark.createDataFrame(
            [
                [1,"2020:01:01 12:00", "2020:01:01 13:00"], # T
                [1,"2020:01:01 12:30", "2020:01:01 13:00"], # T
                [1,"2020:01:01 14:00", "2020:01:01 15:00"], # F
                [2,"2020:01:01 09:00", "2020:01:01 13:00"], # F
                [2,"2020:01:01 18:00", "2020:01:01 19:00"], # F
            ],
            ["id", "start", "end"]
        )
        w = Window.partitionBy("id").orderBy("start").rowsBetween(Window.unboundedPreceding, 0)
        df.withColumn("OVERLAPS", ?????.over(w)).show()

I would like to generate a column that signal whenever two rows with the same id overlaps in time:

                [1,"2020:01:01 12:00", "2020:01:01 13:00", True], 
                [1,"2020:01:01 12:30", "2020:01:01 13:00", True], 
                [1,"2020:01:01 14:00", "2020:01:01 15:00", False], 
                [2,"2020:01:01 09:00", "2020:01:01 13:00", False], 
                [2,"2020:01:01 18:00", "2020:01:01 19:00", False],

Ideally I could create time intervals, in pandas it would be using pd.IntervalIndex and I would have intervals which I could query, for example "2020:01:01 13:00" in [2020:01:01 12:30, 2020:01:01 14:00)

But I've been struggling to do it on spark.

pyspark python
1 Answers

So I gave this a go, in scala, though, and my solution is ugly and could probably improved upon (the whole column renaming and such especially). I made one assumption, which is that the times are already parsed into some integer format, then I basically solved this with a little more complicated join:

val s = Seq(
    (1, 1200, 1300),
    (1, 1230, 1300),
    (1, 1400, 1500),
    (2, 900, 1300),
    (2, 1800, 1900)
)

val df = s.toDF("id", "start", "end").withColumn("i", monotonicallyIncreasingId)
val dfr = df.select($"id", $"start".alias("start_r"), $"end".alias("end_r"), $"i".alias("i_r")) 

val overlaps = df.join(
    dfr,
    df("id") === dfr("id") 
        and df("start") <= dfr("end_r") 
        and df("end") >= dfr("start_r") 
        and (df("i") !== dfr("i_r"))
    ).select($"i".alias("i_overlaps"), lit(true).alias("overlaps"))

val result = df.join(overlaps, df("i") === overlaps("i_overlaps"), "left").
        drop("i_overlaps").
        withColumn("overlaps", when($"overlaps".isNull, false).otherwise($"overlaps"))

result.show

result:

+---+-----+----+---+--------+                                                   
| id|start| end|  i|overlaps|
+---+-----+----+---+--------+
|  1| 1200|1300|  0|    true|
|  1| 1230|1300|  1|    true|
|  2|  900|1300|  3|   false|
|  1| 1400|1500|  2|   false|
|  2| 1800|1900|  4|   false|
+---+-----+----+---+--------+

You can ofc drop the i column again. In this solution it is also not needed to repartition or sort the dataframe at all as this will be handled within the join. Probably, this type of join is not the most efficient, though :)

user_5293210
Related