如何在PySpark中为DataFrame行标记before/after/event标签

我有一个包含ID、Timestamp、Event列的PySpark DataFrame，需要按ID分组实现以下标记逻辑：

• Event=1的行标记为event
• 该行之前的所有行标记为before
• 该行之后的所有行标记为after

尝试了自定义label函数，但Event=1后的部分行仍被错误标记为before，求正确实现方法。

输入输出示例 #

输入DataFrame #

期望输出 #

错误代码分析 #

之前的自定义函数依赖lag逐行判断，这种方式无法批量标记event之后的所有行——lag只能取上一行的值，一旦中间某行判断失效，后续行都会错误标记为before。此外代码中还出现了未定义的列isHypoProtectEnabled，这也是问题之一。

错误代码：

def label(df_):
    remove = ['type1']
    df_ = (
        df_
        .withColumn('type1', F.when((F.col("Event") == 0) & (F.lag(F.col("Event"), 1).over(Window.partitionBy('ID').orderBy('Timestamp')) == 1), 
                                         F.lit('after')))
        .withColumn('type2', F.when((F.col("isHypoProtectEnabled") == 0) & ((F.lag(F.col("Event"), 1).over(Window.partitionBy('ID').orderBy('Timestamp')) == 1) | (F.lag(F.col("type1"), 1).over(Window.partitionBy('ID').orderBy('Timestamp')) == 'after')), 
                                         F.lit('after')).otherwise(F.lit('before')))
    )
    df_ = df_.drop(*remove)
    return df_

正确实现方法 #

方法一：基于事件时间戳关联判断（适用于单事件场景） #

先提取每个ID对应的Event=1的时间戳，再通过时间比较标记行类型：

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

def label_events(df):
    # 提取每个ID的Event=1时间戳（假设每个ID仅一个Event=1行）
    event_time_df = df.filter(F.col("Event") == 1).select("ID", F.col("Timestamp").alias("event_time"))
    
    # 关联原表并判断类型
    result_df = df.join(event_time_df, on="ID", how="left") \
        .withColumn("Type", 
                    F.when(F.col("Event") == 1, F.lit("event"))
                    .when(F.col("Timestamp") < F.col("event_time"), F.lit("before"))
                    .when(F.col("Timestamp") > F.col("event_time"), F.lit("after"))
                    .otherwise(F.lit("before"))) \
        .drop("event_time")
    
    return result_df

方法二：基于窗口累计事件数判断（通用场景） #

通过窗口函数累计Event=1的数量，直接判断当前行处于事件前/事件中/事件后阶段：

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

def label_events(df):
    # 定义窗口：按ID分组、时间排序，累计当前行及之前的Event=1数量
    window_spec = Window.partitionBy("ID").orderBy("Timestamp").rowsBetween(Window.unboundedPreceding, Window.currentRow)
    
    result_df = df.withColumn("event_count", F.sum(F.col("Event")).over(window_spec)) \
        .withColumn("Type", 
                    F.when(F.col("Event") == 1, F.lit("event"))
                    .when(F.col("event_count") == 0, F.lit("before"))
                    .otherwise(F.lit("after"))) \
        .drop("event_count")
    
    return result_df

说明：

• 方法一逻辑简洁，适合每个ID仅一个Event=1的场景
• 方法二更通用，即使ID存在多个Event=1行，也会将第一个事件后的所有行标记为after（可根据需求调整累计逻辑处理多事件场景）

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