python - 使用 nogil 处理 Cython 中的列表列表

Question

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

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

在现实中，子列表的数量是数万个。每个子列表的长度可能相差很大，从零或一个值到数百个值不等。

我想将输入数据作为一些二维结构传递给将处理它的 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. 将列表的列表放入二维 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