Pandas带DatetimeIndex的大数据集执行shift操作耗时过长解决方案问询

性能瓶颈原因 #

带freq参数的groupby.shift性能极差的核心原因是，pandas会对每个分组单独执行日期偏移、索引重排操作，当分组数量大（你的场景中col3有接近10万个唯一值）时，重复计算开销会被指数级放大，才会出现数十分钟的耗时。

优化核心思路 #

绕过groupby.shift的内置日期对齐逻辑，手动批量完成日期偏移后统一做全局索引对齐，避免每个分组的重复计算。你已经摸索出的手动偏移+索引对齐的思路完全正确，在此基础上可以做进一步优化适配多字段、多阶偏移的场景：

优化实现代码 #

首先导入依赖：

import pandas as pd
import numpy as np
from pandas.tseries.offsets import MonthEnd

适配多字段、1-6阶批量偏移的实现： #

# 预存原表的索引键，后续对齐用
base_idx = pd.MultiIndex.from_frame(df[['col3', 'col2']])

shift_result_list = []
for i in range(1, 7):
    # 批量偏移日期，不需要修改其他字段
    shifted_df = df.assign(col2 = lambda x: x['col2'] + MonthEnd(i))
    # 设置双索引、加字段后缀
    shifted_df = shifted_df.set_index(['col3', 'col2']).add_suffix(f'_shift{i}')
    shift_result_list.append(shifted_df)

# 合并所有偏移结果，对齐原表索引后重置索引
final_res = pd.concat(shift_result_list, axis=1).reindex(base_idx, fill_value=0).reset_index()

额外优化技巧（可选，可进一步提升10%-30%性能） #

• 如果你的日期已经统一为MonthEnd格式，可替换MonthEnd(i)为pd.DateOffset(months=i)，部分场景下执行速度更快
• 多字段场景下如果不需要保留所有原字段，可在assign后只筛选需要偏移的字段+col3、col2，减少后续concat的计算量
• 如果用left merge替代reindex，可提前将原表的col3、col2设为索引，直接用索引merge比字段merge效率更高

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