How to get average of last 5 rows in a PySpark dataframe

Viewed 574

I am monitoring a few hundred user activity online and trying to get average records posted for the last 5 rows (ignoring the missing dates in between). I am using Window.partitionBy to partition between users and rowsBetween to get the last 5 records but instead of giving me average of the last 5 records for each subsequent rows it seems to include all other rows into the average.

This is what I have tried:

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

win = Window.partitionBy('user_id').orderBy(f.col('record_date')).rowsBetween(-5, 0)

test = df.sort(
  f.asc('record_date')
).withColumn(
  'last5_avg', f.mean('total').over(win)
)

This is the result I get:

enter image description here

This is the result I want:

enter image description here

1 Answers

I think you meant to do (-4, 0):

win = Window.partitionBy('user_id').orderBy(f.col('record_date')).rowsBetween(-4, 0)

because you want to include 5 rows, not 6 rows. Note that the range is inclusive.

Related