如何使用Numba对可变长度列表进行并行处理

你可以初始化一个预分配长度的Numba类型化列表，来支持并行循环中的索引赋值操作。Numba的Typed List允许预先指定容量并直接通过索引位置赋值，完美适配你的场景。

具体实现代码 #

from numba import njit, prange
from numba.typed import List
import numba

@njit()
def inner_function(A: numba.float32[:]) -> List[int]:
    # 示例逻辑：返回长度可变的uint32列表
    res = List()
    for val in A:
        if val > 0.5:
            res.append(int(val * 100))
    return res

@njit(parallel=True)
def looping_function(A: numba.float32[:,:]) -> List[List[int]]:
    # 初始化预分配长度的类型化列表，元素类型为List[int]
    result = List()
    # 预先填充空的类型化列表，确定结果列表的长度
    for _ in range(A.shape[0]):
        result.append(List.empty_list(numba.types.int_))
    
    # 并行循环赋值
    for i in prange(A.shape[0]):
        tmp = inner_function(A[i])
        result[i] = tmp
    
    return result

关键说明 #

• 初始化时通过循环添加空的List[int]，预先确定result的长度，这样后续就能直接用result[i]进行索引赋值，避免串行append的性能瓶颈。
• 必须使用Numba的typed List（而非Python原生列表），因为原生列表在Numba并行模式下不支持线程安全的索引赋值操作。
• 并行循环中使用prange，Numba会自动处理线程调度，由于各次inner_function调用无数据依赖，不会出现竞争问题。

额外优化点 #

如果提前知道inner_function返回列表的最大可能长度，也可以考虑用二维数组+长度记录数组的组合（比如一个uint32[:, max_len]的数组存储数据，一个int[:]数组存储每个子列表的实际长度），但这种方案需要额外处理截断逻辑，灵活性不如Typed List。

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