Pandas按列规则批量生成新列 如何优化计算性能？

问题背景 #

• 计算需求：针对DataFrame df，将a列分别和b1_{`b6`共6个b系列列做规则匹配，生成`a_b1`}a_b66个新列。赋值规则为：两列值均大于0时新列取1，均小于0时取-1，其余情况取0。
• 现有实现：提取a列numpy数组后，逐列循环每个b系列列，用双层嵌套np.where计算结果再赋值回原DataFrame。
• 性能要求：业务函数get_result累计需运行400万次，当前单次耗时1.56ms左右，需要进一步提速。

可落地的优化手段 #

1. 用二维数组广播替代逐列循环，消除循环冗余开销 #

原实现逐列提取b列、逐列调用np.where会产生大量重复的类型判断、数组构造开销。直接把6个b列一次性取为二维numpy数组，利用numpy广播机制一次性完成所有列的计算，最后统一赋值回DataFrame，可直接将单次耗时压到0.3ms以内。
核心实现代码：

# 提前把固定列名定义为常量，不要在函数内重复生成，减少字符串开销
B_COLS = ["b1", "b2", "b3", "b4", "b5", "b6"]
NEW_COLS = [f"a_{c}" for c in B_COLS]

def get_result(df):
    # 取a列并升维为(行数,1)，适配和(行数,6)的b数组广播
    a_arr = df["a"].to_numpy(copy=False)[:, None]
    # 一次性取所有b列组成二维数组
    b_arr = df[B_COLS].to_numpy(copy=False)
    # 用int8类型存结果，仅占1字节，比默认int64内存效率高8倍
    res = np.zeros(b_arr.shape, dtype=np.int8)
    # 批量打标，无逐列循环
    res[(a_arr > 0) & (b_arr > 0)] = 1
    res[(a_arr < 0) & (b_arr < 0)] = -1
    # 一次性赋值所有新列，减少pandas多次赋值的元数据操作开销
    df[NEW_COLS] = res
    return df

2. 剥离非必要的DataFrame操作 #

如果下游逻辑不需要保留DataFrame结构，直接返回计算得到的numpy数组res，可以砍掉pandas赋值时的列对齐、类型检查、索引对齐开销，性能还能再提升40%左右。如果必须返回DataFrame，不要在函数内做动态列名生成、列筛选这类固定操作，全部提前到函数外作为常量传入。

3. 极致性能场景用numba编译计算逻辑 #

如果还需要进一步压缩耗时，用numba把核心计算逻辑编译为机器码，首次运行编译完成后，单次计算耗时可以压到0.1ms以内。注意jit函数内只处理numpy数组，不要传入pandas对象：

from numba import njit

# 开启缓存，避免每次启动程序重新编译
@njit(cache=True, fastmath=True)
def _calc_core(a_arr, b_arr):
    n_rows, n_bcols = b_arr.shape
    res = np.zeros((n_rows, n_bcols), dtype=np.int8)
    for i in range(n_rows):
        av = a_arr[i]
        for j in range(n_bcols):
            bv = b_arr[i, j]
            if av > 0 and bv > 0:
                res[i, j] = 1
            elif av < 0 and bv < 0:
                res[i, j] = -1
    return res

调用时只需要提前把a、b的numpy数组取出传入即可，编译后的循环效率比numpy广播更高，尤其适配行数较多的DataFrame场景。

避坑提示 #

• 不要用.values取numpy数组，优先用to_numpy(copy=False)，避免不必要的内存拷贝
• 结果数组固定用np.int8类型，三个取值完全覆盖需求，更小的内存占用意味着更快的读写速度
• 不要在高频运行的函数里写动态逻辑，比如循环生成列名、动态判断列是否存在，这类操作占比会随着计算逻辑提速变得越来越明显

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