Ufdjrw

I'm trying to extract combined data intervals based on a time series in scala and spark

I have the following data in a dataframe:

Id | State | StartTime | EndTime
---+-------+---------------------+--------------------
 1 | R | 2019-01-01T03:00:00 | 2019-01-01T11:30:00
 1 | R | 2019-01-01T11:30:00 | 2019-01-01T15:00:00
 1 | R | 2019-01-01T15:00:00 | 2019-01-01T22:00:00
 1 | W | 2019-01-01T22:00:00 | 2019-01-02T04:30:00
 1 | W | 2019-01-02T04:30:00 | 2019-01-02T13:45:00
 1 | R | 2019-01-02T13:45:00 | 2019-01-02T18:30:00
 1 | R | 2019-01-02T18:30:00 | 2019-01-02T22:45:00

I need to extract the data into time intervals based on the id and state. The resulting data needs to look like:

Id | State | StartTime | EndTime
---+-------+---------------------+--------------------
 1 | R | 2019-01-01T03:00:00 | 2019-01-01T22:00:00
 1 | W | 2019-01-01T22:00:00 | 2019-01-02T13:45:00
 1 | R | 2019-01-02T13:45:00 | 2019-01-02T22:45:00

Note that the first three records have been grouped together because the equipment is contiguously in an R state from 2019-01-01T03:00:00 to 2019-01-01T22:00:00, then it switches to a W state for the next two records from 2019-01-01T22:00:00 to 2019-01-02T13:45:00and then back to an R state for the last two records.

Since I had a similar case recently I would like to provide the complete solution for this one. The part of the code:

val df2 = df
 .select('Id,'State,'StartTime,'EndTime,
 row_number().over(idSpec).as("idRowNumber"),
 row_number().over(idStateSpec).as("idStateRowNumber"))

Has output:

+---+-----+-------------------+-------------------+-----------+----------------+
| Id|State| StartTime| EndTime|idRowNumber|idStateRowNumber|
+---+-----+-------------------+-------------------+-----------+----------------+
| 1| R|2019-01-01 03:00:00|2019-01-01 11:30:00| 1| 1|
| 1| R|2019-01-01 11:30:00|2019-01-01 15:00:00| 2| 2|
| 1| R|2019-01-01 15:00:00|2019-01-01 22:00:00| 3| 3|
| 1| W|2019-01-01 22:00:00|2019-01-02 04:30:00| 4| 1|
| 1| W|2019-01-02 04:30:00|2019-01-02 13:45:00| 5| 2|
| 1| R|2019-01-02 13:45:00|2019-01-02 18:30:00| 6| 4|
| 1| R|2019-01-02 18:30:00|2019-01-02 22:45:00| 7| 5|
+---+-----+-------------------+-------------------+-----------+----------------+

Notice that the difference between idRowNumber and idStateRowNumber will be identical for each combination of (Id, State) hence we can create a new column called category and group by this one in order to get min StartTime and max EndTime for each group. The complete code should look like the next one:

val idSpec = Window.partitionBy('Id).orderBy('StartTime)
val idStateSpec = Window.partitionBy('Id,'State).orderBy('StartTime)

val df2 = df
 .select('Id,'State,'StartTime.cast("timestamp"),'EndTime.cast("timestamp"),
 row_number().over(idSpec).as("idRowNumber"),
 row_number().over(idStateSpec).as("idStateRowNumber"))
 .withColumn("Category", $"idRowNumber" - $"idStateRowNumber")
 .groupBy("Category", "Id", "State")
 .agg(min("StartTime").as("StartTime"), max("EndTime").as("EndTime"))
 .drop("Category")

And the output:

+---+-----+-------------------+-------------------+
| Id|State| StartTime| EndTime|
+---+-----+-------------------+-------------------+
| 1| R|2019-01-01 03:00:00|2019-01-01 22:00:00|
| 1| W|2019-01-01 22:00:00|2019-01-02 13:45:00|
| 1| R|2019-01-02 13:45:00|2019-01-02 22:45:00|
+---+-----+-------------------+-------------------+

So it turns out that the answer to this is Combine rows when the end time of one is the start time of another (Oracle) translated to Spark.

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.col,row_number
import spark.implicits._

val idSpec = Window.partitionBy('Id).orderBy('StartTime)
val idStateSpec = Window.partitionBy('Id,'State).orderBy('StartTime)
val df2 = df
 .select('Id,'State,'StartTime,'EndTime,
 row_number().over(idSpec).as("idRowNumber"),
 row_number().over(idStateSpec).as("idStateRowNumber"))
 .groupBy('Id,'State,'idRowNumber - 'idStateRowNumber)
 .agg(min('StartTime).as("StartTime"), max('EndTime).as("EndTime"))