Pyspark splitting session based on particular start and end events

user_13550521 0
Viewed 217

My df is ordered (by session and timestamp) and looks like the following;

    df = spark.createDataFrame(
        [[1, '2020-01-01 12:30:00.000', 'foo'], [1, '2020-01-01 12:31:00.000', 'bar'], [1, '2020-01-01 12:32:00.000', 'foo'],
    [1, '2020-01-01 12:33:00.000', 'foo'], [2, '2020-01-01 13:00:00.000', 'bar'], [2, '2020-01-01 13:01:00.000', 'foo'],
    [2, '2020-01-01 13:02:00.000', 'bar'], [2, '2020-01-01 13:03:00.000', 'foo']],
        ['session_id', 'timestamp', 'event']
    )
    df.show(truncate=False)
+----------+-----------------------+-----+
|session_id|timestamp              |event|
+----------+-----------------------+-----+
|1         |2020-01-01 12:30:00.000|foo  |
|1         |2020-01-01 12:31:00.000|bar  |
|1         |2020-01-01 12:32:00.000|foo  |
|1         |2020-01-01 12:33:00.000|foo  |
|2         |2020-01-01 13:00:00.000|bar  |
|2         |2020-01-01 13:01:00.000|foo  |
|2         |2020-01-01 13:02:00.000|bar  |
|2         |2020-01-01 13:03:00.000|foo  |
+----------+-----------------------+-----+

Where I want my sessions to follow a particular pattern. They need to start with event 'bar' and have one (or more) 'foo' events. Whenever a new 'bar' event takes place, I want to classify this as a new session. Events that fall out of this pattern should be dropped, e.g. the initial 'foo' event. Desired output should look like:

    df_res = spark.createDataFrame(
        [[1, '2020-01-01 12:31:00.000', 'bar'], [1, '2020-01-01 12:32:00.000', 'foo'],
    [1, '2020-01-01 12:33:00.000', 'foo'], [2, '2020-01-01 13:00:00.000', 'bar'], [2, '2020-01-01 13:01:00.000', 'foo'],
    [3, '2020-01-01 13:02:00.000', 'bar'], [3, '2020-01-01 13:03:00.000', 'foo']],
        ['session_id', 'timestamp',  'event']
    )
    df_res.show(truncate=False)
    +----------+------------------------+-----+
|session_id|timestamp               |event|
+----------+------------------------+-----+
|1         |2020-01-01 12:31:00.000 |bar  |
|1         |2020-01-01 12:32:00.000 |foo  |
|1         |2020-01-01 12:33:00.000 |foo  |
|2         |2020-01-01 13:00:00.000 |bar  |
|2         |2020-01-01 13:01:00.000 |foo  |
|3         |2020-01-01 13:02:00.000 |bar  |
|3         |2020-01-01 13:03:00.000 |foo  |
+----------+------------------------+-----+

I've tried doing a groupby and collect_list and then splitting or flattening, but i'm unsure how to proceed. Any help would be welcome!

df.groupBy("session_id").agg(F.collect_list("event").alias("list_event"))
#does not work
# tst_udf = udf(lambda l: split(l, 'bar'))
# df = df.withColumn("tst", tst_udf(col('list_event')))

-edit My end goal is to pivot this table and have one row per session, where I have variables about the 'bar' and (multiple) 'foo' events.

pyspark apache-spark-sql python
1 Answers

Try this:

welcome to SO

from pyspark.sql import functions as F
from pyspark.sql.window import Window

w=Window().orderBy("timestamp")
w2=Window().partitionBy("session_id").orderBy("timestamp")
w3=Window().partitionBy("session_id")
df.withColumn("timestamp", F.to_timestamp("timestamp", 'yyyy-MM-dd HH:mm:ss.SSS'))\
  .withColumn("session_id", F.sum(F.when((F.col("event")=='bar'),F.lit(1))\
                                         .otherwise(F.lit(0))).over(w))\
  .withColumn("rowNum", F.row_number().over(w2))\
  .withColumn("max", F.max("rowNum").over(w3))\
  .withColumn("first", F.when((F.col("rowNum")==1)&(F.col("event")=='foo'), F.lit(1))\
                       .otherwise(F.lit(0)))\
  .filter('max>=2 and first=0').drop(*['rowNum','sample_timestamp','max','first']).show()

#+----------+-------------------+-----+
#|session_id|          timestamp|event|
#+----------+-------------------+-----+
#|         1|2020-01-01 12:31:00|  bar|
#|         1|2020-01-01 12:32:00|  foo|
#|         1|2020-01-01 12:33:00|  foo|
#|         2|2020-01-01 13:00:00|  bar|
#|         2|2020-01-01 13:01:00|  foo|
#|         3|2020-01-01 13:02:00|  bar|
#|         3|2020-01-01 13:03:00|  foo|
#+----------+-------------------+-----+
user_101073890
Related