Какой лучший способ получить структурированный массив / dataframe, как структуры в Numba?

Question

У меня есть пустой массив, на который я ссылаюсь по столбцу, например, df['x'], df['y'].

Каков наилучший способ передать это Numba, чтобы я мог запустить функцию в режиме nopython?

Или как лучше всего обращаться с фреймами данных в Numba, чтобы я мог получить доступ к столбцу по имени?

Answer 1

Предоставить 1d массивы в качестве аргументов

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)

Распаковать 2d массив в функции 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)