使用PySpark添加lead列：获取Now列后续值并匹配行数

现有如下数据集：

Sl Now
    1   D
    2   D
    3   D
    4   R
    5   R
    6   C
    7   C
    8   C
    9   D
    10  P
    11  R
    12  R
    13  D

期望添加名为lead的列，该列显示Now列的下一个不同值，同一连续相同值组的所有行都填充这个后续值，最后一行无后续值则为空，效果如下：

Sl Now   lead
    1   D   R
    2   D   R
    3   D   R
    4   R   C
    5   R   C
    6   C   D
    7   C   D
    8   C   D
    9   D   P
    10  P   R
    11  R   D
    12  R   D
    13  D   

可以通过PySpark实现，核心思路是先对连续相同的Now值分组，再获取每个分组的下一个分组值，最后关联回原数据填充。具体实现步骤和代码如下：

实现步骤 #

1. 生成分组ID：使用窗口函数，对比当前行与前一行的Now值，标记分组边界，生成唯一分组ID。
2. 获取分组的后续值：提取每个分组的ID和对应Now值，用lead函数获取该分组的下一个分组Now值。
3. 关联填充lead列：将分组的后续值关联回原数据集，为每个分组的所有行填充对应后续值。

PySpark代码示例 #

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

# 初始化SparkSession
spark = SparkSession.builder.appName("LeadColumnDemo").getOrCreate()

# 创建原始数据
data = [
    (1, "D"), (2, "D"), (3, "D"),
    (4, "R"), (5, "R"),
    (6, "C"), (7, "C"), (8, "C"),
    (9, "D"), (10, "P"),
    (11, "R"), (12, "R"), (13, "D")
]
df = spark.createDataFrame(data, ["Sl", "Now"])

# 步骤1：生成分组ID
window_group = Window.orderBy("Sl")
df_with_group = df.withColumn(
    "group_id",
    F.sum(F.when(F.lag("Now").over(window_group) != F.col("Now"), 1).otherwise(0)).over(window_group)
)

# 步骤2：获取每个分组的后续值
group_next_value = df_with_group.groupBy("group_id").agg(F.first("Now").alias("current_value"))
window_next = Window.orderBy("group_id")
group_next_value = group_next_value.withColumn(
    "lead",
    F.lead("current_value").over(window_next)
)

# 步骤3：关联回原数据，填充lead列
result_df = df_with_group.join(group_next_value, on="group_id", how="left").drop("group_id", "current_value")

# 展示结果
result_df.orderBy("Sl").show()

代码说明 #

• 生成group_id时，用lag函数获取前一行的Now值，与当前行对比，不同则累加1，最终得到的分组ID能确保连续相同Now值的行属于同一组。
• 通过分组聚合拿到每个分组的Now值，再用lead函数获取下一个分组的Now值，也就是当前分组的后续值。
• 最后将后续值关联回原数据，每个分组的所有行都会填充对应的后续值，最后一个分组没有后续值，lead列自动为空。

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