PySpark需求：标记delta>10的无效行并按person生成连续性递增列

需要对DataFrame完成以下操作：

1. 当delta列值大于10时，标记该行无效，对应continuity列值为NaN
2. 按person字段分组后，每检测到一行无效行，后续所有有效行的continuity值自动递增1

输入DataFrame #

期望输出DataFrame #

实现代码 #

import pandas as pd

# 构造输入DataFrame
data = {
    'person': ['X', 'X', 'X', 'X', 'X', 'X', 'Y', 'Y', 'Y', 'Y', 'Z', 'Z', 'Z', 'Z', 'Z'],
    'delta': [2, 3, 4, 20, 50, 5, 1, 20, 2, 3, 9, 30, 2, 15, 3]
}
df = pd.DataFrame(data)

# 标记无效行（delta>10）
df['is_invalid'] = df['delta'] > 10

# 按person分组计算连续性分组值：累积无效行数量+1作为分组标识
df['continuity'] = df.groupby('person')['is_invalid'].cumsum() + 1

# 将无效行的continuity设为NaN
df.loc[df['is_invalid'], 'continuity'] = pd.NA

# 移除辅助列
df = df.drop(columns=['is_invalid'])

print(df)

代码说明 #

• 新增is_invalid辅助列快速标记无效行；
• 按person分组后对is_invalid做累积求和，每遇到一个无效行，累积值加1，再加上初始值1，得到有效行的连续性分组值；
• 最后将无效行的continuity设为NaN，并移除辅助列，得到目标结果。

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