PySpark DataFrame基于前一行数据更新drift_MS列结果异常问题

代码问题分析 #

你写的代码存在三个核心问题：

1. 逐行依赖的递归计算逻辑无法用普通lag窗口函数实现。你计算drift_MS时用lag(col('drift_MS'),1)取上一行结果，但drift_MS是本次计算新增的列，lag只能取到原始表中该列的值（原始表没有这列，所以拿到的全是null），计算自然不符合预期。
2. 未处理初始值：第一行的前置_MS为NaN、前置drift_MS为null，你没有设置初始值0，后续计算会受null传导影响。
3. NaN比较异常：第一行_MS为NaN，NaN和任意数值比较都会返回False，会干扰第二行的判断逻辑。

正确实现方案 #

你不需要逐行拿上一行的drift_MS计算，只需要先算出每一行需要累加的增量值，再对增量做窗口累加求和即可得到正确的drift_MS，实现代码如下：

import pyspark.sql.functions as f
from pyspark.sql.window import Window
from pyspark.sql.functions import col, when, isnan, lag, sum

# 定义基础排序窗口
w1 = Window.partitionBy('SEQ_ID').orderBy(col('TIME_STAMP').asc())
# 定义累加窗口：从分区第一行到当前行
w_cumsum = w1.rowsBetween(Window.unboundedPreceding, 0)

# 先把_MS列的NaN转为null，避免比较逻辑异常
df = df.withColumn('_MS', when(isnan(col('_MS')), None).otherwise(col('_MS')))
# 取前一行的_MS值
df = df.withColumn('prev_MS', lag(col('_MS'), 1).over(w1))

# 计算每一行的增量值
df = df.withColumn('delta',
    when(col('prev_MS').isNull(), 0)  # 前一行无值（第一行）增量为0
    .when((col('_MS') >=3) & (col('prev_MS') < col('_MS')), 100)
    .when((col('_MS') <3) & (col('prev_MS') < col('_MS')), 1)
    .otherwise(0)
)

# 累加增量得到最终drift_MS
df2 = df.withColumn('drift_MS', sum(col('delta')).over(w_cumsum)).drop('prev_MS', 'delta')

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