python - Max Value in N days before end of week/month/quarter

Question

I have a dataframe df which contains daily data for many ids, sample:

| yyyy_mm_dd | id   | availability |
|------------|------|--------------|
| 2020-01-01 | 1334 | 300          |
| 2020-01-02 | 1334 | 676          |
| 2020-01-03 | 1334 | 463          |
| ...        | 1334 | ...          |
| 2020-03-29 | 1334 | 564          |
| 2020-03-30 | 1334 | 765          |
| 2020-03-31 | 1334 | 7564         |
| ...        | 1334 | ...          |
| 2020-06-28 | 1334 | 4634         |
| 2020-06-29 | 1334 | 65           |
| 2020-06-30 | 1334 | 643          |
| ...        | 1334 | ...          |

How could I identify the maximum availability in the last N days prior to end of a quarter? I.e. if N = 3 then I would like to find the maximum availability per id, in the last 3 days to the prior of each quarter end.

I know I could hardcode the dates but I would like to avoid that if possible. Is this something that windowing could help with?

Expected output based on above input:

| id   | year | quarter | max_availability |
|------|------|---------|------------------|
| 1334 | 2020 | 01      | 7564             |
| 1334 | 2020 | 02      | 4634             |

question from:https://stackoverflow.com/questions/66063623/max-value-in-n-days-before-end-of-week-month-quarter

Answer 1

You can try this code below, which gets the last 10 rows before the end of a quarter and calculates the maximum availability grouped by id.

from pyspark.sql import functions as F, Window

df2 = df.withColumn(
    'rn',
    F.row_number().over(
        Window.partitionBy(
            'id',
            F.quarter('yyyy_mm_dd'), 
            F.year('yyyy_mm_dd')
        ).orderBy(F.desc('yyyy_mm_dd'))
    )
).filter('rn <= 10').groupBy(
    'id',
    F.quarter('yyyy_mm_dd').alias('quarter'), 
    F.year('yyyy_mm_dd').alias('year')
).agg(
    F.max('availability').alias('max_availability')
)

df2.show()
+----+-------+----+----------------+
|  id|quarter|year|max_availability|
+----+-------+----+----------------+
|1334|      1|2020|            7564|
|1334|      2|2020|            4634|
+----+-------+----+----------------+