Time based window function in Pyspark

user_14795898 0
Viewed 728

My goal is to calculate another column, keeping the same number of rows as the original DataFrame, where I can show the mean balance for each user for the last 30 days.

I guess it can be done using Window Functions, partitioning by user and somehow limiting the rows which are between the current date and 30 days before, but I don't know how to implement it in PySpark.

I have the following Spark DataFrame:

userId date balance
A 09/06/2020 100
A 03/07/2020 200
A 05/08/2020 600
A 30/08/2020 1000
A 15/09/2020 500
B 03/01/2020 100
B 05/04/2020 200
B 29/04/2020 600
B 01/05/2020 1600

My desired output DataFrame would be:

userId date balance mean_last_30days_balance
A 09/06/2020 100 100
A 03/07/2020 200 150
A 05/08/2020 600 600
A 30/08/2020 1000 800
A 15/09/2020 500 750
B 03/01/2020 100 100
B 05/04/2020 200 200
B 29/04/2020 600 400
B 01/05/2020 1600 800 
from datetime import datetime
from pyspark.sql import types as T

data = [("A",datetime.strptime("09/06/2020",'%d/%m/%Y'),100),
        ("A",datetime.strptime("03/07/2020",'%d/%m/%Y'),200),
        ("A",datetime.strptime("05/08/2020",'%d/%m/%Y'),600),
        ("A",datetime.strptime("30/08/2020",'%d/%m/%Y'),1000),
        ("A",datetime.strptime("15/09/2020",'%d/%m/%Y'),500),
        ("B",datetime.strptime("03/01/2020",'%d/%m/%Y'),100),
        ("B",datetime.strptime("05/04/2020",'%d/%m/%Y'),200),
        ("B",datetime.strptime("29/04/2020",'%d/%m/%Y'),600),
        ("B",datetime.strptime("01/05/2020",'%d/%m/%Y'),1600)]

schema = T.StructType([T.StructField("userId",T.StringType(),True),
                       T.StructField("date",T.DateType(),True), 
                       T.StructField("balance",T.StringType(),True)
                      ])
 
sdf_prueba = spark.createDataFrame(data=data,schema=schema)
sdf_prueba.printSchema()
sdf_prueba.orderBy(F.col('userId').asc(),F.col('date').asc()).show(truncate=False)
pyspark apache-spark-sql window-functions
1 Answers

You can use the RANGE BETWEEN keyword:

sdf_prueba.createOrReplaceTempView("table1")

spark.sql(
    """SELECT *, mean(balance) OVER (
        PARTITION BY userid 
        ORDER BY CAST(date AS timestamp)  
        RANGE BETWEEN INTERVAL 30 DAYS PRECEDING AND CURRENT ROW
     ) AS mean FROM table1""").show()


+------+----------+-------+-----+
|userId|      date|balance| mean|
+------+----------+-------+-----+
|     A|2020-06-09|    100|100.0|
|     A|2020-07-03|    200|150.0|
|     A|2020-08-05|    600|600.0|
|     A|2020-08-30|   1000|800.0|
|     A|2020-09-15|    500|750.0|
|     B|2020-01-03|    100|100.0|
|     B|2020-04-05|    200|200.0|
|     B|2020-04-29|    600|400.0|
|     B|2020-05-01|   1600|800.0|
+------+----------+-------+-----+

If you want to use the pyspark API, you need to convert days to unix seconds in order to use rangeBetween

one_month_in_seconds = 2629743 # ?
w = (
    Window.partitionBy("userid")
    .orderBy(unix_timestamp(col("date").cast("timestamp")))
    .rangeBetween(-one_month_in_seconds, Window.currentRow)
)

sdf_prueba.select(col("*"), mean("balance").over(w).alias("mean")).show()

+------+----------+-------+-----+
|userId|      date|balance| mean|
+------+----------+-------+-----+
|     A|2020-06-09|    100|100.0|
|     A|2020-07-03|    200|150.0|
|     A|2020-08-05|    600|600.0|
|     A|2020-08-30|   1000|800.0|
|     A|2020-09-15|    500|750.0|
|     B|2020-01-03|    100|100.0|
|     B|2020-04-05|    200|200.0|
|     B|2020-04-29|    600|400.0|
|     B|2020-05-01|   1600|800.0|
+------+----------+-------+-----+
user_89869750
Related