如何用Pandas/Numpy高效处理大数据框的递归循环计算

问题背景 #

有个450万行的DataFrame，需要按key_aux和Date升序排序后计算YTD_v2：

• 若当前行和上一行key_aux相同，且当前行Month≠1，则YTD_v2 = 上一行YTD_v2 + 当前行Value
• 否则YTD_v2 = 当前行Value
  原循环代码跑了13分37秒，还触发SettingWithCopyWarning，需要更高效的实现。

原代码：

df_pivot.sort_values(['key_aux', 'Date'], ascending=[True, True], inplace=True)

df_pivot['YTD_v2'] = 0

for i in range(0, len(df_pivot)):
    df_pivot['YTD_v2'].iloc[i] = np.where((df_pivot['key_aux'].iloc[i]==df_pivot['key_aux'].iloc[i-1]) & (df_pivot['Month'].iloc[i] != 1), 
                                          df_pivot['YTD_v2'].iloc[i-1]+df_pivot['Value'].iloc[i], 
                                          df_pivot['Value'].iloc[i])

运行警告：

<timed exec>:2: SettingWithCopyWarning: 
A value is trying to be set on a copy of a slice from a DataFrame

优化方案（向量化操作，秒级完成） #

直接用pandas的分组和累计求和实现，完全避免循环：

# 先排序，建议不用inplace，避免潜在的视图/副本问题
df_pivot = df_pivot.sort_values(['key_aux', 'Date'], ascending=[True, True])

# 1. 按key_aux分组，创建累计重置标记：Month=1或者是组内第一行时，重置累计
df_pivot['reset_flag'] = df_pivot.groupby('key_aux')['Month'].transform(
    lambda x: x.eq(1).cumsum()
)

# 2. 按key_aux和reset_flag分组，对Value做累计求和得到YTD_v2
df_pivot['YTD_v2'] = df_pivot.groupby(['key_aux', 'reset_flag'])['Value'].cumsum()

# 可选：删除中间生成的reset_flag列
df_pivot = df_pivot.drop(columns='reset_flag')

方案说明 #

1. 排序优化：放弃inplace=True，直接赋值给原变量，避免因视图/副本导致的SettingWithCopyWarning
2. 重置标记生成：每个key_aux组内，只要Month=1就生成一个新的累计组ID，这样后续的累计求和会在每个新组内重新开始
3. 向量化累计求和：groupby+cumsum是pandas底层优化的向量化操作，比Python循环快几十甚至上百倍，450万行数据基本几秒就能处理完

为什么原代码慢？ #

Python循环逐行处理DataFrame本质是逐元素操作，完全没有利用pandas的向量化优势，450万行的循环会触发数百万次的索引查找和赋值，效率极低。而向量化操作是批量处理数据，底层用C实现，性能差距巨大。

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