如何在PySpark中用lead/lag比较行值并设置FLAG列

需求拆解 #

从你的输入输出可以明确核心操作：

• 先按Year列对DataFrame排序
• 对每个重复的Year值，将该组内第二行的FLAG设为100，其余行保持None

代码实现 #

import pandas as pd

# 模拟你的原始DataFrame
df_prod = pd.DataFrame({
    'Year': [2020, 2020, 2019, 2021, 2021, 2022],
    'FLAG': [None, None, None, None, None, None]
})

# 1. 按年份排序并重置索引
df_prod = df_prod.sort_values('Year').reset_index(drop=True)

# 2. 为每个年份组标记行序号，设置对应FLAG值
df_prod['FLAG'] = df_prod.groupby('Year').cumcount().apply(lambda x: 100 if x >= 1 else None)

print(df_prod)

执行结果 #

Year   FLAG
0  2019   None
1  2020   None
2  2020    100
3  2021   None
4  2021    100
5  2022   None

代码说明 #

1. 排序处理：sort_values('Year')将数据按年份升序排列，reset_index(drop=True)避免排序后索引混乱
2. 分组标记行号：groupby('Year').cumcount()为每个年份组内的行分配从0开始的序号，比如两个2020的行序号分别为0和1
3. FLAG赋值：通过apply判断行序号，序号≥1时（即每组的第二行及以后）设置FLAG为100，否则保留None

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