jitclass模式下如何声明heapify所用嵌套列表的类型规范

问题复现 #

尝试在jitclass成员方法中调用heapq的初始实现代码：

spec = [('sese', types.ListType(types.ListType(types.int32)))]
@jitclass(spec)
class CK ():
    def __init__(self):
        self.sese = [[0,0,0,0]]
    
    def main_ (self):
        heapq.heapify (self.sese)

运行代码后抛出如下报错：

Cannot cast list(list(int64)<iv=None>)<iv=None> to ListType[ListType[int32]]:

报错原因 #

1. 类型不匹配：spec中声明sese是元素为int32列表的嵌套列表，但代码中直接赋值的Python原生列表[[0,0,0,0]]会被numba默认推断为int64类型的嵌套列表，numba在jit模式下不会做隐式的int64到int32的列表类型转换，因此抛出类型强转失败的错误。
2. 初始化方式错误：jitclass中声明为numba List类型的属性，不能直接赋值Python原生list对象，必须使用numba提供的typed List进行初始化，否则后续容器操作也会出现类型兼容问题。

正确实现代码 #

numba nopython模式下原生支持heapq.heapify、heappush、heappop等常用堆操作方法，只要严格匹配类型要求即可正常运行，修正后的代码如下：

from numba import jitclass, types, int32
from numba.typed import List
import heapq

spec = [('sese', types.ListType(types.ListType(types.int32)))]

@jitclass(spec)
class CK ():
    def __init__(self):
        # 显式构造numba typed List，严格对齐int32类型声明
        self.sese = List()
        inner_item = List([int32(0), int32(0), int32(0), int32(0)])
        self.sese.append(inner_item)
    
    def main_ (self):
        heapq.heapify(self.sese)

注意事项 #

• 所有存入sese的元素必须是int32类型的numba typed List，不能直接传入Python原生list，也不能混入int64等其他类型的数值
• 目前numba支持的heapq方法仅包括heapify、heappush、heappop、heapreplace、heappushpop，其他heapq工具方法无法在jit模式下直接调用
• 堆内元素必须是可比较的同类型对象，否则会触发比较操作的类型错误

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