基于PySpark窗口函数实现按条件校验滞后行的状态筛选

核心逻辑 #

无需拆分两种场景单独处理，只需要对每个ID分组后按时间升序排序，判断每条STAT=200的记录后续所有状态的最小值是否符合要求即可，可直接兼容你已经实现的「最后一条是200」的场景。

具体实现步骤 #

• 第一步：导入依赖的函数与窗口类

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, min
from pyspark.sql.window import Window

• 第二步：定义按ID分组、时间升序排序的窗口

# 按ID分区，TIME升序排列，窗口范围取当前行之后的所有行
w = Window.partitionBy("ID").orderBy("TIME").rowsBetween(Window.currentRow + 1, Window.unboundedFollowing)

• 第三步：新增后续状态最小值列，过滤符合要求的记录

# 此处替换为你的源数据集变量名即可，示例中假设源数据集名为df
result_df = df.withColumn("next_min_stat", min("STAT").over(w)) \
    .filter(
        (col("STAT") == 200) & 
        # 满足任意一个条件即可：1. 后续无记录（是当前ID的最后一条状态）；2. 后续所有状态的最小值≥200
        (col("next_min_stat").isNull() | (col("next_min_stat") >= 200))
    ) \
    # 可选操作：不需要中间计算列可直接删除
    .drop("next_min_stat")

效果验证 #

• 对应示例1：ID1的200状态后存在150，next_min_stat取值为150<200，会被过滤，仅保留ID2的200记录，符合要求
• 对应示例2：ID1的200状态后为350、400，next_min_stat取值为350≥200，会被保留，两个ID的200记录都留存，符合要求

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