pandas调用自定义函数报only integer scalar arrays错的高效实现方案问询

问题说明 #

你遇到的问题本质是pandas的apply逐行运算在数据量大的时候性能极低，而直接把整列数组传入自定义函数会因为函数内部用了只接受标量的方法（比如range()）报错，且你需要保留原有循环逻辑不做大幅修改。

最优解决方案：使用Numba JIT编译 #

Numba可以直接把Python的循环代码编译成机器码执行，完全不需要修改你原来的核心循环逻辑，性能可以提升几十到上百倍，是这类场景的首选方案。

步骤1：安装依赖 #

pip install numba

步骤2：针对简单测试函数的改写示例 #

你的原始test函数逻辑不需要修改，只需要套上Numba装饰器，同时适配数组输入即可：

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

# 你原来的生成数据代码
df4s = """
    contract  RB  BeginDate  ValIssueDate   EndDate   Valindex0
2  A00118  46   19850100      19880901  99999999          50
3  A00118  47   19000100      19880901  19831231          47
5  A00118  47   19850100      19880901  99999999          50
6  A00253  48   19000100      19820101  19811231          47
7  A00253  48   19820100      19820101  19841299          47
8  A00253  48   19850100      19820101  99999999          50
9  A00253  50   19000100      19820101  19781231          47
10 A00253  50   19790100      19820101  19841299          47
11 A00253  50   19850100      19820101  99999999          50

"""
df4 = pd.read_csv(StringIO(df4s.strip()), sep='\s+', 
                  dtype={"RB": str, "BeginDate": int, "EndDate": int,'ValIssueDate':int,'Valindex0':int})

# 改写后的test函数
@njit
def test(RB_arr):
    res = np.empty(RB_arr.shape[0], dtype=np.int64)
    for idx in range(RB_arr.shape[0]):
        RB = RB_arr[idx]
        n=1
        for i in range(RB,50):
            n+=1
        res[idx] = n
    return res

# 调用方式
rb_arr = df4['RB'].astype(int).values
df4['n'] = test(rb_arr)

运行后得到的结果和你用apply的结果完全一致，速度快得多。

步骤3：针对实际业务函数getnpx的改写示例 #

你需要先把用到的pandas列提前转换成numpy数组，再传入Numba装饰的函数即可，核心循环逻辑完全保留：

# 提前提取需要的数组，避免在numba函数里操作pandas对象（numba对pandas支持有限）
# 假设你df中有age列，以及列名为"1"到"135"的对应系数列
age_arr = df4['age'].astype(int).values
# 按顺序提取1到135的列，转成numpy数组，shape为(行数, 135)
rate_arr = df4[[str(i) for i in range(1, 136)]].values
interest = 0.03 # 替换成你的实际利率参数

@njit
def getnpx(rate_arr, age_arr, interest):
    n_rows = rate_arr.shape[0]
    res = np.empty(n_rows, dtype=np.float64)
    for row_idx in range(n_rows):
        age = age_arr[row_idx]
        val = 1.0
        initval = 1.0
        for i in range(age, 136):
            # 列是按1到135顺序排列的，所以i对应的索引是i-1
            val = val * rate_arr[row_idx, i-1]
            intval = val / (1 + interest) ** (i + 1 - age)
            initval = initval + intval
        res[row_idx] = initval
    return res

# 调用得到结果
df4['npx'] = getnpx(rate_arr, age_arr, interest)

其他说明 #

• 如果你不想引入Numba依赖，也可以用numpy的广播和累积乘积函数cumprod来实现完全向量化的运算，但需要重构循环逻辑，改造成本比Numba方案高。
• Numba方案第一次运行的时候会有编译开销，后续多次调用会直接用编译好的机器码，速度更快。

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