Spark 2.3: Remove rows if between 3 minutes of two consecutive occurrence

Viewed 77

Suppose we have a dataframe

val df = Seq(
  ("id1", "2020-08-02 16:42:00", "2020-08-02 16:43:00", "event1", 1),
  ("id1", "2020-08-02 16:43:00", "2020-08-02 16:44:00", "event2", 0),
  ("id1", "2020-08-02 16:44:00", "2020-08-02 16:45:00", "event1", 0),
  ("id1", "2020-08-02 16:45:00", "2020-08-02 16:47:00", "event3", 1),
  ("id1", "2020-08-02 16:47:00", "2020-08-02 16:51:00", "event4", 0),
  ("id1", "2020-08-02 16:51:00", "2020-08-02 16:52:00", "event3", 0))
.toDF("id", "start_time", "end_time", "event_id", "indicator")

df.show()

+---+-------------------+-------------------+--------+---------+
| id|         start_time|           end_time|event_id|indicator|
+---+-------------------+-------------------+--------+---------+
|id1|2020-08-02 16:42:00|2020-08-02 16:43:00|  event1|        1|
|id1|2020-08-02 16:43:00|2020-08-02 16:44:00|  event2|        0|
|id1|2020-08-02 16:44:00|2020-08-02 16:45:00|  event1|        0|
|id1|2020-08-02 16:45:00|2020-08-02 16:47:00|  event3|        1|
|id1|2020-08-02 16:47:00|2020-08-02 16:51:00|  event4|        0|
|id1|2020-08-02 16:51:00|2020-08-02 16:52:00|  event3|        0|
+---+-------------------+-------------------+--------+---------+

For every indicator = 1, I want to see if the start_time when the same event_id happens again is within 3 minutes of the end_time. And if it is, then remove rows that are between the two time range.

For example, event1 in row 1 has end_time = 2020-08-02 16:43:0, and the next event1 happens in row 3 at start_time = 2020-08-02 16:44:00. Since they happen within 3 min, we remove the second row and group event1 into one row.

Final output should look like

+---+-------------------+-------------------+--------+---------+
| id|         start_time|           end_time|event_id|indicator|
+---+-------------------+-------------------+--------+---------+
|id1|2020-08-02 16:42:00|2020-08-02 16:45:00|  event1|        1|
|id1|2020-08-02 16:45:00|2020-08-02 16:47:00|  event3|        1|
|id1|2020-08-02 16:47:00|2020-08-02 16:51:00|  event4|        0|
|id1|2020-08-02 16:51:00|2020-08-02 16:52:00|  event3|        0|
+---+-------------------+-------------------+--------+---------+

I started off by finding the next occurring time when the event happens and finding the time difference in sec

val timeThreshold = 3*60 // in seconds
val window = Window.partitionBy("id", "event_id").orderBy("start_time")
  .rowsBetween(1, Window.unboundedFollowing)

val updatedDF = df
  .withColumn("end_time_long", to_timestamp(col("end_time")).cast(LongType))
  .withColumn("next_time",
    when(col("indicator") === 1, first("start_time").over(window)).otherwise("null"))
  .withColumn("next_time_long", to_timestamp(col("next_time")).cast(LongType))
  .withColumn("time_difference", col("next_time_long") - col("end_time_long"))

testDFUpdate.drop("next_time_long", "end_time_long").orderBy("start_time").show()

+---+-------------------+------------------+-------+---------+-------------------+-------------------+
| id|         start_time|          end_time|item_id|indicator|          next_time|time_difference_sec|
+---+-------------------+------------------+-------+---------+-------------------+-------------------+
|id1|2020-08-02 16:42:00|2020-08-02 16:43:0|  item1|        1|2020-08-02 16:44:00|                 60|
|id1|2020-08-02 16:43:00|2020-08-02 16:44:0|  item2|        0|               null|               null|
|id1|2020-08-02 16:44:00|2020-08-02 16:45:0|  item1|        0|               null|               null|
|id1|2020-08-02 16:45:00|2020-08-02 16:47:0|  item3|        1|2020-08-02 16:51:00|                240|
|id1|2020-08-02 16:47:00|2020-08-02 16:51:0|  item4|        0|               null|               null|
|id1|2020-08-02 16:51:00|2020-08-02 16:52:0|  item3|        0|               null|               null|
+---+-------------------+------------------+-------+---------+-------------------+-------------------+

