Pyspark: finding blocks of timestamp data without pause - pyspark

Question

I am dealing with a problem in PySpark.
Input:
+-------+---------------------+
|user_id|activity_timestamp |
+-------+---------------------+
| 1| 2021/06/01 19:00 |
| 1| 2021/06/01 19:01 |
| 2| 2021/06/01 19:01 |
| 1| 2021/06/01 19:02 |
| 2| 2021/06/01 19:02 |
| 1| 2021/06/01 19:10 |
| 1| 2021/06/01 19:11 |
+-------+---------------------+
Desired output:
For each user, detect periods of continuous activity timestamp (gaps between timestamps are smaller than for ex. 5 minutes)
+-------+---------------------------+---------------------------+
|user_id| activity _start | activity _stop |
+-------+---------------------------+---------------------------+
| 1| 2021/06/01 19:00 | 2021/06/01 19:02 |
| 1| 2021/06/01 19:10 | 2021/06/01 19:11 |
| 2| 2021/06/01 19:01 | 2021/06/01 19:02 |
Progress so far: I have used a Window function to find the time of the previous activity and from that I have calculated the time that passed since previous activity. But I am struggling to create the desired output.
time_window = Window.partitionBy("user_id").orderBy("user_id", "activity_timestamp")
df = (df
.withColumn('prev_time', F.lag(F.col('activity_timestamp')).over(time_window))
.withColumn('time_gap', F.col('activity_timestamp').cast("long") - F.col('prev_time').cast("long"))
)

Answer 1

Starting with the time difference between two following rows within a window works, just like time_gap in the question.
The next steps are:
a new column is created with a value that depends on this difference: if the difference is smaller than 5 minutes, the new column gets the value 0 and 1 otherwise
sum the new column over the window. All rows that belong the to same period will get the same number: only the first row in the period contains a 1 and all other rows contain a 0. In the code below, this column is called id.
group the dataframe by user_id and id and take the minimum and maximum of activity_timestamp as start and end of the period
from pyspark.sql import functions as F
df.selectExpr("*", """
sum(
case
when lag(activity_timestamp) over (PARTITION BY user_id ORDER BY activity_timestamp) is null then 1
when cast(activity_timestamp as long)-cast(lag(activity_timestamp) over (PARTITION BY user_id ORDER BY activity_timestamp) as long) > 60 * 5 then 1
else 0
end
) over (PARTITION BY user_id ORDER BY activity_timestamp) as id
""") \
.groupBy("user_id", "id") \
.agg(F.min("activity_timestamp").alias("activity_start"),
F.max("activity_timestamp").alias("activity_end")) \
.drop("id") \
.show()
Output:
+-------+-------------------+-------------------+
|user_id| activity_start| activity_end|
+-------+-------------------+-------------------+
| 1|2021-06-01 19:00:00|2021-06-01 19:02:00|
| 1|2021-06-01 19:10:00|2021-06-01 19:11:00|
| 2|2021-06-01 19:01:00|2021-06-01 19:02:00|
+-------+-------------------+-------------------+

Answer 2

First convert the time in minutes, then get the adjacent time in the same row. From here we can calculate the time difference and then do a convert add a column which has 0 if the difference between times is less than 5 mins and 1 if it is larger. Just summing over the this new column reusing the window leads to the desired result.
from pyspark.sql.window import Window
from pyspark.sql import functions as F
from pyspark.sql import types as T
w = Window.partitionBy('user_id').orderBy('unixtime')
df.withColumn('unixtime', F.unix_timestamp('activity_timestamp') / 60)\
.withColumn('lag', F.lag('unixtime').over(w))\
.withColumn('diff', F.col('lag') - F.col('unixtime'))\
.withColumn('group', F.sum(F.when(F.col('diff') > -5., 0).otherwise(1)).over(w))\
.groupBy('user_id', 'group').agg(F.first('activity_timestamp').alias('activity_start'),
F.last('activity_timestamp').alias('activity_end'))\
.drop('group').show()
Output
+-------+-------------------+-------------------+
|user_id| activity_start| activity_end|
+-------+-------------------+-------------------+
| 1|2021-06-01 19:00:00|2021-06-01 19:02:00|
| 1|2021-06-01 19:10:00|2021-06-01 19:11:00|
| 2|2021-06-01 19:01:00|2021-06-01 19:02:00|
+-------+-------------------+-------------------+