熊猫:groupby使用numba申请

时间:2020-08-04 20:13:33

标签: python pandas pandas-groupby numba

使用pandas v1.1.0。

在pandas文档中,有一个很好的示例,说明了如何使用numba来加快rolling.apply()操作here

import pandas as pd
import numpy as np

def mad(x):
    return np.fabs(x - x.mean()).mean()

df = pd.DataFrame({"A": np.random.randn(100_000)},
                  index=pd.date_range('1/1/2000', periods=100_000, freq='T')
).cumsum()

df.rolling(10).apply(mad, engine="numba", raw=True)

我想对其进行修改以适合groupby操作:

df['day'] = df.index.day
df.groupby('day').agg(mad)

工作正常。

但是

df.groupby('day').agg(mad, engine='numba')

错误和给予

---------------------------------------------------------------------------
NumbaUtilError                            Traceback (most recent call last)
<ipython-input-21-ee23f1eec685> in <module>
----> 1 df.groupby('day').agg(mad, engine='numba')

~\AppData\Local\Continuum\anaconda3\envs\ds-cit-dev\lib\site-packages\pandas\core\groupby\generic.py in aggregate(self, func, engine, engine_kwargs, *args, **kwargs)
    939 
    940         if maybe_use_numba(engine):
--> 941             return self._python_agg_general(
    942                 func, *args, engine=engine, engine_kwargs=engine_kwargs, **kwargs
    943             )

~\AppData\Local\Continuum\anaconda3\envs\ds-cit-dev\lib\site-packages\pandas\core\groupby\groupby.py in _python_agg_general(self, func, engine, engine_kwargs, *args, **kwargs)
   1068 
   1069             if maybe_use_numba(engine):
-> 1070                 result, counts = self.grouper.agg_series(
   1071                     obj,
   1072                     func,

~\AppData\Local\Continuum\anaconda3\envs\ds-cit-dev\lib\site-packages\pandas\core\groupby\ops.py in agg_series(self, obj, func, engine, engine_kwargs, *args, **kwargs)
    623 
    624         if maybe_use_numba(engine):
--> 625             return self._aggregate_series_pure_python(
    626                 obj, func, *args, engine=engine, engine_kwargs=engine_kwargs, **kwargs
    627             )

~\AppData\Local\Continuum\anaconda3\envs\ds-cit-dev\lib\site-packages\pandas\core\groupby\ops.py in _aggregate_series_pure_python(self, obj, func, engine, engine_kwargs, *args, **kwargs)
    681 
    682         if maybe_use_numba(engine):
--> 683             numba_func, cache_key = generate_numba_func(
    684                 func, engine_kwargs, kwargs, "groupby_agg"
    685             )

~\AppData\Local\Continuum\anaconda3\envs\ds-cit-dev\lib\site-packages\pandas\core\util\numba_.py in generate_numba_func(func, engine_kwargs, kwargs, cache_key_str)
    215     nopython, nogil, parallel = get_jit_arguments(engine_kwargs)
    216     check_kwargs_and_nopython(kwargs, nopython)
--> 217     validate_udf(func)
    218     cache_key = (func, cache_key_str)
    219     numba_func = NUMBA_FUNC_CACHE.get(

~\AppData\Local\Continuum\anaconda3\envs\ds-cit-dev\lib\site-packages\pandas\core\util\numba_.py in validate_udf(func)
    177         or udf_signature[:min_number_args] != expected_args
    178     ):
--> 179         raise NumbaUtilError(
    180             f"The first {min_number_args} arguments to {func.__name__} must be "
    181             f"{expected_args}"

NumbaUtilError: The first 2 arguments to mad must be ['values', 'index']

我猜测engine=numba期望数据会略有不同。

2 个答案:

答案 0 :(得分:0)

尝试一下

df.groupby('day')。agg(mad(df.groupby(day)),engine ='numba') 不确定sure,但这是说第1个第2个参数必须为[“ values”,'index']我认为它将与数据框一起使用。

答案 1 :(得分:0)

我自己也遇到过这个问题。显然,要使用 pandas + numba 引擎,您需要f(value, index) 格式实现自定义函数。

根据文档(GroupBy.transform):

<块引用>

如果选择了“numba”引擎,则函数必须是用户定义的 以值和索引作为第一个和第二个参数的函数 分别在函数签名中。每个组的索引将是 传递给用户定义的函数并可选择使用。

我有一个简单的函数 f(x) 返回一个 int,我想在 groupby 中使用它。让它与 numba 一起工作所需要的只是将函数修改为 f(values, index),这样 numba 例程就有一个有效的参数来将索引传递给函数。

以前的功能(工作正常,但不适用于 numba):

def equal_weight(arr) -> int:
    '''
    returns a float of 1/n where 'n' is the number of rows
    '''
    return 1 / len(arr)

新功能,兼容numba引擎:

def equal_weight(values, index) -> int:
    '''
    returns a float of 1/n where 'n' is the number of rows
    '''
    return 1 / len(values)