在 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 窗口函数 - rangeBetween 日期）

对于以月为单位的范围，这不起作用，因为一个月中的天数不是常数。知道如何使用 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|
# +----------+----------+----+

Answer 2

这并不理想，但应该足以估计大约月中左右的时间，以使相关月份进入该范围。

# define monthly windows of interest
window_dict = {
"w1": {"start": -40, "end": -1},
"w2": {"start": -70, "end": -1},
"w3": {"start": -100, "end": -1},
"w4": {"start": -130, "end": -1}
}

# set window - this one is for w3 
window_3M = (Window.partitionBy(item_id)
             .orderBy(F.expr("unix_date(to_date(ds, 'yyyy-MM-dd'))"))
             .rangeBetween(window_dict["w3"]["start"],window_dict["w3"]["end"])
)

# apply as needed
new_df = old_df.withColumn("SumOf3M", F.sum("some_column").over(window_3M))