将自定义numba njit函数应用于熊猫滚动对象

Question

在pandas 1.0.0版本中，.apply现在具有ability以使用numba jit函数。

使用.apply时是否可以利用此.rolling功能？

例如，我有一个pandas.DataFrame和一个numba_mean函数，我希望将其应用于3个周期的滚动平均值。

import pandas as pd
import numpy as np
import numba as nb

df = pd.DataFrame({"A" : np.random.rand(10)})
@nb.jit 
def numba_mean(x): 
    return np.sum(x) / len(x) 

df.A.rolling(3).apply(numba_mean)

但是我得到以下错误，这意味着它不适用于nopython模式。

启用了循环提升功能后，编译回落到对象模式，原因是函数“ numba_mean”由于以下原因而导致类型推断失败：

将numba_mean应用于njit而不是jit时的错误回溯

此错误可能是由以下参数引起的： -参数0：无法确定<class 'pandas.core.series.Series'>

的Numba类型

我想知道是否有一种方法可以将 numbafied 函数与pandas rolling对象一起使用？

Answer 1

您需要指定engine keyword以便让熊猫知道您要使用Numba：

df.A.rolling(3).apply(numba_mean, engine='numba', raw=True)

Pandas可以jit为您提供该功能，但是当我自己执行此操作时，可以得到更快的结果。也许Numba会针对每次对.apply()的调用进行重新编译，并且以下时间包括编译时间。

具有如此简单的功能，并且每次调用的数据量很少（窗口= 3），因此不太可能通过内置函数获得任何加速。

import pandas as pd
import numpy as np
import numba as nb

@nb.njit(nogil=True)
def numba_mean(x):
    return np.sum(x) / x.size

def numpy_mean(x): 
    return np.sum(x) / x.size

df = pd.DataFrame({"A" : np.random.rand(10000)})