Rolling Sum on Spark Dataframe 被其他列过滤

Question

我有一个 spark dataframe 看起来像这样，其中包含文章编号、国家/地区代码和日期的每个组合的一行，其中存在该组合的金额值。 这个 dataframe 中大约有 400,000 行。

articlenumber   countrycode   date         amount
--------------------------------------------------
4421-222-222    DE            2020-02-05   200
1234-567-890    EN            2019-05-23   42
1345-457-456    EN            2019-12-12   107

现在我需要一个额外的列“数量 12M”，它根据以下规则计算每行的值：

在每一行中，“金额 1200 万”应包含“金额”中所有值的总和，其中文章编号和国家/地区代码与该特定行中的值匹配，并且日期位于该行日期之前的 12 个月之间。

我是否需要为尚未获得值的日期/国家/地区/文章编号组合添加数量为 0 的行？

由于我不是编程方面的专家（工科学生），我需要一些帮助如何在处理 dataframe 的 python 脚本中实现这一点。

感谢您对此的任何想法。

Answer 1

编辑：

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

w = Window.partitionBy('articlenumber', 'countrycode').orderBy('date').orderBy('yearmonth').rangeBetween(-11, 0)

df.withColumn('yearmonth', f.expr('(year(date) - 2000) * 12 + month(date)')) \
  .withColumn('amount 12M', f.sum('amount').over(w)) \
  .orderBy('date').show(10, False)

+-------------+-----------+----------+------+---------+----------+
|articlenumber|countrycode|date      |amount|yearmonth|amount 12M|
+-------------+-----------+----------+------+---------+----------+
|4421-222-222 |DE         |2019-02-05|100   |230      |100       |
|4421-222-222 |DE         |2019-03-01|50    |231      |150       |
|1234-567-890 |EN         |2019-05-23|42    |233      |42        |
|1345-457-456 |EN         |2019-12-12|107   |240      |107       |
|4421-222-222 |DE         |2020-02-05|200   |242      |250       |
+-------------+-----------+----------+------+---------+----------+

---

我不确定确切的 12 个月，但这会起作用。

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

w = Window.partitionBy('articlenumber', 'countrycode').orderBy('unix_date').rangeBetween(- 365 * 86400, 0)

df.withColumn('unix_date', f.unix_timestamp('date', 'yyyy-MM-dd')) \
  .withColumn('amount 12M', f.sum('amount').over(w)) \
  .orderBy('date').show(10, False)

+-------------+-----------+----------+------+----------+----------+
|articlenumber|countrycode|date      |amount|unix_date |amount 12M|
+-------------+-----------+----------+------+----------+----------+
|4421-222-222 |DE         |2019-02-05|100   |1549324800|100       |
|4421-222-222 |DE         |2019-02-06|50    |1549411200|150       |
|1234-567-890 |EN         |2019-05-23|42    |1558569600|42        |
|1345-457-456 |EN         |2019-12-12|107   |1576108800|107       |
|4421-222-222 |DE         |2020-02-05|200   |1580860800|350       |
+-------------+-----------+----------+------+----------+----------+