使用Numba输出两列结果赋值给pandas DataFrame新增列时TypingError如何解决

报错原因 #

返回两个一维数组组成的元组时，直接赋值给Pandas多列切片会触发Numba类型推断错误，nopython模式下无法匹配Pandas的赋值类型规则。此外要注意删除测试用的无效代码（比如你示例中单输出版本里的print(a[x])，x未定义会额外触发报错）。

解决方法 #

有两种常用的修复方案，均不需要用到df.loc，性能完全满足25万行数据的计算需求：

方案1：先解包返回值再分别赋值

直接把函数返回的两个数组先赋值给临时变量，再分别写入DataFrame的两列，代码修改量最小：

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

# 构造测试数据
df = pd.DataFrame({'inputa': [1,2,3,4]})

@jit(nopython=True)
def foo(inputa):
    a = np.zeros(inputa.shape, dtype=np.int64)
    b = np.zeros(inputa.shape, dtype=np.int64)
    for i in range(1, a.shape[0]):
        a[i] = a[i-1] + inputa[i]
        b[i] = inputa[i-1] + a[i]
    return a, b

# 先解包再分别赋值
a_arr, b_arr = foo(df['inputa'].values)
df['a'] = a_arr
df['b'] = b_arr

print(df)

方案2：合并为二维数组后批量赋值

把两个一维数组合并为n行2列的二维数组，直接赋值给两列，代码更简洁：

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

# 构造测试数据
df = pd.DataFrame({'inputa': [1,2,3,4]})

@jit(nopython=True)
def foo(inputa):
    a = np.zeros(inputa.shape, dtype=np.int64)
    b = np.zeros(inputa.shape, dtype=np.int64)
    for i in range(1, a.shape[0]):
        a[i] = a[i-1] + inputa[i]
        b[i] = inputa[i-1] + a[i]
    return np.column_stack((a, b))

# 二维数组直接对应两列赋值
df[['a', 'b']] = foo(df['inputa'].values)

print(df)

两种方案运行后都可以得到预期结果：

inputa  a   b
0       1  0   0
1       2  2   3
2       3  5   7
3       4  9  12

额外优化建议 #

创建数组时显式指定dtype，可以避免Numba类型推断的额外开销，也能进一步降低类型不匹配的报错概率。25万行数据的递归计算在nopython模式下运行耗时通常低于10ms，远快于原生df.loc的实现。

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