加速自定义聚合函数

Posted 2023-03-11

【中文标题】加速自定义聚合函数

【英文标题】：Speed up custom aggregation functions

【发布时间】：2016-01-30 22:11:39

【问题描述】：

我有一个非常简单的设置：pandas 数据框 df 中的市场数据（报价），如下所示：

index period ask bid
00:00:00.126 42125 112.118 112.117
00:00:00.228 42125 112.120 112.117
00:00:00.329 42125 112.121 112.120
00:00:00.380 42125 112.123 112.120
00:00:00.432 42125 112.124 112.121
00:00:00.535 41126 112.124 112.121
00:00:00.586 41126 112.122 112.121
00:00:00.687 41126 112.124 112.121
00:00:01.198 41126 112.124 112.120
00:00:01.737 41126 112.124 112.121
00:00:02.243 41126 112.123 112.121

现在我使用 pandas.groupy 来汇总期间

g=df.groupby('period')

按时间段很容易获得最低和最高价格，例如

import numpy as np
res=g.agg('ask': [np.amax, np.amin])

这也相当快。现在，我还想要每个时期的第一个和最后一个价格。这就是麻烦的开始。当然可以：

res=g.agg('ask': lambda x: x[0])

它基本上可以工作，但是对于大型数据集来说它非常慢。基本上，Python 函数调用的调用开销非常巨大。

有谁知道类似于 np.amax 的 numpy 函数，它将返回组的第一个或最后一个元素？我找不到一个。 iloc[0] 不能解决问题，因为它是一个对象的方法，因此，我不能将它作为函数传递给 g.agg，因为我在这个阶段没有对象（这就是需要 lambda)。

现在，我并不懒惰，我尝试使用 cython 为自己做这件事。

import numpy as np
cimport numpy as np

cpdef double first(np.ndarray array_series):
    return array_series[0]

但是 pandas 不会接受它作为聚合函数，因为它传递的是 pd.core.series-object 而不是 np.ndarray。 （没关系，一个派生自另一个，编译器不承认这一点。）

有谁知道如何编写一个 cython 函数来接受 pandas 系列而无需 python 调用开销？

【问题讨论】：

你试过df.groupby('period').first()和df.groupby('period').last()吗？

谢谢，这是一个很好的提示。它有效，但我不能将 first() 函数传递给 g.agg(...)，可以吗？我希望这样更好，因为我想一次应用许多不同的聚合函数（amin、amax、first、...）。我猜，使用它然后手动组装我的最终数据集将是一种解决方法。

可以，请看我的回答

【参考方案1】：

IIUC 那么你可以做first 和last:

In [270]:
g=df.groupby('period')
res=g.agg('ask': [np.amax, np.amin, 'first', 'last'])
res

Out[270]:
            ask                           
           amax     amin    first     last
period                                    
41126   112.124  112.122  112.124  112.123
42125   112.124  112.118  112.118  112.124

【讨论】：

你真棒。谢谢！！

但是更进一步。如果我想要一个非标准的自定义聚合函数怎么办。我可以对它进行cythonize吗？

可能，但我不是 cython 专家

使用自定义函数的方法是对其进行cythonize/numbize，参见docs。虽然对分组表达式执行此操作有点棘手。

【参考方案2】：

另一种方法是简单地重新采样并使用 OHLC (open=first,close=last,high=max,low=min)

In [56]: df = DataFrame('A' : np.arange(10), 'B' : pd.date_range('20130101',periods=5).tolist()*2)

In [57]: df
Out[57]: 
   A          B
0  0 2013-01-01
1  1 2013-01-02
2  2 2013-01-03
3  3 2013-01-04
4  4 2013-01-05
5  5 2013-01-01
6  6 2013-01-02
7  7 2013-01-03
8  8 2013-01-04
9  9 2013-01-05

In [58]: df.set_index('B').resample('D',how='ohlc')
Out[58]: 
              A               
           open high low close
B                             
2013-01-01    0    5   0     5
2013-01-02    1    6   1     6
2013-01-03    2    7   2     7
2013-01-04    3    8   3     8
2013-01-05    4    9   4     9