如何基于缺失区间总和在PySpark中插值时间序列缺失值？

高效解决方案（纯PySpark内置函数） #

核心思路是通过窗口函数为每个缺失区间关联前后的有效total_mileage值，计算区间总里程差后均分，再推导每个缺失行的miles和total_mileage，全程不用UDF，适配百万级以上数据集。

步骤1：标记缺失行并关联前后有效累计里程 #

按用户分组、时间排序，标记原始缺失状态，同时用窗口函数向前/向后填充最近的有效total_mileage：

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

def interpolate_mileage(df):
    # 按用户分组、时间排序的基础窗口
    user_ts_window = Window.partitionBy("user_id").orderBy("timestamp")
    
    # 标记原始是否缺失（用于最终的was_missing列）
    df = df.withColumn("was_missing", F.when(F.col("miles").isNull() & F.col("total_mileage").isNull(), 1).otherwise(0))
    
    # 定义向前填充窗口：取当前行及之前最后一个非null的total_mileage
    forward_fill_window = user_ts_window.rowsBetween(Window.unboundedPreceding, 0)
    # 定义向后填充窗口：取当前行及之后第一个非null的total_mileage
    backward_fill_window = user_ts_window.rowsBetween(0, Window.unboundedFollowing)
    
    # 为每个行获取前后最近的有效累计里程
    df = df.withColumn("prev_valid_total", F.last(F.col("total_mileage"), ignorenulls=True).over(forward_fill_window))
    df = df.withColumn("next_valid_total", F.first(F.col("total_mileage"), ignorenulls=True).over(backward_fill_window))
    
    # 计算每个分组内的行号，用于后续计算缺失区间内的位置
    df = df.withColumn("row_num", F.row_number().over(user_ts_window))
    
    return df

步骤2：识别缺失区间并计算插值值 #

为每个缺失区间计算总里程差、区间长度，推导每个缺失行的miles和total_mileage：

def calculate_interpolated_values(df):
    # 按用户分组的窗口，用于计算区间的起始和结束行号
    user_window = Window.partitionBy("user_id")
    
    # 为每个行找到所属缺失区间的起始行号（前一个非缺失行的row_num）
    df = df.withColumn("gap_start_row", F.when(F.col("was_missing") == 1, F.last(F.when(F.col("was_missing") == 0, F.col("row_num")), ignorenulls=True).over(user_window.rowsBetween(Window.unboundedPreceding, 0))).otherwise(F.col("row_num")))
    
    # 为每个行找到所属缺失区间的结束行号（后一个非缺失行的row_num）
    df = df.withColumn("gap_end_row", F.when(F.col("was_missing") == 1, F.first(F.when(F.col("was_missing") == 0, F.col("row_num")), ignorenulls=True).over(user_window.rowsBetween(0, Window.unboundedFollowing))).otherwise(F.col("row_num")))
    
    # 计算缺失区间的总里程差和区间内的行数
    df = df.withColumn("total_gap_miles", F.col("next_valid_total") - F.col("prev_valid_total"))
    df = df.withColumn("gap_length", F.col("gap_end_row") - F.col("gap_start_row"))
    
    # 计算插值后的miles：非缺失行用原值，缺失行用总里程差均分
    df = df.withColumn("interpolated_miles", 
                       F.when(F.col("was_missing") == 0, F.col("miles"))
                       .otherwise(F.col("total_gap_miles") / F.col("gap_length")))
    
    # 计算插值后的total_mileage：基于前一个有效累计里程 + 之前所有插值miles的和
    running_total_window = Window.partitionBy("user_id", "gap_start_row").orderBy("row_num")
    df = df.withColumn("running_miles_sum", F.sum(F.col("interpolated_miles")).over(running_total_window))
    df = df.withColumn("interpolated_total", 
                       F.when(F.col("was_missing") == 0, F.col("total_mileage"))
                       .otherwise(F.col("prev_valid_total") + F.col("running_miles_sum") - F.col("interpolated_miles")))
    
    # 清理临时列，保留最终需要的字段
    final_df = df.select("user_id", "timestamp", 
                         F.col("interpolated_miles").alias("miles"), 
                         F.col("interpolated_total").alias("total_mileage"), 
                         "was_missing")
    
    return final_df

步骤3：整合调用 #

将两个函数组合处理数据集：

# 调用示例
processed_df = calculate_interpolated_values(interpolate_mileage(raw_df))
processed_df.show()

方案优势 #

• 全程使用PySpark内置窗口函数，无UDF或Pandas转换，性能适配大规模数据集
• 基于total_mileage前后有效值直接计算区间总差，避免分步插值的低效
• 按用户和缺失区间分组计算，逻辑清晰且并行度高

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