为含object dtype的np.array函数加Numba njit装饰器遇错求优化

能否同时用Numba装饰f和g？ #

不行。Numba的nopython模式（@njit默认启用）无法处理dtype=object的数组——这类数组存储的是Python对象，Numba无法进行静态类型推断和编译，直接装饰会触发你遇到的TypingError。

最快的映射方式（适配5亿次调用的HPC场景） #

核心思路是将A的不规则结构转换为Numba可静态编译的类型，消除Python对象开销。最优方案是填充为二维数组+单独存储子数组长度，具体实现如下：

1. 预处理数据

将原不规则列表转换为固定形状的二维数组（用0填充短数组），同时记录每个子数组的实际长度：

import numpy as np
from numba import njit

# 原始数据
A_list = [[2, 5], [4, 5, 6, 7], [0, 8], [6, 7], [1, 8], [0, 1], [1, 3], [1, 3], [2, 4]]
B = np.array([1]*9, dtype=int)

# 预处理A：生成填充后的二维数组+长度数组
max_sub_len = max(len(sub) for sub in A_list)
A_padded = np.zeros((len(A_list), max_sub_len), dtype=int)
A_lengths = np.array([len(sub) for sub in A_list], dtype=int)

for idx, sub_arr in enumerate(A_list):
    A_padded[idx, :len(sub_arr)] = sub_arr

2. 用Numba装饰优化后的函数

修改g和f，使用预处理后的静态数组：

@njit(fastmath=True)
def g(sub_len, B_len):
    # 若需处理子数组元素，可传入A_padded的行（如A_padded[i]），配合sub_len遍历有效元素
    return 19.12 / (sub_len + B_len)

@njit(fastmath=True)
def f(A_lengths, B_len):
    total = 0.0
    loop_count = len(A_lengths)
    for i in range(loop_count):
        total += g(A_lengths[i], B_len)
    return total

# 调用（预计算B的长度避免循环内重复计算）
B_len = len(B)
result = f(A_lengths, B_len)
print(result)

备选方案：Numba Typed List

如果不想做填充，可使用Numba的typed.List存储子数组，但性能略低于静态数组方案：

from numba.typed import List

# 转换为typed.List
A_typed = List()
for sub in A_list:
    A_typed.append(np.array(sub, dtype=int))

@njit(fastmath=True)
def g(a, B):
    return 19.12 / (len(a) + len(B))

@njit(fastmath=True)
def f(A, B):
    total = 0.0
    for i in range(len(B)):
        total += g(A[i], B)
    return total

result = f(A_typed, B)

性能说明 #

静态数组方案的内存连续，Numba可完全编译为机器码，几乎消除Python层开销，最适合5亿次调用的超大规模场景；Typed List方案虽更灵活，但存在一定的动态类型访问开销，仅在无法预处理时使用。

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