如何高效在DataFrame上调用含循环的自定义函数？

问题分析 #

你当前用df.apply(..., raw=True)的方式可行，但行数增加时Python层逐行调用的开销会逐渐凸显。np.vectorize报错是因为它本质是逐元素处理而非逐行，fn1接收到的是单个numpy.float64值，无法迭代切片；直接用df.to_numpy()替换行操作则混淆了二维数组与单行的维度，导致后续逻辑不兼容。

优化方案 #

1. 优化现有apply调用 #

你已经用了raw=True，这是正确的——它会把行以numpy.ndarray形式传入，而非pandas.Series，减少对象转换开销。可以进一步修改fn1，直接用numpy数组切片替代list(row)，避免额外类型转换：

def fn1(row, a, b):
    # 直接用numpy数组切片，无需转list
    x = row[:a] + something  # 确保something适配numpy数组操作
    y = row[b:] + something
    result = some_evaluatefn(x, y)
    return result

# 保持原调用，确保内部逻辑兼容numpy数组
df['f'] = df.apply(lambda row: fn1(row, a, b), axis=1, raw=True)

2. NumPy+Numba循环加速 #

如果fn1内的循环无法避免，建议将DataFrame转为NumPy二维数组，用Numba将循环编译为机器码执行，大幅降低Python层开销：

import numpy as np
from numba import jit

# 转换为numpy二维数组
arr = df.to_numpy()

# Numba装饰器编译循环为机器码
@jit(nopython=True)
def accelerated_fn(arr, a, b):
    n_rows = arr.shape[0]
    results = np.empty(n_rows, dtype=np.float64)  # 提前分配结果数组
    
    for idx in range(n_rows):
        row = arr[idx]
        x = row[:a]
        # 写入你的循环操作（Numba支持多数numpy操作与基础Python循环）
        y = row[b:]
        # 写入你的循环操作
        results[idx] = some_evaluatefn(x, y)
    
    return results

# 调用并赋值
df['f'] = accelerated_fn(arr, a, b)

这种方式能将速度提升数倍至数十倍，尤其适合行数较多的场景。

3. 向量化替换部分循环 #

如果fn1内的循环逻辑可拆解为NumPy向量化操作，可直接对整个二维数组批量处理。例如，若something是标量或同维度数组，x = arr[:, :a] + something能一次性完成所有行的x计算，再结合向量化的some_evaluatefn，效率会达到最优。

关于效率的问题 #

效率不完全取决于fn1内部实现，调用方式的优化能显著降低开销：

• df.apply本质是Python层逐行循环，每次调用函数都有对象创建、参数传递的额外开销；
• NumPy/Numba的方式将循环放到更低层级（C/LLVM）执行，避免了Python解释器的 overhead；
• 但如果fn1内部循环逻辑本身非常复杂且无法向量化，核心瓶颈仍在内部循环实现——此时用Numba加速内部循环是最优解。

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