But I cannot think of a way to aggregate/join and create a new desired end_time column with event2 removed.

1 Answers

You can use lag function also to implement this logic. Bring in all the previous start_time and end_time to the same row using lag function as below.

scala> someDF.show
+---+-------------------+-------------------+--------+---------+
| id|         start_time|           end_time|event_id|indicator|
+---+-------------------+-------------------+--------+---------+
|id1|2020-08-02 16:42:00|2020-08-02 16:43:00|  event1|        1|
|id1|2020-08-02 16:43:00|2020-08-02 16:44:00|  event2|        0|
|id1|2020-08-02 16:44:00|2020-08-02 16:45:00|  event1|        0|
|id1|2020-08-02 16:45:00|2020-08-02 16:47:00|  event3|        1|
|id1|2020-08-02 16:47:00|2020-08-02 16:51:00|  event4|        0|
|id1|2020-08-02 16:51:00|2020-08-02 16:52:00|  event3|        0|
+---+-------------------+-------------------+--------+---------+

scala> someDF.createOrReplaceTempView("somedf")

scala> val a = spark.sql("""select id,cast(start_time as timestamp),cast(end_time as timestamp),
cast(lag(start_time,1) over(partition by event_id order by start_time) as timestamp) as prev_start_time,
cast(lag(end_time,1) over(partition by event_id order by end_time) as timestamp) as prev_end_time,event_id,indicator from somedf""")
a: org.apache.spark.sql.DataFrame = [id: string, start_time: timestamp ... 5 more fields]

scala> a.show
+---+-------------------+-------------------+-------------------+-------------------+--------+---------+
| id|         start_time|           end_time|    prev_start_time|      prev_end_time|event_id|indicator|
+---+-------------------+-------------------+-------------------+-------------------+--------+---------+
|id1|2020-08-02 16:45:00|2020-08-02 16:47:00|               null|               null|  event3|        1|
|id1|2020-08-02 16:51:00|2020-08-02 16:52:00|2020-08-02 16:45:00|2020-08-02 16:47:00|  event3|        0|
|id1|2020-08-02 16:43:00|2020-08-02 16:44:00|               null|               null|  event2|        0|
|id1|2020-08-02 16:42:00|2020-08-02 16:43:00|               null|               null|  event1|        1|
|id1|2020-08-02 16:44:00|2020-08-02 16:45:00|2020-08-02 16:42:00|2020-08-02 16:43:00|  event1|        0|
|id1|2020-08-02 16:47:00|2020-08-02 16:51:00|               null|               null|  event4|        0|
+---+-------------------+-------------------+-------------------+-------------------+--------+---------+

Now find the time difference between the start_time and prev_end_time and flag those records where the time difference is less than 3 mins

scala> a.withColumn("time_diff", unix_timestamp($"start_time") - unix_timestamp($"prev_end_time"))
.withColumn("flag", when(col("time_diff") < 60*3,lit("FALSE")).otherwise("TRUE")).orderBy("event_id", "start_time").show
+---+-------------------+-------------------+-------------------+-------------------+--------+---------+---------+-----+
| id|         start_time|           end_time|    prev_start_time|      prev_end_time|event_id|indicator|time_diff| flag|
+---+-------------------+-------------------+-------------------+-------------------+--------+---------+---------+-----+
|id1|2020-08-02 16:42:00|2020-08-02 16:43:00|               null|               null|  event1|        1|     null| TRUE|
|id1|2020-08-02 16:44:00|2020-08-02 16:45:00|2020-08-02 16:42:00|2020-08-02 16:43:00|  event1|        0|       60|FALSE|
|id1|2020-08-02 16:43:00|2020-08-02 16:44:00|               null|               null|  event2|        0|     null| TRUE|
|id1|2020-08-02 16:45:00|2020-08-02 16:47:00|               null|               null|  event3|        1|     null| TRUE|
|id1|2020-08-02 16:51:00|2020-08-02 16:52:00|2020-08-02 16:45:00|2020-08-02 16:47:00|  event3|        0|      240| TRUE|
|id1|2020-08-02 16:47:00|2020-08-02 16:51:00|               null|               null|  event4|        0|     null| TRUE|
+---+-------------------+-------------------+-------------------+-------------------+--------+---------+---------+-----+

Now Filter out only the records matching flag = TRUE

Related