使用字典时Numba比CPython慢25%的原因及用法排查

我针对Numba操作字典的性能做了测试，代码如下：

from numpy.random import randint
import numba as nb

@nb.njit
def foo_numba(a, b, c):
    N = 100**2
    d = {}
    for i in range(N):
        d[(randint(N), randint(N), randint(N))] = (a, b, c)
    return d

@nb.njit
def test_numba(numba_dict):
    s = 0
    for k in numba_dict:
        s += numba_dict[k][2]
    return s

def foo(a, b, c):
    N = 100**2
    d = {}
    for i in range(N):
        d[(randint(N), randint(N), randint(N))] = (a, b, c)
    return d

def test(numba_dict):
    s = 0
    for k in numba_dict:
        s += numba_dict[k][2]
    return s

a = randint(10, size=10)
b = randint(10, size=10)
c = 1.3

t_numba = foo_numba(a, b, c)
dummy = test_numba(t_numba)
%timeit test_numba(t_numba)
t = foo(a, b, c)
%timeit test(t) 

测试结果显示Numba版本反而更慢：

870 µs ± 6.36 µs per loop (mean ± std. dev. of 10 runs, 1,000 loops each)
654 µs ± 35.8 µs per loop (mean ± std. dev. of 10 runs, 1,000 loops each)

问题解析 #

为什么Numba代码更慢？ #

这不是你用错了Numba，而是Numba的JIT字典对值类型的兼容性问题：

• 当你把numpy数组作为值的一部分存入Numba字典时，Numba无法将这个值优化为原生的紧凑类型，只能以Python对象引用的方式存储。遍历字典时，每次取值都要进行Python对象的解引用和类型检查，额外开销极大。
• CPython的原生字典对Python对象的处理已经经过长期优化，在存储对象引用的场景下，其遍历效率反而超过了Numba的“原生代码+Python对象交互”的混合模式。

为什么转成元组后性能恢复？ #

当把a或b转为元组时，Numba可以将字典的值识别为纯原生类型组成的元组（整数元组+浮点数），不需要再处理Python对象引用。此时Numba能生成完全优化的原生机器码，遍历效率自然超过CPython。

优化建议 #

• 用Numba操作字典时，尽量让键和值都使用Numba支持的原生类型（如整数、浮点数、原生类型元组），避免存储numpy数组、自定义类等Python对象。
• 如果必须存储数组类数据，建议改用Numba的typed.List或者将数组数据展平为原生类型序列，不要直接放入字典。

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