使用 nogil 在 Cython 中处理列表列表

Question

在 Python 中，我有一个列表列表作为输入：

input = [[0,1,2],[0,3,4,5],[0,6]]

实际上，子列表的数量是数万个。每个子列表的长度可以变化很大，从零或一个值到数百个。

我想将输入数据作为一些 2D 结构传递给将处理它的 Cython 模块。我希望在多核上处理数据，因此我使用prange 和nogil=True：

from cython.parallel import prange

cpdef int my_func(long[:,:] arr):
    cdef int i,j
    for i in prange(arr.shape[0], nogil=True):
        for j in range(arr.shape[1]):
                    # Do something
            pass
    return 42

我看到以下解决方案：

1. 将列表列表放入 2D ndarray。但是由于每个子列表的长度变化很大，ndarray并不是一个理想的数据结构
2. 修改my_func 以接受列表列表。问题是部分代码在没有 GIL 的情况下执行，因此无法访问 python 对象。

有没有人有关于如何解决这个问题的建议，最好是代码？

Answer 1

我可能会选择扁平数组，其中单个列表的开头存储在辅助数组中，与csr-matrices 相似。

这是一个如何从列表列表构造此数据结构的示例（使用 numpy，但您也可以使用 array.array；它也没有真正针对速度进行优化），只是为了给您一个想法：

import numpy as np
def flatten_list_of_lists(lst_of_lsts):
    N = sum(map(len, lst_of_lsts))  # number of elements in the flattened array   
    starts = np.empty(len(lst_of_lsts)+1, dtype=np.uint64)  # needs place for one sentinel
    values = np.empty(N, dtype=np.int64)

    starts[0], cnt = 0, 0
    for i,lst in enumerate(lst_of_lsts):
        for el in lst:
            values[cnt] = el
            cnt += 1       # update index in the flattened array for the next element
        starts[i+1] = cnt  # remember the start of the next list

    return starts, values

因此，对于您的示例，产生以下结果：

#         starts                                 values
(array([0, 3, 7, 9], dtype=uint64), array([0, 1, 2, 0, 3, 4, 5, 0, 6]))

可以看到：有 3 个子列表，从 0、3 和 7 开始，分别为 3（区别 starts[1]-starts[0]）、4 和 2 元素长。

这就是这些数据的使用方式：

%%cython
from cython.parallel import prange

cpdef int my_func(unsigned long long[::1] starts, long long[::1] values):
    cdef int i,j
    for i in prange(len(starts)-1, nogil=True):
        for j in range(starts[i], starts[i+1]):
                    # Do something with values[j]
            pass
    return 42