如何在Python中高效实现带不等式连接与分组的SQL逻辑？

首先，先定义数据（将日期转换为datetime类型以便正确比较）：

import pandas as pd

dates = pd.to_datetime(['31-12-2015', '31-12-2016', '31-12-2017', '31-12-2018'], format='%d-%m-%Y')
df1 = pd.DataFrame({'id': [1,1,1,1,2,2,2,2,3,3,3,3,4,4,4,4],
                    't': dates*4,
                    'stage': [1,2,2,3,1,1,2,3,1,1,1,3,2,1,1,3]})

df2 = df1.loc[df1['stage'] == 1].copy()

你需要实现的SQL逻辑可翻译为：

从df2（即stage=1的记录）出发，左连接df1，匹配条件为相同id且df2的日期早于df1的日期；随后按df2的id和t分组，判断分组内是否存在df1中stage=2的记录，存在则flag=1，否则flag=0。

高效实现方案 #

方案1：分组预处理（推荐，大数据集更高效）

先提取每个id下所有stage=2的日期集合，再对df2的每条记录做存在性检查，避免全量笛卡尔积式的连接：

# 预处理：按id收集所有stage=2的日期
stage2_dates = df1[df1['stage'] == 2].groupby('id')['t'].agg(set).to_dict()

# 为df2添加flag字段：检查当前id下是否有stage=2的日期晚于当前t
df2['flag'] = df2.apply(lambda row: int(any(d > row['t'] for d in stage2_dates.get(row['id'], set()))), axis=1)

# 最终结果
result = df2[['id', 't', 'flag']]

方案2：模拟SQL的连接+分组逻辑

如果更贴近SQL思路，可以用merge配合分组聚合：

# 左连接df2和df1，保留id匹配且df2.t < df1.t的记录
merged = df2.merge(df1, on='id', how='left', suffixes=('_a', '_b'))
merged = merged[merged['t_a'] < merged['t_b']]

# 按id和t_a分组，判断是否存在stage_b=2的记录
result = merged.groupby(['id', 't_a']).agg(
    flag=('stage_b', lambda x: int((x == 2).any()))
).reset_index().rename(columns={'t_a': 't'})

两种方案的输出结果一致，方案1在数据量较大时性能更优，因为避免了生成大量中间连接数据。

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