使用pandas v1.1.0。
在 pandas 文档中,有一个很好的示例,介绍如何使用 numba 来加速
rolling.apply()
操作这里
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
预计数据会略有不同。
您尝试过 Bodo 吗?它构建在 Numba 之上并直接支持 Pandas。例如:
pip install bodo
import pandas as pd
import numpy as np
import bodo
def mad(x):
return np.fabs(x - x.mean()).mean()
np.random.seed(0)
df = pd.DataFrame({"A": np.random.randn(100_000)},
index=pd.date_range('1/1/2000', periods=100_000, freq='T')
).cumsum()
@bodo.jit(distributed=False)
def f(df):
df['day'] = df.index.day
df2 = df.groupby('day').agg(mad)
return df2
df2 = f(df)
print(df2)
这个示例似乎太小,无法从编译中受益,但这可能对您的实际用例有所帮助。
我自己也遇到过这个问题。显然,要使用 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)
df.groupby('day').agg(mad(df.groupby(day)), engine='numba')
不确定 😕 但它说第 2 个参数必须是 ["values", 'index']
我认为它可以与数据框一起使用。