周/月/季度结束前 N 天的最大值

Question

我有一个数据框df，其中包含许多ids 的每日数据，示例：

| 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 | ...          |

如何确定季度结束前最后 N 天的最大 availability？ IE。如果 N = 3，那么我想在每个季度结束前的最后 3 天内找到每个 id 的最大值 availability。

我知道我可以对日期进行硬编码，但我想尽可能避免这种情况。 windowing 可以帮忙吗？

基于上述输入的预期输出：

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

Answer 1

您可以试试下面的这段代码，它获取季度结束前的最后 10 行，并计算按 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|
+----+-------+----+----------------+