如何在Numba中实现包含元组列表元素的优先级队列？

问题背景 #

尝试用Numba创建包含元组列表元素的PriorityQueue，参考已有实现但出现类型不匹配错误。

原实现代码 #

import typing
from heapq import heappush, heappop
import numba as nb
from numba.experimental import jitclass


itemType = nb.typed.List.empty_list(nb.types.Tuple((nb.types.int64, nb.types.int64)))
entry_def = (0.0, 0, nb.typed.List([(0,0)]))
entry_type = nb.typeof(entry_def)

@jitclass
class PriorityQueue:
    pq: typing.List[entry_type]
    id: int
    entry: entry_type

    def __init__(self):
        self.pq = nb.typed.List.empty_list((0.0, 0, nb.typed.List([(0,0)])))

    def put(self, priority: float, id: int, item: itemType):
        entry = (priority, id, item)
        heappush(self.pq, entry)

    def pop(self):
        if self.pq:
            priority, id, item = heappop(self.pq)
            return priority, id, item
        raise KeyError("pop from an empty priority queue")

期望功能 #

>>> q = PriorityQueue() 
>>> q.put(5.0, 1, [(0,1)])
>>> q.put(2.0, 2, [(0,1), (1,2)])
>>> q.put(3.0, 3, [(0,1), (0,1), (1,1)])
>>> node = q.pop()
>>> node
(2.0, 2, [(0, 1), (1, 2)])

报错信息 #

Traceback (most recent call last):
  File "C:\VS Code\myproject.py", line 34, in <module>
    q.put(5.0, 1, [(0,1)])
  File "C:\Users\Me\AppData\Local\Programs\Python\Python310\lib\site-packages\numba\experimental\jitclass\boxing.py", line 61, in wrapper
    return method(*args, **kwargs)
  File "C:\Users\Me\AppData\Local\Programs\Python\Python310\lib\site-packages\numba\core\dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "C:\Users\Me\AppData\Local\Programs\Python\Python310\lib\site-packages\numba\core\dispatcher.py", line 409, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
- Resolution failure for literal arguments:
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function heappush>) found for signature:

 >>> heappush(ListType[Tuple(float64, int64, ListType[UniTuple(int64 x 2)])], Tuple(float64, int64, reflected list(UniTuple(int64 x 2))<iv=None>))

There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload in function 'heappush': File: numba\cpython\heapq.py: Line 150.
    With argument(s): '(ListType[Tuple(float64, int64, ListType[UniTuple(int64 x 2)])], Tuple(float64, int64, reflected list(UniTuple(int64 x 2))<iv=None>))':
   Rejected as the implementation raised a specific error:
     TypingError: heap type must be the same as item type
  raised from C:\Users\Me\AppData\Local\Programs\Python\Python310\lib\site-packages\numba\cpython\heapq.py:119

During: resolving callee type: Function(<built-in function heappush>)
During: typing of call at C:\VS Code\myproject.py (24)


File "myproject.py", line 24:
    def put(self, priority: float, id: int, item: itemType):
        <source elided>
        entry = (priority, id, item)
        heappush(self.pq, entry)
        ^

- Resolution failure for non-literal arguments:
None

During: resolving callee type: BoundFunction((<class 'numba.core.types.misc.ClassInstanceType'>, 'put') for instance.jitclass.PriorityQueue#228d8eebd30<pq:ListType[Tuple(float64, int64, ListType[UniTuple(int64 x 2)])],id:int64,entry:Tuple(float64, int64, ListType[UniTuple(int64 x 2)])>)
During: typing of call at <string> (3)


File "<string>", line 3:
<source missing, REPL/exec in use?>

问题原因 #

报错核心是类型不匹配：

• 优先队列pq存储的是包含nb.typed.List的元组，但调用put时传入的是Python原生列表（reflected list），Numba无法将原生列表隐式转换为nb.typed.List，导致堆元素类型和队列声明类型不一致。
• jitclass的类型定义冗余，id和entry字段未在初始化时使用，属于无效声明。

修复方案 #

1. 严格统一类型：使用nb.typed.List而非原生列表 #

调用put时必须传入Numba的typed.List，或在方法内部自动转换原生列表。

2. 修正jitclass的类型声明 #

移除冗余字段，明确队列的元素类型。

修复后的代码 #

import numba as nb
from numba.experimental import jitclass
from heapq import heappush, heappush, heappop

# 定义元组类型：两个int64组成的元组
tuple_type = nb.types.UniTuple(nb.int64, 2)
# 定义列表类型：存储上述元组的typed.List
item_type = nb.typed.List.empty_list(tuple_type).type
# 定义队列条目类型：(优先级float64, id int64, 元组列表item_type)
entry_type = nb.types.Tuple((nb.float64, nb.int64, item_type))

# jitclass的类型规范，只保留必要的pq字段
pq_spec = [
    ('pq', nb.typed.List.empty_list(entry_type).type),
]

@jitclass(pq_spec)
class PriorityQueue:
    def __init__(self):
        # 初始化空队列，类型为entry_type的列表
        self.pq = nb.typed.List.empty_list(entry_type)

    def put(self, priority: float, id: int, item: item_type):
        entry = (priority, id, item)
        heappush(self.pq, entry)

    def pop(self):
        if self.pq:
            return heappop(self.pq)
        raise KeyError("pop from an empty priority queue")

# 使用示例：将原生列表转为nb.typed.List
if __name__ == "__main__":
    q = PriorityQueue()
    # 转换原生列表为typed.List
    item1 = nb.typed.List([(0,1)], dtype=tuple_type)
    q.put(5.0, 1, item1)
    item2 = nb.typed.List([(0,1), (1,2)], dtype=tuple_type)
    q.put(2.0, 2, item2)
    item3 = nb.typed.List([(0,1), (0,1), (1,1)], dtype=tuple_type)
    q.put(3.0, 3, item3)
    node = q.pop()
    print(node)  # 输出：(2.0, 2, [(0, 1), (1, 2)])

可选优化：在put方法内自动转换原生列表 #

如果希望直接传入原生列表，可在put内部添加转换逻辑：

# 修改后的put方法
def put(self, priority: float, id: int, item):
    # 检查是否为原生列表，转换为typed.List
    if isinstance(item, list):
        nb_item = nb.typed.List.empty_list(tuple_type)
        for tpl in item:
            nb_item.append(tpl)
        item = nb_item
    entry = (priority, id, item)
    heappush(self.pq, entry)

此时调用可直接传入原生列表：

q.put(5.0, 1, [(0,1)])

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