Pyspark: finding blocks of timestamp data without pause

user_8103683 0
Viewed 68

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"))
)
pyspark
2 Answers

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|
+-------+-------------------+-------------------+
user_21298010

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|
+-------+-------------------+-------------------+
user_95343900
Related