PySpark保留每个ID最后一次出现'a'后的所有时间戳

解决方案代码 #

import pandas as pd

# 构造原始数据集
data = {
    'id': [1,1,1,1,1,1,2,2,2,2],
    'col1': ['a','b','c','a','b','d','a','b','a','c'],
    'timestamp': ['01.01.2022 9:00:00', '01.01.2022 9:01:00', '01.01.2022 9:02:00',
                  '01.01.2022 10:00:00', '01.01.2022 10:01:00', '01.01.2022 10:02:00',
                  '01.01.2022 12:00:00', '01.01.2022 12:01:00', '01.01.2022 13:00:00',
                  '01.01.2022 13:01:00']
}

df = pd.DataFrame(data)

# 转换timestamp为datetime类型，确保时间比较有效
df['timestamp'] = pd.to_datetime(df['timestamp'], format='%d.%m.%Y %H:%M:%S')

# 计算每个id最后一次出现col1='a'的时间
last_a_time = df[df['col1'] == 'a'].groupby('id')['timestamp'].max().reset_index(name='last_a_ts')

# 合并原始数据和last_a_time，筛选符合条件的记录
result = df.merge(last_a_time, on='id').query('timestamp >= last_a_ts').drop('last_a_ts', axis=1)

print(result)

步骤说明 #

1. 时间格式转换：将原始字符串类型的timestamp转为datetime类型，这是正确比较时间先后的前提，否则字符串的大小判断会违背实际时间逻辑。
2. 提取最后一次'a'的时间：先筛选出所有col1='a'的行，再按id分组取每组最大的timestamp，得到每个id最后一次出现'a'的时间点。
3. 筛选目标记录：把原始数据和上一步得到的时间表合并，通过query筛选出每个id中时间不早于最后一次'a'时间的记录，最后移除辅助字段last_a_ts即可得到结果。

补充说明 #

你之前考虑的groupby取最大时间戳的方法有局限性：它取的是每个id所有记录的最大时间，而非仅针对col1='a'的行。只有当最后一次'a'恰好是该id的最后一条记录时，这个方法才有效，而我们的方案专门针对'a'的出现时间做筛选，能覆盖所有场景。

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