Pandas/Python中如何重写for循环提升大批量计算执行速度

你当前的代码运行慢、甚至可能出现结果不符合预期，核心是两个问题：

• 逐行for循环搭配iloc单单元格赋值是Pandas生态下性能最差的操作模式，每次索引、赋值都会产生高额的固定开销，数据量达到十万级以上时运行数小时是常态；同时x.iloc[i][列名] = 值属于链式赋值，大概率会写入临时副本而非原DataFrame，会出现计算结果不生效的问题。
• 你的计算逻辑本质是按Org_Item（单SKU维度）分组的时序递推运算，递推初始状态仅和每个SKU分组的第一行有关，不需要全局逐行判断分组边界。

以下方案按改造成本从低到高、提速幅度从小到大排序，计算逻辑和你原有规则完全对齐，不会出现结果偏差。

方案1：Numba编译递推逻辑（改造成本最低，提速100~1000倍） #

该方案不需要调整原有计算逻辑的思路，仅需将循环逻辑编译为机器码执行，绕开Pandas的所有索引开销，百万行级数据可在数秒内跑完。

import pandas as pd
import numpy as np
from numba import njit

# 必须先按SKU+日期排序，保证递推顺序正确
x = x.sort_values(by=["Org_Item", "Date"]).reset_index(drop=True)

# 提前提取计算需要的列为numpy数组，消除运行时索引开销
org_item_arr = x["Org_Item"].values
available_oh_arr = x["Available OH"].values
incoming_inv_arr = x["Incoming Inventory"].values
weekly_fcst_arr = x["Weekly Fcst"].values
target_inv_arr = x["Target Inventory Level"].values
stock_upper_arr = x["Stock Upper Limit"].values
date_arr = x["Date"].values
next_14d_arr = x["Next 14 Days"].values

# 初始化结果数组
proj_receipt_res = np.zeros(len(x), dtype=np.float64)
closing_inv_res = np.zeros(len(x), dtype=np.float64)
shortage_res = np.zeros(len(x), dtype=np.float64)
excess_res = np.zeros(len(x), dtype=np.float64)

# Numba编译递推函数，逻辑与原代码完全一致
@njit
def calc_inv_metrics(org_item, available_oh, incoming_inv, weekly_fcst, 
                     target_inv, stock_upper, date_col, next_14d,
                     proj_receipt, closing_inv, shortage, excess):
    n = len(org_item)
    for i in range(n):
        # 全局首行/新SKU分组首行，走初始值计算逻辑
        if i == 0 or org_item[i] != org_item[i-1]:
            proj_receipt[i] = 0
            temp_stock = available_oh[i] + incoming_inv[i] + proj_receipt[i] - weekly_fcst[i]
            closing_inv[i] = max(0, temp_stock)
            shortage[i] = abs(min(0, temp_stock))
            excess[i] = max(0, closing_inv[i] - stock_upper[i])
        # 同SKU内按周递推
        else:
            if date_col[i] < next_14d[i]:
                proj_receipt[i] = 0
            else:
                proj_receipt[i] = max(0, target_inv[i] + weekly_fcst[i] - closing_inv[i-1] - incoming_inv[i])
            temp_stock = closing_inv[i-1] + incoming_inv[i] + proj_receipt[i] - weekly_fcst[i]
            closing_inv[i] = max(0, temp_stock)
            shortage[i] = abs(min(0, temp_stock))
            excess[i] = max(0, closing_inv[i] - stock_upper[i])
    return proj_receipt, closing_inv, shortage, excess

# 执行计算（首次运行会有1~2秒编译耗时，后续运行无编译开销）
proj_receipt_res, closing_inv_res, shortage_res, excess_res = calc_inv_metrics(
    org_item_arr, available_oh_arr, incoming_inv_arr, weekly_fcst_arr,
    target_inv_arr, stock_upper_arr, date_arr, next_14d_arr,
    proj_receipt_res, closing_inv_res, shortage_res, excess_res
)

# 结果写回原DataFrame
x["Projected Receipt"] = proj_receipt_res
x["Closing Inventory"] = closing_inv_res
x["Expected Shortage"] = shortage_res
x["Expected Excess"] = excess_res

注意：如果日期列是Pandas内置的datetime类型，Numba可直接兼容，不需要额外做格式转换。

方案2：分组迭代+数组运算（无额外依赖，提速50~100倍） #

如果不想安装Numba依赖，可以先按Org_Item分组，组内直接操作numpy数组做递推，避免全局逐行判断分组边界、减少iloc索引开销，性能远高于原有写法，适合数据量在十万级以内的场景。

import pandas as pd
import numpy as np

# 先排序保证顺序正确
x = x.sort_values(by=["Org_Item", "Date"]).reset_index(drop=True)
# 初始化结果列
x[["Projected Receipt", "Closing Inventory", "Expected Shortage", "Expected Excess"]] = 0.0

# 按SKU分组迭代
for sku_id, group in x.groupby("Org_Item", sort=False):
    idx = group.index
    # 提取组内字段为数组
    oh = group["Available OH"].values
    inc = group["Incoming Inventory"].values
    fcst = group["Weekly Fcst"].values
    tgt = group["Target Inventory Level"].values
    upper = group["Stock Upper Limit"].values
    dt = group["Date"].values
    n14 = group["Next 14 Days"].values
    grp_len = len(idx)

    pr = np.zeros(grp_len, dtype=np.float64)
    ci = np.zeros(grp_len, dtype=np.float64)
    short = np.zeros(grp_len, dtype=np.float64)
    exc = np.zeros(grp_len, dtype=np.float64)

    # 组内首行计算
    tmp = oh[0] + inc[0] + pr[0] - fcst[0]
    ci[0] = max(0, tmp)
    short[0] = abs(min(0, tmp))
    exc[0] = max(0, ci[0] - upper[0])

    # 组内递推
    for i in range(1, grp_len):
        pr[i] = 0 if dt[i] < n14[i] else max(0, tgt[i] + fcst[i] - ci[i-1] - inc[i])
        tmp = ci[i-1] + inc[i] + pr[i] - fcst[i]
        ci[i] = max(0, tmp)
        short[i] = abs(min(0, tmp))
        exc[i] = max(0, ci[i] - upper[i])

    # 结果批量写回
    x.loc[idx, "Projected Receipt"] = pr
    x.loc[idx, "Closing Inventory"] = ci
    x.loc[idx, "Expected Shortage"] = short
    x.loc[idx, "Expected Excess"] = exc

• 不要尝试用纯apply或者全局向量化方法实现该逻辑：由于Projected Receipt和Closing Inventory是互相依赖的递推值，必须按时间顺序逐行计算，不存在纯向量化的实现方式。
• 计算前必须校验排序逻辑：必须保证同一个Org_Item下的行按日期升序排列，否则递推时取到的上一行值不属于同一SKU或时间顺序错误，会导致结果完全失效。

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