Pandas统计两个DataFrame中每个用户对应时间点前的历史事件数量

核心痛点 #

原实现使用apply逐行遍历用户表，每次遍历都全量过滤事件表，时间复杂度为O(用户表行数 * 事件表行数)，80万用户表对应20万事件表的场景下自然耗时极长。

优化方案 #

使用pandas向量化操作替代逐行遍历，时间复杂度仅为O(M log M + N log N)（M、N分别为两张表的行数），相同数据量下运行耗时可压缩到10秒以内。

推荐优先使用pd.merge_asof方案，这是pandas专门为有序键匹配设计的API，天生适配“匹配小于等于当前时间的对应记录”的场景：

import pandas as pd

def count_events_optimized(df_clients: pd.DataFrame, df_event: pd.DataFrame, col_event_date: str = 'DATE_EVENT',
                 col_start_date: str = 'START_DATE', col_event:str = 'nb_event'):
    # 两张表都按用户ID+日期排序，merge_asof要求排序
    df_clients_sorted = df_clients.sort_values(by=['ID_CLIENT', col_start_date]).reset_index(drop=True)
    df_event_sorted = df_event.sort_values(by=['ID_CLIENT', col_event_date]).reset_index(drop=True)
    # 事件表新增计数标记
    df_event_sorted['cnt'] = 1
    # 按用户分组，匹配所有小于等于START_DATE的事件
    merged = pd.merge_asof(
        df_clients_sorted,
        df_event_sorted,
        left_on=col_start_date,
        right_on=col_event_date,
        by='ID_CLIENT'
    )
    # 分组统计每个时间节点的事件总数
    res = merged.groupby(['ID_CLIENT', col_start_date], as_index=False)['cnt'].sum().rename(columns={'cnt': f'{col_event}_tot'})
    # 无匹配事件的场景填充0
    res[f'{col_event}_tot'] = res[f'{col_event}_tot'].fillna(0).astype(int)
    return res

# 测试用例
df_users = pd.DataFrame(data={
    'ID_CLIENT': ['A', 'A', 'A', 'B'],
    'START_DATE': ['2015-12-31', '2016-12-31', '2017-12-31', '2016-12-31'],
})
df_users["START_DATE"] = pd.to_datetime(df_users["START_DATE"])

df_events = pd.DataFrame(data={
    'ID_CLIENT': ['A', 'A', 'A', 'A', 'B'],
    'DATE_EVENT': ['2017-01-01', '2017-05-01', '2018-02-01', '2016-05-02', '2015-01-01']
})
df_events["DATE_EVENT"] = pd.to_datetime(df_events["DATE_EVENT"])

tmp = count_events_optimized(df_users, df_events)
print(tmp)

运行输出完全符合预期：

ID_CLIENT START_DATE  nb_event_tot
0         A 2015-12-31             0
1         A 2016-12-31             1
2         A 2017-12-31             3
3         B 2016-12-31             1

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