带日期间隔的两张表间滚动求和高效实现方案问询

问题概述 #

现有两张数据表counts和picks，需要以counts为基表，为每一行数据按对应group计算count_date_time当日及往前7天内，picks表中number_picked的滚动求和，最终将结果合并到counts表中。

需求示例：

• group A、时间为2022-11-11 14:00的行，求和结果为4
• group B、时间为2022-11-14 09:00的行，求和结果为16

当前问题：逐行循环合并临时表的实现方式效率极低，无法支撑实际场景中counts表20万行、picks表百万级的规模，需要高效的向量化实现方案。

复现代码 #

import pandas as pd

counts = pd.DataFrame(columns=['count_date_time', 'group', 'outcome'], data=[["November 11, 2022 2:00 PM",'A',1], ["November 14, 2022 9:00 AM",'B',0]])
counts['count_date_time'] = pd.to_datetime(counts['count_date_time'])

picks = pd.DataFrame(columns=['pick_date_time','group','number_picked'], data=[["November 1, 2022 10:00 AM","A",3],
["November 1, 2022 11:00 AM","A",7],
["November 7, 2022 2:00 PM","A",4],
["November 12, 2022 3:00 PM","A",2],
["November 2, 2022 11:00 AM","B",3],
["November 8, 2022 4:00 AM","B",2],
["November 10, 2022 6:00 PM","B",4],["November 12, 2022 6:00 PM","B",10]])
picks['pick_date_time'] = pd.to_datetime(picks['pick_date_time'])

高效解决方案 #

以下两种方案均采用向量化操作，完全避免逐行循环，可高效处理大规模数据：

方案一：使用merge_asof区间匹配+分组求和 #

利用merge_asof实现同组内的时间区间匹配，再分组计算求和结果，步骤清晰且性能优异：

import pandas as pd

# 1. 为counts添加7天窗口起始时间和唯一标识
counts['start_time'] = counts['count_date_time'] - pd.Timedelta(days=7)
counts['id'] = range(len(counts))

# 2. 对两张表按group和时间排序（merge_asof要求必须排序）
picks_sorted = picks.sort_values(['group', 'pick_date_time'])
counts_sorted = counts.sort_values(['group', 'count_date_time'])

# 3. 匹配同组内pick_date_time在[start_time, count_date_time]区间的记录
merged = pd.merge_asof(
    picks_sorted,
    counts_sorted,
    by='group',
    left_on='pick_date_time',
    right_on='count_date_time',
    direction='backward'
)
merged = merged[merged['pick_date_time'] >= merged['start_time']]

# 4. 按counts的唯一id分组求和，合并回原表
rolling_sum = merged.groupby('id')['number_picked'].sum().reset_index(name='rolling_7d_sum')
result = counts.merge(rolling_sum, on='id', how='left').fillna(0)
result = result.drop(['id', 'start_time'], axis=1)

print(result)

方案二：分组滚动求和+近邻匹配 #

先对picks按组计算7天滚动求和，再通过merge_asof匹配到counts对应的时间点：

import pandas as pd

# 1. 对picks按group和时间排序，计算分组7天滚动求和
picks_sorted = picks.sort_values(['group', 'pick_date_time'])
picks_sorted['rolling_7d_sum'] = picks_sorted.groupby('group').rolling(
    window='7D',
    on='pick_date_time',
    closed='right'
)['number_picked'].sum().reset_index(level=0, drop=True)

# 2. 对counts排序，使用merge_asof匹配同组内最近的有效滚动求和结果
counts_sorted = counts.sort_values(['group', 'count_date_time'])
result = pd.merge_asof(
    counts_sorted,
    picks_sorted[['group', 'pick_date_time', 'rolling_7d_sum']],
    by='group',
    left_on='count_date_time',
    right_on='pick_date_time',
    direction='backward'
)

# 3. 处理无匹配数据的情况，恢复原表顺序
result['rolling_7d_sum'] = result['rolling_7d_sum'].fillna(0)
result = result.drop('pick_date_time', axis=1)
result = result.set_index(counts.index).sort_index()

print(result)

两种方案最终输出结果均符合需求：

count_date_time group  outcome  rolling_7d_sum
0 2022-11-11 14:00:00     A        1              4.0
1 2022-11-14 09:00:00     B        0             16.0

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