在 PySpark 中使用 rangeBetween 考虑月而不是天

Question

我正在寻找如何将这块 SQL 代码翻译成 PySpark 语法。

SELECT MEAN(some_value) OVER (
    ORDER BY yyyy_mm_dd
    RANGE BETWEEN INTERVAL 3 MONTHS PRECEDING AND CURRENT ROW
    ) AS mean
FROM
    df

如果以上是以天为单位表示的范围，则可以使用类似的方法轻松完成

.orderBy(F.expr("datediff(col_name, '1000')")).rangeBetween(-7, 0)

（另见 ZygD 的解决方案： Spark Window Functions - rangeBetween dates ）

然而，对于几个月的范围，这不起作用，因为一个月的天数不是一个常数。 知道如何使用 PySpark 语法执行考虑月份的范围吗？

Answer 1

您可以“借用”完整的 SQL 列表达式并在 PySpark 中使用它。

输入：

from pyspark.sql import functions as F
df = spark.createDataFrame(
    [('2022-05-01', 1),
     ('2022-06-01', 2),
     ('2022-07-01', 3),
     ('2022-08-01', 4),
     ('2022-09-01', 5)],      
    ['yyyy_mm_dd', 'some_value']
).withColumn('yyyy_mm_dd', F.to_date('yyyy_mm_dd'))

脚本：

df = df.withColumn('mean', F.expr("mean(some_value) over (order by yyyy_mm_dd range between interval 3 months preceding and current row)"))

df.show()
# +----------+----------+----+
# |yyyy_mm_dd|some_value|mean|
# +----------+----------+----+
# |2022-05-01|         1| 1.0|
# |2022-06-01|         2| 1.5|
# |2022-07-01|         3| 2.0|
# |2022-08-01|         4| 2.5|
# |2022-09-01|         5| 3.5|
# +----------+----------+----+