Cython并行二分搜索优化千万级DataFrame过滤问题求助

已通过Cython化实现函数替代千万级行（100,000,000行）pandas.DataFrame的过滤操作，现需向量化该操作并利用并行处理加速。

原Python版本通过对比df列（Z、T、M）与单一值（Z1、Z2、T1、T2、M1、M2）实现过滤：

# Python
# df: pandas.DataFrame
# Z, T, M: df的列，float64类型
# I: df的列，uint32类型
# Z1, Z2, T1, T2, M1, M2: float64类型

res = df.query(
    f"{Z1} < Z < {Z2} and {T1} < T < {T2} and {M1} < M < {M2}",
    engine="numexpr",
)[["T", "I", "M"]]

由@Nick ODell提供的Cython解决方案：

import numpy as np
import pandas as pd
cimport numpy as cnp
cimport cython
cnp.import_array()

@cython.boundscheck(False)
@cython.wraparound(False)
cpdef filter_df_cy_no_pd(df, double Z1, double Z2, double T1, double T2, double M1, double M2):
    # 获取各列的NumPy数组
    cdef cnp.float64_t[::1] Z = df['Z'].values
    cdef cnp.float64_t[::1] T = df['T'].values
    cdef cnp.float64_t[::1] M = df['M'].values
    cdef cnp.int32_t[::1] I = df['I'].values

    cdef cnp.float64_t[::1] T_out = np.empty_like(T)
    cdef cnp.int32_t[::1] I_out = np.empty_like(I)
    cdef cnp.float64_t[::1] M_out = np.empty_like(M)

    cdef Py_ssize_t n = Z.shape[0]
    cdef cython.bint in_range = 0
    cdef Py_ssize_t out_idx = 0
    for i in range(n):

        in_range = (
            (Z1 < Z[i]) & (Z[i] < Z2) &
            (T1 < T[i]) & (T[i] < T2) &
            (M1 < M[i]) & (M[i] < M2)
        )
        if in_range:
            T_out[out_idx] = T[i]
            I_out[out_idx] = I[i]
            M_out[out_idx] = M[i]
            out_idx += 1

    T_out_arr = np.asarray(T_out)[:out_idx]
    I_out_arr = np.asarray(I_out)[:out_idx]
    M_out_arr = np.asarray(M_out)[:out_idx]

    return pd.DataFrame({
        'T': T_out_arr,
        'I': I_out_arr,
        'M': M_out_arr,
    }, copy=False)

尝试并行化该函数的操作均以编译失败告终。

尝试方案 #

• 将df参数替换为：

  cnp.uint32_t[::1] I
  cnp.float64_t[::1] Z
  cnp.float64_t[::1] T
  cnp.float64_t[::1] M
  

• 将Z1、Z2、T1、T2、M1、M2参数改为cnp.float64_t[::1]类型，这些参数由Python中创建的pandas.Series通过.to_numpy()转为numpy.ndarray传入。
• 将Cython函数定义从cpdef改为cdef，并添加noexcept nogil。
• 将for i in range(n)改为for i in prange(n, nogil=True)，同时在module.pyx顶部添加from cython.parallel import prange。
• 函数返回T_out_arr, I_out_arr, M_out_arr。
• 在外部cpdef函数的上下文管理器中调用该函数：

  cdef Py_ssize_t p, r
  r = Z1.shape[0]
  cdef object[:, :] output = np.empty((r, 3), dtype=object)
  # `3` 对应内部cdef函数filter_cy()返回的3个数组
  
  with nogil:
      for p in range(r):
          output_nested[p, :] = filter_cy(
              # 参数...
          )
  

报错信息 #

遇到的部分报错如下：

• Function with Python return type cannot be declared nogil
• Calling gil-requiring function not allowed without gil
• Accessing Python attribute not allowed without gil
• Cannot read reduction variable in loop body

明白多数问题源于内部cdef函数中存在Python对象，但难以找到替代方案。不确定是否应并行化调用内部函数的外部循环，但理想状态下需加速作为性能瓶颈的内部过滤操作。

此外，尚未着手在内部函数中实现二分搜索。思路是先在Python函数中对每个numpy.ndarray排序，再在内部函数中获取匹配索引，取索引交集后在外部函数中提取对应值，但不确定这是否为最优方案，也不清楚具体实现方式。

恳请各位提供帮助。

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