Rolling Average

PySpark data engineering interview problem. Difficulty: intermediate. Pattern: Window Functions. About 16 minutes. Part of the Pro drill bank.

3-day rolling average of daily_amount. Treat this as a production helper: match the contracted return shape, including empty and duplicate inputs.

A rolling (or moving) average smooths a number that jumps around. For each day, you average the value of that day and the 2 days before it. On the first days there are fewer than 3 days to use, so the average uses what is available. Sort the rows of df by date and compute rolling_avg_3d, the average of daily_amount over the current row and the 2 rows before it (rowsBetween(-2, 0)). Return date, daily_amount and rolling_avg_3d, sorted by date. The editor creates the table df and ends with result = df. Change the last line to build your answer.

Requirements

• Use a window ordered by date with rowsBetween(-2, 0).
• Add rolling_avg_3d with F.avg("daily_amount").over(w).
• Return date, daily_amount, rolling_avg_3d, sorted by date.
• Store the DataFrame in result.

Constraints

• rowsBetween(-2, 0) means: 2 rows before up to the current row.
• There are 5 daily rows, one per day, with no missing days.
• The first 2 days average over fewer than 3 rows.

Examples

Input: daily amounts Output: date | daily_amount | rolling_avg_3d 2024-01-01 | 10.0 | 10.0 2024-01-02 | 20.0 | 15.0 2024-01-03 | 30.0 | 20.0 2024-01-04 | 40.0 | 30.0 2024-01-05 | 50.0 | 40.0 Early days use a shorter available frame.

Topics: lakebench, pyspark, rowsBetween, avg.

More PySpark interview questions · All interview problems · Learn data engineering