PySpark实现按ID计算指定年份前3年的money求和

给定包含id、year、money列的PySpark DataFrame，示例数据如下：

id    year    money
1      2019    10
1      2018    15
1      2013    13
1      2009    10
1      2015    10
1      2014    11

需求：对每个id和year，计算**该年份之前连续3年（不含当前年份）**的money总和。例如：

• 2019年需计算2018、2017、2016年的money之和，仅2018年有数据，总和为15；
• 2015年需计算2014、2013、2012年的money之和，前两年有数据，总和为24。

期望输出DataFrame：

id     year    sum_money
1      2019      15
1      2018      10
1      2015      24
1      2014      13
1      2013      0
1      2009      0

提问：如何实现该需求？lag函数是否支持此类指定年份的查询？或是否有其他更简便的方法？当前已采用按年份降序计算累积求和，再通过年份窗口差值得到结果，希望找到更简单的实现方式。

1. 关于lag函数的局限性 #

lag函数是基于窗口内的行偏移量（比如取前1行、前2行的数据），而非基于year的数值范围。由于数据中年份不连续，直接用lag无法准确匹配到“前3年”的所有数据，因此不适合直接用lag实现该需求。

2. 更简便的实现：基于数值范围的窗口聚合 #

可以利用PySpark窗口函数的范围边界（rangeBetween），按year的数值范围定义窗口，直接计算符合条件的money总和，步骤如下：

步骤1：定义窗口

按id分组，按year排序，窗口范围设置为当前年份减3（不包含）到当前年份减1（包含），即rangeBetween(-3, -1)。这里的范围是基于排序后的year数值差：

• -3表示当前year减去3（即year - 3）
• -1表示当前year减去1（即year - 1）

步骤2：计算总和并处理空值

用sum聚合窗口内的money，再用coalesce将空值替换为0（无符合条件数据时总和为0）。

完整代码示例

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import sum, coalesce

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

# 示例数据
data = [
    (1, 2019, 10),
    (1, 2018, 15),
    (1, 2013, 13),
    (1, 2009, 10),
    (1, 2015, 10),
    (1, 2014, 11)
]
df = spark.createDataFrame(data, ["id", "year", "money"])

# 定义窗口：按id分组，year排序，范围是year-3到year-1（不含当前year）
window_spec = Window.partitionBy("id").orderBy("year").rangeBetween(-3, -1)

# 计算sum_money，空值替换为0
result_df = df.withColumn("sum_money", coalesce(sum("money").over(window_spec), 0))

# 展示结果
result_df.orderBy("year", ascending=False).show()

代码说明

• rangeBetween(-3, -1)：当按year排序时，这个范围会匹配所有满足current_year - 3 < year < current_year的行，也就是当前年份的前3年（不含当前年）。
• coalesce(..., 0)：如果窗口内没有数据（比如2009年，前3年没有任何记录），sum会返回null，用coalesce替换为0，符合需求。

3. 对比现有方法 #

这种基于范围窗口的方式无需额外处理累积求和与年份差值，直接通过窗口定义匹配目标年份范围，代码更简洁直观，性能也更优（减少了中间计算步骤）。

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