使用Numba处理NumPy结构化数组按条件选列时遇unicode_type索引错误

尝试从NumPy结构化数组中根据传入函数的条件选择列，无条件判断时代码正常运行，但加入条件判断后Numba抛出类型错误。

可正常运行的代码 #

import numpy as np
import numba

@numba.njit(fastmath=True, cache=True)
def fun(a, b=0):
    c = 'name'
    #if b:
    #    c = 'age'
    return a[c]

a = np.array([('Rex', 9, 81.0), ('Fido', 3, 27.0)],
             dtype=[('name', 'U10'), ('age', 'i4'), ('weight', 'f4')])
             
fun(a)

无法运行的代码 #

import numpy as np
import numba

@numba.njit(fastmath=True, cache=True)
def fun(a, b=0):
    c = 'name'
    if b:
        c = 'age'
    return a[c]

a = np.array([('Rex', 9, 81.0), ('Fido', 3, 27.0)],
             dtype=[('name', 'U10'), ('age', 'i4'), ('weight', 'f4')])
             
fun(a)

错误信息 #

---------------------------------------------------------------------------
TypingError                               Traceback (most recent call last)
c:\Users\User\Workspaces\temp.ipynb Cell 9 in ()
      8     return a[c]
     10 a = np.array([('Rex', 9, 81.0), ('Fido', 3, 27.0)],
     11              dtype=[('name', 'U10'), ('age', 'i4'), ('weight', 'f4')])
---> 13 fun(a)

File ~\AppData\Roaming\Python\Python38\site-packages\numba\core\dispatcher.py:468, in _DispatcherBase._compile_for_args(self, *args, **kws)
    464         msg = (f"{str(e).rstrip()} \n\nThis error may have been caused "
    465                f"by the following argument(s):\n{args_str}\n")
    466         e.patch_message(msg)
---> 468     error_rewrite(e, 'typing')
    469 except errors.UnsupportedError as e:
    470     # Something unsupported is present in the user code, add help info
    471     error_rewrite(e, 'unsupported_error')

File ~\AppData\Roaming\Python\Python38\site-packages\numba\core\dispatcher.py:409, in _DispatcherBase._compile_for_args..error_rewrite(e, issue_type)
    407     raise e
    408 else:
---> 409     raise e.with_traceback(None)

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function() found for signature:
 
 >>> getitem(unaligned array(Record(name[type=[unichr x 10];offset=0],age[type=int32;offset=40],weight[type=float32;offset=44];48;False), 1d, C), unicode_type)
 
There are 22 candidate implementations:
      - Of which 20 did not match due to:
      Overload of function 'getitem': File: : Line N/A.
        With argument(s): '(unaligned array(Record(name[type=[unichr x 10];offset=0],age[type=int32;offset=40],weight[type=float32;offset=44];48;False), 1d, C), unicode_type)':
       No match.
      - Of which 2 did not match due to:
      Overload in function 'GetItemBuffer.generic': File: numba\core\typing\arraydecl.py: Line 166.
        With argument(s): '(unaligned array(Record(name[type=[unichr x 10];offset=0],age[type=int32;offset=40],weight[type=float32;offset=44];48;False), 1d, C), unicode_type)':
       Rejected as the implementation raised a specific error:
         NumbaTypeError: unsupported array index type unicode_type in [unicode_type]
  raised from C:\Users\User\AppData\Roaming\Python\Python38\site-packages\numba\core\typing\arraydecl.py:72

During: typing of intrinsic-call at C:\Users\User\AppData\Local\Temp\ipykernel_37200\1621110578.py (8)

File "..\..\..\..\..\AppData\Local\Temp\ipykernel_37200\1621110578.py", line 8:

原因分析 #

Numba的nopython模式需要在编译阶段确定所有变量的类型。当用条件判断动态赋值字符串变量c时，Numba会将c推断为unicode_type，但它不支持用动态字符串索引结构化数组——只有编译时就能确定的常量字符串才能用于索引结构化数组的列，这就是注释掉条件判断后代码能运行的原因。

解决方法 #

方法1：分支直接返回对应列 #

在条件分支中分别返回指定列，让Numba编译时明确每个分支的返回类型：

import numpy as np
import numba

@numba.njit(fastmath=True, cache=True)
def fun(a, b=0):
    if b:
        return a['age']
    else:
        return a['name']

a = np.array([('Rex', 9, 81.0), ('Fido', 3, 27.0)],
             dtype=[('name', 'U10'), ('age', 'i4'), ('weight', 'f4')])
             
fun(a)
fun(a, 1)

方法2：使用列索引位置替代名称 #

结构化数组的列可通过整数索引访问，先将列名映射为索引，再用条件判断选择索引：

import numpy as np
import numba

@numba.njit(fastmath=True, cache=True)
def fun(a, b=0):
    # 提前映射列索引：name是0，age是1
    col_idx = 0
    if b:
        col_idx = 1
    return a[:, col_idx]

a = np.array([('Rex', 9, 81.0), ('Fido', 3, 27.0)],
             dtype=[('name', 'U10'), ('age', 'i4'), ('weight', 'f4')])
             
fun(a)
fun(a, 1)

方法3：显式指定返回类型（仅限返回类型一致场景） #

如果待选列的类型相同，可通过numba.njit的return_type参数显式指定返回类型，但此方法不适用于返回类型不同的场景（比如示例中name是字符串、age是整数）：

# 仅适用于返回列类型一致的情况
import numpy as np
import numba

@numba.njit(numba.int32[:](numba.types.Record, numba.int64), fastmath=True, cache=True)
def fun(a, b=0):
    c = 'age' if b else 'age' # 示例仅返回整数类型列
    return a[c]

内容的提问来源于stack exchange，提问作者D.Manasreh