How to use the windows over function PySpark with time period constraint and other conditions

user_14183627 0
Viewed 481

Can I get some help on how to write this logic up in pyspark?

Suppose I have the the table as attached image shown.

So given: date, userid, visit, grouping (old) as inputs, I want to create a new column called grouping (new) such that it does the following:

for any given user id:

  1. first check and see what's the grouping (old). If it is != Bad, then grouping (New) = grouping (old)

  2. If the grouping (old) = Bad, then apply the most recent date's most recent visit's grouping (old) such that it is != Bad

  3. However, if the most recent grouping (old) from a prior date that is != Bad is more than 30 days away, then make grouping (new) = Bad (as the data is out of date)

enter image description here

what I've attempted which didn't work as expected:

days = lambda i: i * 86400 

user_30d_tracker = 
  Window.partitionBy("userid")
    .orderBy(f.col("date").cast("timestamp").cast("long"))
    .rangeBetween(-days(30), 0)
    .rowsBetween(Window.unboundedPreceding, Window.currentRow - 1)


df = (df.withColumn("Grouping(old)_YN",
      f.when(f.col("Grouping(old)")==f.lit("Bad"), "No")
       .otherwise(f.lit("Yes"))))

df = df.withColumn("Grouping_new", 
          f.max(f.when(f.col("Grouping(old)_YN") == f.lit("Yes"), 
          f.col("Grouping(old)"))).over(user_30d_tracker))
pyspark apache-spark-sql apache-spark
1 Answers

Suppose this is what your dataframe looks like

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

data = [[123, "20200101", 1, "Good_sub1"],
        [123, "20200101", 2, "Bad"],
        [123, "20200115", 1, "Bad"],
        [123, "20200115", 2, "Bad"],
        [123, "20200116", 1, "Good_sub2"],
        [123, "20200116", 2, "Bad"],
        [123, "20200116", 3, "Good_sub3"],
        [123, "20220901", 1, "Bad"]]

df = spark.createDataFrame(data, 
                           "userid:int, date:string, visit:int, `grouping(old)`:string")
df.show()
# +------+--------+-----+-------------+
# |userid|    date|visit|grouping(old)|
# +------+--------+-----+-------------+
# |   123|20200101|    1|    Good_sub1|
# |   123|20200101|    2|          Bad|
# |   123|20200115|    1|          Bad|
# |   123|20200115|    2|          Bad|
# |   123|20200116|    1|    Good_sub2|
# |   123|20200116|    2|          Bad|
# |   123|20200116|    3|    Good_sub3|
# |   123|20220901|    1|          Bad|
# +------+--------+-----+-------------+

days = lambda i: i * 86400 

user_30d_tracker = Window.partitionBy("userid")\
                         .orderBy(f.col("date").cast("timestamp").cast("long"))\
                         .rangeBetween(-days(30), 0)\
                         .rowsBetween(Window.unboundedPreceding, Window.currentRow - 1)

Let's take a look at the user_30d_tracker window. This window needs some changes considering the points below

  1. The order of the window seems to need the values of visit column beside the unix timestamp of date. With only timestamp of date in the orderBy clause, spark will not guarantee the rows with visit = 1 will come before visit = 2 for example. So, somehow we need to include visit column in the orderBy clause.
  2. Additional frame rowsBetween will overwrite the rangeBetween frame, hence it will be different from what were expected.

An option would be only using rangeBetween frame for the window. But, since rangeBetween frame only allows 1 column to be used in orderBy expression, we can use a workaround by adding visit value to the unix timestamp of date (this is like treating visit as how many seconds had passed since date started).

user_30d_tracker =  Window\
                        .partitionBy("userid")\
                        .orderBy(f.unix_timestamp("date", "yyyyMMdd") + f.col("visit"))\
                        .rangeBetween(-days(30), 0)

Then, to get the most recent non-'Bad' grouping(old) value, it would be better using last function (with ignorenulls=True) instead of max since it takes the recent value in the window, not the maximum of sorted strings. After that, use coalesce to fill the null values in the new column.

df = (df
      .withColumn("Grouping(old)_YN",
                  f.when(f.col("Grouping(old)") == f.lit("Bad"), "No")
                  .otherwise(f.lit("Yes")))
      .withColumn("Grouping_new", 
                  f.last(f.when(f.col("Grouping(old)_YN") == f.lit("Yes"), 
                                f.col("Grouping(old)")), ignorenulls=True).over(user_30d_tracker))
      .withColumn("Grouping_new", f.coalesce(f.col("Grouping_new"), f.col("Grouping(old)")))
      )
df.show()

# +------+--------+-----+-------------+----------------+------------+
# |userid|    date|visit|grouping(old)|Grouping(old)_YN|Grouping_new|
# +------+--------+-----+-------------+----------------+------------+
# |   123|20200101|    1|    Good_sub1|             Yes|   Good_sub1|
# |   123|20200101|    2|          Bad|              No|   Good_sub1|
# |   123|20200115|    1|          Bad|              No|   Good_sub1|
# |   123|20200115|    2|          Bad|              No|   Good_sub1|
# |   123|20200116|    1|    Good_sub2|             Yes|   Good_sub2|
# |   123|20200116|    2|          Bad|              No|   Good_sub2|
# |   123|20200116|    3|    Good_sub3|             Yes|   Good_sub3|
# |   123|20220901|    1|          Bad|              No|         Bad|
# +------+--------+-----+-------------+----------------+------------+
user_94778430
Related