PySpark中日期范围递归计算与结果转置的高效实现方案

我有一份包含10年销售数据的数据集kf，需要按部门统计每月销售额，并且以任意年月为基准（比如2023年4月），生成sale_1（当月销售额总和）、sale_2（前一个月总和）……以此类推的结构，无数据时填充null。

目前我通过循环12次执行以下PySpark代码生成主表，但不知道如何转置或纵向堆叠结果，方便用户按年月筛选查询数据快照，求PySpark的高效实现方案！

现有代码：

kf=main file with 10 years data
mf=kf.dropDuplicates(['dept_ID'])
gf=kf
month_list={'1','2','3','4','5','6','7','8','9','10','11','12'}
window = Window().partitionBy("dept_ID")
for i in month_list:
   df = gf.filter(gf.month==i).withColumn("sale_"+i, sum(coalesce('sale_amt'), lit(0))).over(window))
   df = df.dropDuplicates(['dept_ID'])
   mf = mf.join(df, mf.dept_ID==df.dept_ID, 'left').drop(df.dept_ID)

数据集kf示例：

dept_ID sale_amt    sale_date   sale_month  sale_year
1   10  4/1/2023    4   2023
1   60  4/1/2023    4   2023
1   30  3/1/2023    3   2023
1   15  3/1/2023    3   2023
1   12  2/1/2023    2   2023
1   10  1/1/2023    1   2023
1   90  1/1/2023    1   2023
1   40  12/1/2022   12  2022
1   40  11/1/2022   11  2022
1   75  10/1/2022   10  2022
1   30  9/1/2022    9   2022
1   50  9/1/2022    9   2022
1   25  8/1/2022    8   2022
1   40  8/1/2022    8   2022
1   70  7/1/2022    7   2022
1   80  5/1/2022    5   2022
1   10  5/1/2022    5   2022
1   45  4/1/2022    4   2022
1   15  4/1/2022    4   2022
2   10  4/1/2023    4   2023
2   60  4/1/2023    4   2023
2   30  3/1/2023    3   2023
2   15  3/1/2023    3   2023
2   12  2/1/2023    2   2023
2   10  1/1/2023    1   2023
2   90  1/1/2023    1   2023
2   40  12/1/2022   12  2023
2   40  11/1/2022   11  2023
2   80  10/1/2022   10  2023
2   30  9/1/2022    9   2023
3   50  9/1/2022    9   2023
3   25  8/1/2022    8   2023
3   40  8/1/2022    8   2023
3   70  7/1/2022    7   2023
3   80  5/1/2022    5   2023
3   10  5/1/2022    5   2023

期望结果：
以年月为每行基准，按部门展示当月及往前推11个月的销售额，无数据则填充null，示例结构如下：

步骤1：按部门+年月聚合销售总额 #

先通过一次分组聚合计算每个部门每一年每个月的销售总额，替代循环中重复的filter和聚合操作：

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 聚合部门-年月的销售总额
agg_df = kf.groupBy("dept_ID", "sale_year", "sale_month") \
    .agg(F.sum("sale_amt").alias("total_sale")) \
    .orderBy("dept_ID", "sale_year", "sale_month")

步骤2：用窗口函数生成偏移销售额 #

定义按部门分区、按年月排序的窗口，通过lag函数获取当前年月及往前推11个月的销售额，对应sale_1到sale_12：

# 定义时间排序窗口：确保年月顺序正确
window_spec = Window.partitionBy("dept_ID") \
    .orderBy(F.to_date(F.concat_ws("-", "sale_year", "sale_month"), "yyyy-MM"))

# 生成基准年月列，并依次计算sale_1到sale_12
result_df = agg_df.withColumn("base_year", F.col("sale_year")) \
    .withColumn("base_month", F.col("sale_month"))

for i in range(1, 13):
    # lag偏移量为i-1：偏移0是当月（sale_1），偏移1是前1个月（sale_2），以此类推
    result_df = result_df.withColumn(f"sale_{i}", F.lag("total_sale", i-1).over(window_spec))

# 清理冗余列
result_df = result_df.drop("sale_year", "sale_month", "total_sale")

步骤3：补全缺失的部门-年月组合（可选） #

如果需要确保所有部门的所有年月都存在（即使没有销售额），可以先生成部门和年月的笛卡尔积，再与聚合结果左连接补全：

# 获取所有唯一部门和年月
depts = kf.select("dept_ID").distinct()
years_months = kf.select("sale_year", "sale_month").distinct().orderBy("sale_year", "sale_month")

# 生成部门-年月笛卡尔积
all_dept_year_month = depts.crossJoin(years_months)

# 左连接聚合结果，补全缺失的销售额为null
full_agg_df = all_dept_year_month.join(agg_df, ["dept_ID", "sale_year", "sale_month"], "left") \
    .withColumn("total_sale", F.coalesce("total_sale", F.lit(None)))

# 重复步骤2生成最终结果
full_result_df = full_agg_df.withColumn("base_year", F.col("sale_year")) \
    .withColumn("base_month", F.col("sale_month"))

for i in range(1, 13):
    full_result_df = full_result_df.withColumn(f"sale_{i}", F.lag("total_sale", i-1).over(window_spec))

full_result_df = full_result_df.drop("sale_year", "sale_month", "total_sale")

方案优势 #

• 避免循环中多次filter和join操作，减少Spark作业调度开销
• 一次聚合+一次窗口操作完成所有计算，充分利用分布式计算能力
• 逻辑简洁清晰，易于维护和扩展（比如调整偏移月份数量）

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