使用Numba njit装饰器触发TypingError，如何修复该问题？

错误根源 #

Numba的@njit装饰器默认启用nopython模式，该模式对函数有严格的类型要求：

1. 所有返回路径的类型必须完全统一（包括数组维度、数据类型）。原代码中：
   • 当arr.size > 0时返回2维数组（形状为(N,2)）
   • 当arr.size == 0时返回1维数组（形状为(9,)）
     维度不匹配导致类型无法统一，触发TypingError。
2. np.array(())空数组初始化+多次np.append的动态数组操作，在Numba中支持性差，会干扰类型推断。

修复方案 #

方案1：统一返回2维数组（推荐，符合原逻辑的数组结构） #

替换低效的动态数组操作，同时保证所有分支返回同维度数组：

import numpy as np
from numba import njit

arr1 = np.array([1, 2, 3, 4])
arr2 = np.array([0, 0, 0, 0])
# arr2 = np.array([5, 6, 7, 8])
print(arr1)
print(arr2)

@njit
def test(arr1, arr2):
    # 直接构造目标2D数组，替代np.append的动态扩容
    arr = np.column_stack((arr1, arr2))
    # 执行过滤逻辑
    arr = arr[arr[:, 1] > 0]
    
    if arr.size > 0:
        return arr
    # 返回与有效结果同维度的空数组，统一返回类型
    return np.zeros((0, 2), dtype=arr1.dtype)

arr = test(arr1, arr2)
print(arr)

方案2：统一返回1维数组（若业务逻辑要求） #

将有效结果展平为1维，与默认返回的1维数组保持维度一致：

import numpy as np
from numba import njit

arr1 = np.array([1, 2, 3, 4])
arr2 = np.array([0, 0, 0, 0])
# arr2 = np.array([5, 6, 7, 8])
print(arr1)
print(arr2)

@njit
def test(arr1, arr2):
    arr = np.column_stack((arr1, arr2))
    arr = arr[arr[:, 1] > 0]
    
    if arr.size > 0:
        # 将2D数组展平为1D
        return arr.flatten()
    # 返回指定长度的1维数组，统一类型
    return np.zeros(9, dtype=arr1.dtype)

arr = test(arr1, arr2)
print(arr)

关键修复点 #

1. 替换动态数组操作：用np.column_stack直接构造目标数组，避免np.append带来的类型推断问题，同时提升代码效率。
2. 统一返回类型：确保所有分支返回的数组维度、数据类型完全一致，满足Numba nopython模式的严格类型检查。
3. 显式指定dtype：避免默认float64与输入int64的类型不一致问题，进一步强化类型统一性。

内容的提问来源于stack exchange，提问作者Аркадий Шмелев