PySpark如何为排序数据集添加前序后序事件及序列id列

你猜测的方向是正确的，用窗口函数的lag、lead结合row_number就能完全实现需求，以下是可直接运行的实现代码：

前置依赖导入 #

from pyspark.sql.functions import lag, lead, row_number
from pyspark.sql.window import Window

步骤1：定义窗口规则 #

我们需要两类窗口实现不同的计算逻辑：

• 会话级窗口：用于同一会话内的前后行数据提取，按会话ID分区、时间戳排序
• 全局排序窗口：用于生成全局连续的唯一ID，保持和原始数据集的排序一致

# 同会话计算窗口
session_window = Window.partitionBy("session_id").orderBy("timestamp")
# 全局排序窗口（按原始数据的排序规则session_id+timestamp排序）
global_window = Window.orderBy("session_id", "timestamp")

步骤2：实现需求字段计算 #

# 假设你的原始数据集变量名为df
result_df = df \
    # 可选需求：生成全局从1开始的唯一ID
    .withColumn("uniq_id", row_number().over(global_window)) \
    # 需求1：同会话当前行的下一条event值，无后续行返回null
    .withColumn("next_event", lead("event", 1).over(session_window)) \
    # 需求2：同会话当前行的上一条event值，无前行返回null
    .withColumn("prev_event", lag("event", 1).over(session_window)) \
    # 可选需求：同会话当前行的下一条uniq_id值，无后续行返回null
    .withColumn("next_id", lead("uniq_id", 1).over(session_window))

补充说明 #

默认lag、lead取不到对应值时会自动填充null，如果你需要显式填充字符串NA，可以给函数增加第三个参数指定填充值：

# 示例：显式填充NA字符串
.withColumn("next_event", lead("event", 1, "NA").over(session_window))

执行result_df.show()即可得到你需要的目标输出格式。

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