PySpark基于日期边界的Window窗口实现6个月值求和问题

问题分析 #

你需要计算每行数据所在日期及之前6个月的VALUE字段总和，用rowsBetween/rangeBetween无法得到预期结果的原因：

• rowsBetween按行号偏移，不考虑日期实际间隔，若数据日期不连续，会包含超出6个月范围的行
• rangeBetween若基于日期的毫秒数计算，6个月的天数不固定（不同月份天数差异），无法精准匹配时间范围

以下提供两种可行的PySpark API实现方式：

方法1：自连接（兼容性好，适用于所有Spark版本） #

通过自连接筛选出当前行日期前6个月到当前日期的所有记录，再分组求和，逻辑直观且准确。

代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sum, add_months

# 初始化SparkSession
spark = SparkSession.builder.appName("Rolling6MonthsSum").getOrCreate()

# 测试数据
data = [
    ("2023-01-01", 10),
    ("2023-03-01", 20),
    ("2023-06-01", 30),
    ("2023-07-01", 40),
    ("2023-08-01", 50),
    ("2023-12-01", 60)
]

# 创建DataFrame并转换日期类型
df = spark.createDataFrame(data, ["dt", "VALUE"]).withColumn("dt", col("dt").cast("date"))

# 自连接计算滚动总和
result_df = df.alias("a").join(
    df.alias("b"),
    col("b.dt").between(add_months(col("a.dt"), -6), col("a.dt")),
    how="left"
).groupBy("a.dt", "a.VALUE").agg(sum("b.VALUE").alias("rolling_6months_sum")).orderBy("a.dt")

# 查看结果
result_df.show()

# 停止SparkSession
spark.stop()

方法2：窗口函数+Filter（Spark 3.1+支持） #

利用Spark 3.1及以上版本支持的窗口FILTER子句，在窗口聚合时筛选符合日期范围的记录，避免行号或固定时间间隔的误差。

代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import col, sum, expr, add_months

# 初始化SparkSession
spark = SparkSession.builder.appName("Rolling6MonthsSum").getOrCreate()

# 测试数据（同方法1）
data = [
    ("2023-01-01", 10),
    ("2023-03-01", 20),
    ("2023-06-01", 30),
    ("2023-07-01", 40),
    ("2023-08-01", 50),
    ("2023-12-01", 60)
]

df = spark.createDataFrame(data, ["dt", "VALUE"]).withColumn("dt", col("dt").cast("date"))

# 定义窗口（按日期排序，范围从起始行到当前行）
window_spec = Window.orderBy("dt").rowsBetween(Window.unboundedPreceding, Window.currentRow)

# 带Filter的窗口聚合（expr写法更直观）
result_df = df.withColumn(
    "rolling_6months_sum",
    expr("sum(VALUE) over (order by dt rows between unbounded preceding and current row) filter (where dt >= add_months(current_row().dt, -6))")
)

# 查看结果
result_df.show()

# 停止SparkSession
spark.stop()

预期结果 #

两种方法都会输出如下结果：

+----------+-----+-------------------+
|        dt|VALUE|rolling_6months_sum|
+----------+-----+-------------------+
|2023-01-01|   10|                 10|
|2023-03-01|   20|                 30|
|2023-06-01|   30|                 60|
|2023-07-01|   40|                100|
|2023-08-01|   50|                140|
|2023-12-01|   60|                180|
+----------+-----+-------------------+

内容的提问来源于stack exchange，提问作者Gabriel Sanches