基于回溯窗口扩展Pandas DataFrame行的高效优化方案问询

原始数据 #

初始DataFrame结构： #

生成代码： #

import pandas as pd

df = pd.DataFrame({'date': pd.date_range(start='20090101', end='20090103')})
df['window'] = [2, 1, 3]

需求说明 #

针对每行的date字段，生成其过去window范围内的所有回溯行：

• 新增lookback字段，取值范围为0到window（包含两端）
• 新增new_date字段，通过内部方法add_biz_dt(date, -lookback)计算得到对应的工作日日期
• 最终每行原始数据会扩展为window + 1行

示例输出： #

现有低效实现 #

当前方案依赖apply生成range序列，再用explode扩展行，最后逐行apply计算日期，在数据量大时效率低下：

df['lookback'] = df['window'].apply(lambda x: range(x+1))
df = df.explode('lookback')
df['new_date'] = df[['date', 'lookback']].apply(lambda x: add_biz_dt(x[0], -x[1]), axis=1)

优化实现方案 #

通过numpy的批量操作替代逐行apply和explode，大幅提升效率：

import pandas as pd
import numpy as np

# 1. 计算每行需要扩展的行数（window + 1）
repeat_counts = df['window'] + 1

# 2. 批量重复原始行数据，替代explode
df_expanded = df.loc[df.index.repeat(repeat_counts)].reset_index(drop=True)

# 3. 批量生成lookback序列，避免逐行生成range
lookback_array = np.concatenate([np.arange(count) for count in repeat_counts])
df_expanded['lookback'] = lookback_array

# 4. 向量化处理日期计算，替代逐行apply
# 若add_biz_dt支持批量输入可直接调用，否则用np.vectorize包装
vectorized_add_biz = np.vectorize(add_biz_dt)
df_expanded['new_date'] = vectorized_add_biz(df_expanded['date'], -df_expanded['lookback'])

优化点说明： #

• 用index.repeat替代explode：底层基于numpy数组操作，比explode的Python循环实现效率更高
• 批量生成lookback：通过numpy数组拼接生成所有回溯值，避免逐行apply生成range的开销
• 向量化日期计算：用np.vectorize将单值方法转为批量处理，比apply(axis=1)的逐行处理效率提升显著

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