基于Pandas的BOM计算函数性能优化技术求助

背景 #

我是Python新手，正在开发物料清单（BOM）计算程序：从Excel订单表获取客户所需物料ID及采购数量，通过同文件内的BOM表计算所有原材料需求，减去库存后以{物料ID: 数量}的字典形式输出。

BOM表结构 #

属性说明 #

• process_id：所用工序ID
• process_No.：工序路径中的工序顺序，不一定为连续自然数，如(51,60,70,100)
• IN/OUT：标识物料为原材料还是产出品

我使用的Pandas DataFrame（命名为df_demand）新增了两列：count_demand表示所需数量，flag用于标记需处理的物料。目前程序功能已实现，但运行速度不理想，通过timeit测试发现两处耗时操作，寻求优化方案：

问题1：calculate_demand_raw函数索引查找耗时过高 #

该函数用于查找当前需求物料的原材料，并按BOM比例计算原材料需求量，代码如下：

def calculate_demand_raw(row, df_demand):
    try:
       if np.isnan(row['quantity_out']):
           raise ValueError('To avoid including those recycled materials with negative outputs')

       list_index = list(df_demand['item_id'].isin([row['item_id']]) &
                         df_demand['process_No.'].isin([row['process_No.']]) &
                         df_demand['IN/OUT'].isin(['IN']))
       index = [i for i, x in enumerate(list_index) if x==True]  
       # Search to find the index of the required generation process

       df_demand.loc[index, 'count_demand'] = row['count_demand']/row['quantity_out']*
                                               df_demand.loc[index, 'quantity_in']
       # calculate quantity of raw materials.
       df_demand.loc[index, 'flag'] = 1

    except ValueError:
        pass  # Prevent the query material is the base material, no process generation
    df_demand.loc[row.name, 'flag'] = 0

df_demand[df_demand['flag'].isin([1])].apply(lambda row: calculate_demand_raw(row, df_demand), axis=1)

timeit显示，函数中符合条件行的索引查找耗时是原材料数量计算的三倍，求缩短索引查找时间的方案。

问题2：fill_demand函数条件查询耗时 #

该函数用于将聚合后的原材料需求填充到df_demand对应工序的count_demand列，代码如下：

def fill_demand(row, qty_sum_demand, df_demand):
    df_demand[row.name, 'count_demand'] += qty_sum_demand.loc[
                                qty_sum_demand['IN/OUT'].isin([row['IN/OUT']]).tolist(),
                                'count_demand'].tolist()
    df_demand.loc[index, 'flag'] = 1

df_demand.loc[index_generated_process].apply(lambda row: 
                                fill_demand(row, qty_sum_demand, df_demand), axis=1)

是否为函数中的条件查询导致耗时？能否转为更快的Numpy向量化操作？

针对问题1：优化索引查找逻辑 #

原代码的核心低效点是逐行遍历+重复全表查询，可以从以下方向优化：

1. 替换低效索引查找方式 #

原代码中手动遍历布尔列表找索引的方式完全没必要，改用Pandas原生的掩码+索引提取，效率提升数倍：

# 替代原函数内的索引查找逻辑
mask = (df_demand['item_id'] == row['item_id']) & \
       (df_demand['process_No.'] == row['process_No.']) & \
       (df_demand['IN/OUT'] == 'IN')
index = df_demand[mask].index

直接用==替代isin([x])，避免单元素列表的冗余操作，同时利用Pandas向量化掩码操作（底层C实现）替代Python层面的遍历。

2. 提前创建复合索引 #

对df_demand按item_id、process_No.、IN/OUT创建复合索引，直接通过索引定位目标行，彻底避免全表扫描：

# 函数外提前创建复合索引
df_demand = df_demand.set_index(['item_id', 'process_No.', 'IN/OUT'])

# 函数内直接用索引定位
def calculate_demand_raw(row, df_demand):
    try:
        if np.isnan(row['quantity_out']):
            raise ValueError()
        
        target_idx = (row['item_id'], row['process_No.'], 'IN')
        if target_idx in df_demand.index:
            ratio = row['count_demand'] / row['quantity_out']
            df_demand.loc[target_idx, 'count_demand'] = ratio * df_demand.loc[target_idx, 'quantity_in']
            df_demand.loc[target_idx, 'flag'] = 1
    
    except ValueError:
        pass
    df_demand.loc[row.name, 'flag'] = 0

3. 彻底去掉逐行apply #

最有效的优化是用分组批量处理替代逐行遍历，完全规避apply(axis=1)的性能损耗：

# 筛选出需要处理的OUT行
out_rows = df_demand[(df_demand['flag'] == 1) & (~df_demand['quantity_out'].isna())]

# 按item_id+process_No.分组批量处理
for (item, proc_no), group in out_rows.groupby(['item_id', 'process_No.']):
    mask = (df_demand['item_id'] == item) & (df_demand['process_No.'] == proc_no) & (df_demand['IN/OUT'] == 'IN')
    if mask.any():
        ratio = group['count_demand'].iloc[0] / group['quantity_out'].iloc[0]
        df_demand.loc[mask, 'count_demand'] = ratio * df_demand.loc[mask, 'quantity_in']
        df_demand.loc[mask, 'flag'] = 1

# 批量标记处理完成的OUT行
df_demand.loc[out_rows.index, 'flag'] = 0

针对问题2：fill_demand函数的向量化改造 #

原函数的低效根源同样是逐行apply+重复查询，直接用向量化操作替代：

1. 聚合后批量映射 #

如果仅按IN/OUT维度聚合需求，直接生成映射字典后批量更新：

# 先聚合qty_sum_demand的需求
agg_sum = qty_sum_demand.groupby('IN/OUT')['count_demand'].sum().to_dict()

# 批量更新df_demand的count_demand
df_demand.loc[index_generated_process, 'count_demand'] += df_demand.loc[index_generated_process, 'IN/OUT'].map(agg_sum)
# 批量设置flag
df_demand.loc[index_generated_process, 'flag'] = 1

2. 多维度匹配用merge批量处理 #

如果需要多列（如item_id+IN/OUT）匹配，用merge合并聚合结果后批量计算：

# 按多维度聚合需求
qty_agg = qty_sum_demand.groupby(['IN/OUT', 'item_id'])['count_demand'].sum().reset_index()

# 合并到df_demand
df_demand = df_demand.merge(qty_agg, on=['IN/OUT', 'item_id'], how='left', suffixes=('', '_agg'))

# 批量更新count_demand
df_demand.loc[index_generated_process, 'count_demand'] += df_demand.loc[index_generated_process, 'count_demand_agg'].fillna(0)
df_demand.loc[index_generated_process, 'flag'] = 1

# 删除临时列
df_demand = df_demand.drop('count_demand_agg', axis=1)

3. 修复原代码语法错误 #

原函数存在语法错误：qty_sum_demand['IN/OUT'].isin([row['IN/OUT']).tolist()缺少右括号，且.tolist()后无法用loc正确取值，即使修复后效率依然低下，建议直接采用上述向量化方案。

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