Pyspark conditionally increasing the index of a column

Question

Is there a good way to create a condition where I can achieve the following outcome for the session column:

+------+-----------------------+-----------------------+------+-----------+
|userId|timestamp              |timestamp_prev         |diff  |session    |
+------+-----------------------+-----------------------+------+-----------+
|123456|2022-01-15 19:30:21.789|2022-01-15 19:29:48.18 |33    |Session 2|
|123456|2022-01-15 19:29:48.18 |2022-01-15 19:29:16.933|32    |Session 2|
|123456|2022-01-15 19:29:16.933|2022-01-15 19:29:08.062|8     |Session 2|
|123456|2022-01-15 19:29:08.062|2022-01-14 08:00:33.126|127715|Session 2|
|123456|2022-01-14 08:00:33.126|2022-01-14 08:00:30.807|3     |Session 1|
|123456|2022-01-14 08:00:30.807|2022-01-14 08:00:12.627|18    |Session 1|
|123456|2022-01-14 08:00:12.627|2022-01-14 08:00:09.703|3     |Session 1|

I have the following code right now that partially does what I want but the condition has not been set correctly. Any tips? Should also be partitioned by each userId.

threshold = 1000 #sec

df_1 = df\
    .withColumn('timestamp', to_timestamp('timestamp'))\
    .withColumn('timestamp_prev', to_timestamp('timestamp_prev'))\
    .withColumn("diff", 
                when(isnull(unix_timestamp('timestamp') - unix_timestamp('timestamp_prev')), 0)
                .otherwise(unix_timestamp('timestamp') - unix_timestamp('timestamp_prev')))
df_2 = df_1\
        .withColumn('session',
                    when(df.diff > threshold, 'New session')
                    .otherwise('Old session'))

+------+-----------------------+-----------------------+------+-----------+
|userId|timestamp              |timestamp_prev         |diff  |session    |
+------+-----------------------+-----------------------+------+-----------+
|123456|2022-01-15 19:30:21.789|2022-01-15 19:29:48.18 |33    |Old session|
|123456|2022-01-15 19:29:48.18 |2022-01-15 19:29:16.933|32    |Old session|
|123456|2022-01-15 19:29:16.933|2022-01-15 19:29:08.062|8     |Old session|
|123456|2022-01-15 19:29:08.062|2022-01-14 08:00:33.126|127715|New session|
|123456|2022-01-14 08:00:33.126|2022-01-14 08:00:30.807|3     |Old session|
|123456|2022-01-14 08:00:30.807|2022-01-14 08:00:12.627|18    |Old session|
|123456|2022-01-14 08:00:12.627|2022-01-14 08:00:09.703|3     |Old session|

user_17682024 · Answer

I found a solution to my own question in a previous post (link to post)

df2 = df\
    .withColumn('to_add_number',
                when(df.diff > threshold, 1)
                .otherwise(0))\
    .withColumn("order_id", 
                monotonically_increasing_id())\
    .withColumn("session", 
                sum("to_add_number").over(Window().partitionBy("userId").orderBy("order_id"))+1)\
    .orderBy('order_id')\
    .drop('to_add_number', 'order_id')

Pyspark conditionally increasing the index of a column

1 Answers