如何在Numba nopython函数中设置NumPy结构化标量整字段

你想要跳过逐元素循环，直接在Numba的nopython模式函数里批量操作结构化数组的整个字段对吧？其实Numba完全支持这种操作，和你在普通NumPy代码里的写法几乎一致，不用再写繁琐的循环了。

优化后的实现代码 #

直接修改你的desired_fn，利用Numba对矢量化操作的支持，就能实现批量更新整个字段：

import numpy as np
import numba as nb

test_numpy_dtype = np.dtype([("blah", np.int64)])
test_numba_dtype = nb.from_dtype(test_numpy_dtype)

@nb.njit
def working_fn(thing):
    for j in range(len(thing)):
        thing[j]['blah'] += j

@nb.njit
def desired_fn(thing):
    # 直接对整个字段做矢量化运算，无需循环
    thing['blah'] += np.arange(len(thing), dtype=np.int64)

验证效果 #

运行下面的测试代码，你会发现desired_fn和working_fn的结果完全一致，但前者的性能更优（尤其是数组规模较大时）：

# 创建测试数组
test_arr = np.zeros(5, dtype=test_numpy_dtype)
# 调用优化后的函数
desired_fn(test_arr)
# 查看结果
print(test_arr['blah'])  # 输出：[0 1 2 3 4]

额外说明 #

如果你的业务逻辑更复杂，只要是Numba支持的矢量化运算（比如加减乘除、广播操作等），都可以直接应用到结构化数组的字段上。需要注意的是，要确保运算中使用的数组和字段的dtype匹配，避免出现类型不兼容的报错。

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