如何用PySpark窗口函数计算无空值连续6个月的平均值

原始数据与问题 #

你的原始PySpark DataFrame定义如下：

df = spark.createDataFrame(
  [(1,1,10), (2,1,10), (3,1,None),(4,1,10),(5,1,10),(6,1,20)  \
   ,(7,1,20), (1,2,10),(2,2,10),(3,2,10),(4,2,20),(5,2,20)],
  ["Month","customer","amount"])

原窗口函数实现会忽略空值计算平均值，无法满足仅当连续6个月的amount无空值时才计算平均值，否则返回Null的需求。

解决方案：窗口函数结合非空计数过滤 #

可以通过在窗口内同时计算平均值和非空值数量，再通过条件判断实现需求，具体步骤如下：

1. 定义窗口规则：按客户分区，按月份排序，窗口范围为当前月份往前推5个月到当前月份
2. 计算窗口内的平均值和非空值计数
3. 仅当非空值计数等于6时保留平均值，否则设为Null

完整代码：

from pyspark.sql import Window
from pyspark.sql.functions import avg, count, when, col

# 定义窗口
windowPartition = Window.partitionBy("customer").orderBy("Month").rangeBetween(Window.currentRow - 5, Window.currentRow)

# 添加计算列并过滤结果
df_result = df.withColumn("avg_temp", avg("amount").over(windowPartition)) \
              .withColumn("non_null_cnt", count("amount").over(windowPartition)) \
              .withColumn("avg_6_month", when(col("non_null_cnt") == 6, col("avg_temp")).otherwise(None)) \
              .drop("avg_temp", "non_null_cnt")

# 查看结果
display(df_result.orderBy("customer", "Month"))

结果说明 #

• 客户1的第1-5行：窗口内数据不足6条或包含空值，avg_6_month为Null；第6行开始窗口内6个月amount均非空，正常计算平均值
• 客户2只有5个月数据，所有行的avg_6_month均为Null，符合需求

扩展：处理非连续月份场景 #

如果你的Month字段存在断档（比如缺某个月的数据），需要取最近的6条记录而非按月份范围计算，只需将窗口的rangeBetween替换为rowsBetween：

windowPartition = Window.partitionBy("customer").orderBy("Month").rowsBetween(Window.currentRow - 5, Window.currentRow)

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