Generate unique session id for each row of dataframe using custom logic

user_1847212 0
Viewed 27

I have dataframe sorted by user_id and timestamp and it looks like the following:

user_id timestamp
1 1661941403

I want to generate unique session_id for each rows by following rules:

  1. If diff between two rows is less than 5 min (and user_id is the same) then it is same session_id
  2. Otherwie if diff between two timestamps is more than 5 min (or user_id is different) then use session_id + 1

So my first attempt is to map dataframe to another dataframe, and use LongAccumulator to get new session_id, e.g.

  val sessionId: LongAccumulator = sparkSession.sparkContext.longAccumulator("sessionId")

df.select(col("user_id"), col("timestamp"), lag("user_id", 1, 0), lag("timestamp", 1, 0))

.rdd.map(
row => {
  val userId = row.getLong(0)
  val time = row.getLong(1)
  val prevTime = row.getLong(2)
  val prevUserId = row.getLong(3)

  if (prevTime != 0 && prevUserId != 0) {
    if (userId == prevUserId && ((time - prevTime) / 60 > 5)) {
      sessionId.add(1)
    } else if (userId != prevUserId) {
      sessionId.add(1)
    }
  }

  (userId,time,sessionId.value)
}).toDF("user_id", "time", "session_id")

However, as far as I understood it is a wrong idea to use accumulator. So what can I do?

pyspark apache-spark-sql apache-spark
1 Answers

With Spark sql the solution is quite straightforward:

// Generate sample data
val df = Seq(
    (1,1661941403),
    (1,1661941408),
    (1,1661941412),
    (1,1661962245)).toDF("user_id","timestamp")
+-------+----------+
|user_id| timestamp|
+-------+----------+
|      1|1661941403|
|      1|1661941408|
|      1|1661941412|
|      1|1661962245|
+-------+----------+

// We'll use lag() analytic function, and for that we declare our window and period against which to compare.
import org.apache.spark.sql.expressions.Window
val w = Window.partitionBy("user_id").orderBy("timestamp")
val five_minutes = lit(5*60)

// If timestamp - prev(timestamp) >= 5 min then the session is new.
// Cumulative sum of all session starts will be our session_id.
df.withColumn(
    "session_is_new", 
    when(
        coalesce(
            col("timestamp") - lag("timestamp",1).over(w), 
            five_minutes
        ) >= five_minutes, 
        1
    ).otherwise(0)
).withColumn(
    "session_id", 
    sum("session_is_new").over(w)
).show
+-------+----------+-----------+----------+
|user_id| timestamp|session_new|session_id|
+-------+----------+-----------+----------+
|      1|1661941403|          1|         1|
|      1|1661941408|          0|         1|
|      1|1661941412|          0|         1|
|      1|1661962245|          1|         2|
+-------+----------+-----------+----------+
user_190322060
Related