如何重构含列表值字典的函数以兼容numba @jit(nopython=True)

numba 0.52版本的nopython模式不支持值为Python原生list的字典对象，原实现直接使用Python内置的dict和list容器，无法通过JIT编译。

使用numba原生提供的类型化Dict和List容器重写函数，初始化容器时显式声明键、值的固定类型，即可完全兼容@jit(nopython=True)模式，重写后代码如下：

from numba import jit, types
from numba.typed import Dict, List

@jit(nopython=True)
def getIslands(labels2D, ignoreSea=True):
    # 初始化类型化字典：键为标签数组对应数据类型，值为存二元坐标元组的类型化列表
    islands = Dict.empty(
        key_type=labels2D.dtype,
        value_type=types.ListType(types.UniTuple(types.int64, 2))
    )
    width = labels2D.shape[1]
    height = labels2D.shape[0]
    for x in range(width):
        for y in range(height):
            label = labels2D[y, x]
            if ignoreSea and label == -1:
                continue
            if label in islands:
                islands[label].append((x, y))
            else:
                # 为新标签初始化坐标列表，写入第一个坐标点
                coord_list = List.empty_list(types.UniTuple(types.int64, 2))
                coord_list.append((x, y))
                islands[label] = coord_list
    return islands

• 重写后的函数逻辑和原实现完全一致，返回的类型化容器可以直接被其他nopython模式下的jit函数调用，不需要额外适配
• 坐标使用固定长度的UniTuple存储，numba对该结构的优化效率高于数组或可变列表，遍历访问性能更好
• 如果需要在非JIT的普通Python代码中使用返回结果，在函数外做一层简单转换即可，转换逻辑运行在Python层面，无额外性能损耗：

# 将numba类型化结果转为普通Python字典+列表
jit_result = getIslands(your_label_array)
py_result = {label: list(coords) for label, coords in jit_result.items()}

• 类型化列表的append操作在numba中做了专门优化，性能和Python原生list相当，不会因为动态追加元素出现明显性能下降。

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