Spark中如何在窗口数据不完整时跳过窗口函数计算？

需求说明 #

基于前一天的时间窗口计算滚动平均，但当窗口内没有覆盖完整24小时数据时（例如初始时段回溯不足一天），返回Null；只有当窗口覆盖完整一天数据时，才计算并返回平均值。

期望输出 #

dateTime group val rolling_average
2023-01-01 00:00:00     1   1            Null
2023-01-01 00:05:00     1   2            Null
2023-01-01 00:10:00     1   3            Null
2023-01-01 00:15:00     1   4            Null
...
2023-01-02 00:00:00     1   2             2.5

当前错误输出 #

dateTime group val rolling_average
2023-01-01 00:00:00     1   1               1
2023-01-01 00:05:00     1   2             1.5
2023-01-01 00:10:00     1   3               2
2023-01-01 00:15:00     1   4             2.5
...
2023-01-02 00:00:00     1   2             2.5

当前窗口定义 #

window = Window.partitionBy("group").orderBy(f.col("dateTime").cast("long")).rangeBetween(-86400, 0)

现有方案的局限 #

此前尝试通过行号判断是否返回Null，但该方案硬编码了数据频率（如5分钟间隔对应288行），无法适配小时级或其他任意时间间隔的数据，通用性不足。

通用解决方案 #

无需修改窗口构造逻辑，通过判断窗口内的最早时间是否覆盖完整24小时范围即可实现需求：

from pyspark.sql import functions as f

# 保留原时间窗口定义
time_window = Window.partitionBy("group").orderBy(f.col("dateTime").cast("long")).rangeBetween(-86400, 0)

df_result = (
    df
    # 计算滚动平均值
    .withColumn("temp_avg", f.avg("val").over(time_window))
    # 获取窗口内的最早时间戳
    .withColumn("window_start", f.min("dateTime").over(time_window))
    # 判断窗口是否覆盖完整24小时：最早时间 <= 当前时间前一天的同一时刻
    .withColumn("rolling_average", 
                f.when(
                    f.col("window_start") <= f.date_sub(f.col("dateTime"), 1),
                    f.col("temp_avg")
                ).otherwise(f.lit(None)))
    # 清理中间临时列
    .drop("temp_avg", "window_start")
)

方案说明 #

• 完全基于时间维度判断，不依赖数据采样频率，适用于任意时间间隔的数据集
• 通过date_sub(f.col("dateTime"), 1)获取当前时间前一天的同一时刻，只要窗口内的最早时间不晚于该时刻，即可确认窗口覆盖了完整24小时数据
• 逻辑简洁清晰，避免了硬编码行号的局限性

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