Python中DataFrame数据更新的性能优化求助

原代码因循环内频繁修改DataFrame导致运行耗时5-8分钟，核心瓶颈是逐次loc操作带来的重复数据拷贝。以下是针对性优化方案，可将运行时间压缩至1分钟内：

核心优化思路 #

• 彻底移除循环内的DataFrame修改操作，改用批量计算+一次性更新
• 用pandas向量化运算替代Python循环，利用底层C实现加速
• 移除循环内的print(cnt)输出（IO操作会大幅拖慢循环速度）
• 优化字典映射逻辑，减少循环内的重复查找

优化后的代码 #

import pandas as pd

i = CAT
ST_DATA_TEMP = ST_DATA.copy()  # 显式复制避免原数据被意外修改
DC_STK_TEMP = DC_STK_DATA.loc[DC_STK_DATA['REV_DC_STK'] > 0].reset_index(drop=True)

# 预构建所有映射字典
gen_clr_to_merchcat = DC_STK_TEMP.set_index('GEN_CLR')['MERCHCAT'].to_dict()
comb_artreq_first = ST_DATA_TEMP.groupby('COMB')['ART_REQ'].first()
comb_artreq_sum = ST_DATA_TEMP.groupby('COMB')['ART_REQ'].sum()
gen_clr_rev_sum = DC_STK_TEMP.groupby('GEN_CLR')['REV_DC_STK'].sum()
mj_comb_sum = MJ_MBQ.groupby('COMB')['MJ_REQ'].sum()
mc_comb_sum = MC_MBQ.groupby('COMB')['MC_REQ'].sum()

# 生成所有组合的DataFrame，替代原字典循环
combinations_df = pd.DataFrame([
    {'comb': str(j) + str(k), 'j': j, 'k': k, 'l': gen_clr_to_merchcat[k]}
    for j in ST_DETAILS['ST_CD']
    for k in DC_STK_TEMP['GEN_CLR']
])

# 合并所有需要的统计值到组合表
combinations_df['ART_REQ'] = combinations_df['comb'].map(comb_artreq_sum).fillna(0)
combinations_df['rev_dc_stk'] = combinations_df['k'].map(gen_clr_rev_sum).fillna(0)
combinations_df['comb_artreq_first'] = combinations_df['comb'].map(comb_artreq_first).fillna(0)
combinations_df['mj_req'] = combinations_df.apply(lambda x: mj_comb_sum.get(str(x['j']) + str(i), 0), axis=1)
combinations_df['mc_req'] = combinations_df.apply(lambda x: mc_comb_sum.get(str(x['j']) + str(x['l']), 0), axis=1)

# 筛选符合所有条件的记录
valid_df = combinations_df[
    (combinations_df['rev_dc_stk'] >= combinations_df['ART_REQ']) &
    (combinations_df['comb_artreq_first'] > 0) &
    (combinations_df['mj_req'] >= combinations_df['ART_REQ']) &
    (combinations_df['mc_req'] >= combinations_df['ART_REQ'])
].copy()
valid_df['ALLOC'] = valid_df['ART_REQ']

# ------------------- 批量更新所有DataFrame -------------------
# 更新ST_DATA_TEMP
st_update_map = valid_df.set_index('comb')['ALLOC']
ST_DATA_TEMP['ALLOC'] = ST_DATA_TEMP['COMB'].map(st_update_map).fillna(ST_DATA_TEMP['ALLOC'])
ST_DATA_TEMP['ART_REQ'] = ST_DATA_TEMP.apply(
    lambda row: row['ART_REQ'] - st_update_map.get(row['COMB'], 0), axis=1
)

# 更新DC_STK_TEMP
dc_update_map = valid_df.groupby('k')['ALLOC'].sum()
DC_STK_TEMP['ALLOC'] = DC_STK_TEMP['GEN_CLR'].map(dc_update_map).fillna(DC_STK_TEMP['ALLOC'])
DC_STK_TEMP['REV_DC_STK'] = DC_STK_TEMP.apply(
    lambda row: row['REV_DC_STK'] - dc_update_map.get(row['GEN_CLR'], 0), axis=1
)

# 更新MJ_MBQ
mj_comb = valid_df.apply(lambda x: str(x['j']) + str(i), axis=1)
mj_update_map = pd.Series(valid_df['ALLOC'].values, index=mj_comb).groupby(level=0).sum()
MJ_MBQ['ALLOC'] = MJ_MBQ['COMB'].map(mj_update_map).fillna(MJ_MBQ['ALLOC'])
MJ_MBQ['MJ_REQ'] = MJ_MBQ.apply(
    lambda row: row['MJ_REQ'] - mj_update_map.get(row['COMB'], 0), axis=1
)

# 更新MC_MBQ
mc_comb = valid_df.apply(lambda x: str(x['j']) + str(x['l']), axis=1)
mc_update_map = pd.Series(valid_df['ALLOC'].values, index=mc_comb).groupby(level=0).sum()
MC_MBQ['ALLOC'] = MC_MBQ['COMB'].map(mc_update_map).fillna(MC_MBQ['ALLOC'])
MC_MBQ['MC_REQ'] = MC_MBQ.apply(
    lambda row: row['MC_REQ'] - mc_update_map.get(row['COMB'], 0), axis=1
)

额外性能提升建议 #

• 如果数据量极大，可考虑用Dask进行并行计算，或用NumPy数组替代部分DataFrame操作
• 确保所有用于筛选和映射的列都已设置索引（如ST_DATA_TEMP的COMB列），可进一步加速查找
• 避免在代码中使用链式索引，改用预先构建的布尔掩码或索引映射

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