python - 在 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)