基于非缺货行的动态窗口移动平均计算需求（Spark）

要实现排除缺货行的动态窗口移动平均，核心是解决静态窗口无法根据缺货状态调整范围的问题。可以通过累计计数非缺货行来确定每行的有效窗口边界，再基于这个边界计算移动平均。

实现步骤 #

1. 标记初始临时缺货：先基于原始数据做临时缺货标记，避免后续计算出现循环依赖
2. 给行分配序号：按时间戳排序后给每行加行号，方便定位窗口位置
3. 累计非缺货行数量：计算从第一行到当前行的非缺货行累计数，以此确定需要回溯的行数
4. 确定动态窗口起始点：找到能凑够当前行+至少2个非缺货行的最早行索引
5. 计算动态移动平均：在确定的窗口范围内，排除缺货行计算平均
6. 更新最终缺货标记：用新计算的移动平均重新标记缺货状态

完整Spark代码 #

from pyspark.sql import Window
from pyspark.sql.functions import col, avg, when, lit, sum as spark_sum, row_number, max as spark_max

# 1. 生成临时缺货标记（先按原逻辑做初始标记，后续会更新）
temp_df = result_df.withColumn(
    "temp_stockout",
    when((col("hour_sales") == 0) | (col("hour_sales") < 0.5 * avg("hour_sales").over(Window.orderBy("formatted_timestamp").rowsBetween(-2, 0))), lit(1)).otherwise(lit(0))
)

# 2. 按时间戳排序，给每行加行号
ranked_df = temp_df.withColumn("row_num", row_number().over(Window.orderBy("formatted_timestamp")))

# 3. 累计非缺货行数量，同时计算窗口起始行号
window_spec = Window.orderBy("row_num").rowsBetween(Window.unboundedPreceding, 0)
cumulative_df = ranked_df.withColumn(
    "cumulative_non_stockout",
    spark_sum(when(col("temp_stockout") == 0, lit(1)).otherwise(lit(0))).over(window_spec)
).withColumn(
    "target_cumulative",
    col("cumulative_non_stockout") - 2  # 要凑够2个非缺货历史行，目标累计数为当前累计数减2
).withColumn(
    # 找到满足累计数条件的最早行号，作为窗口起始
    "start_row_num",
    spark_max(when(col("cumulative_non_stockout") <= col("target_cumulative"), col("row_num"))).over(window_spec)
).na.fill({"start_row_num": 1})  # 前几行不足2个非缺货行时，默认从第一行开始

# 4. 基于动态起始行号计算移动平均，自动排除缺货行
dynamic_window = Window.orderBy("row_num").rowsBetween(col("start_row_num") - col("row_num"), 0)
final_df = cumulative_df.withColumn(
    "3_avg_sales",
    avg(when(col("temp_stockout") == 0, col("hour_sales")).otherwise(None)).over(dynamic_window)
)

# 5. 更新最终缺货标记
final_df = final_df.withColumn(
    "stockout_0.5",
    when((col("hour_sales") == 0) | (col("hour_sales") < 0.5 * col("3_avg_sales")), lit(1)).otherwise(lit(0))
).drop("temp_stockout", "row_num", "cumulative_non_stockout", "target_cumulative", "start_row_num")

核心细节 #

• 用row_num定位行位置，避免时间戳重复导致的窗口计算错误
• 通过累计非缺货行数量动态调整窗口范围，确保每次计算平均时窗口内至少有2个非缺货观测值
• 计算平均时用when函数将缺货行的hour_sales设为None，Spark的avg会自动忽略空值，实现排除缺货行的效果

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