在 Numba 中获取结构化数组/数据帧之类的结构的最佳方法是什么？

Question

我有一个按列引用的 numpy 数组，例如，df['x']、df['y']。

将它提供给 Numba 以便我可以在 nopython 模式下运行该功能的最佳方法是什么？

或者在 Numba 中处理数据框的最佳方法是什么，以便我可以按名称访问列？

Answer 1

提供一维数组作为参数

numba 旨在直接使用 NumPy 数组。因此，您不应该将数据帧或结构化数组提供给 numba 函数。您可以将数组作为单独的参数提供。例如：

from numba import njit

@njit
def func(A, B):
    # some logic
    arr = A + B
    return arr

df['z'] = func(df['x'].values, df['y'].values)

在numba 函数中解压二维数组

这是一种特殊情况，您的数据框系列都具有相同的类型。如果您不确定您的系列类型，请查看df.dtypes。您可以提供单个数组并在 numba 内执行解包：

@njit
def func(df_values):
    A, B = df_values[:, 0], df_values[:, 1]
    # some logic
    arr = A + B
    return arr

df['z'] = func(df[['x', 'y']].values)