在 Pandas/Numpy 中最有效地使用 groupby-apply 和用户定义的函数

Question

我缺少有关在 Pandas 或 Numpy 的 groupby-apply 设置中使用用户定义函数的最有效（阅读：最快）方式的信息。我已经做了一些自己的测试，但我想知道是否还有其他我还没有遇到过的方法。

以DataFrame为例：

import numpy as np
import pandas as pd

idx = pd.MultiIndex.from_product([range(0, 100000), ["a", "b", "c"]], names = ["time", "group"])
df = pd.DataFrame(columns=["value"], index = idx)

np.random.seed(12)
df["value"] = np.random.random(size=(len(idx),))

print(df.head())

               value
time group          
0    a      0.154163
     b      0.740050
     c      0.263315
1    a      0.533739
     b      0.014575

我想计算（例如，下面可以是任意用户定义的函数）每组随时间变化的百分比。我可以在纯 Pandas 实现中执行此操作，如下所示：

def pct_change_pd(series, num):
    return series / series.shift(num) - 1

out_pd = df.sort_values(['group', 'time']).groupby(["group"]).apply(pct_change_pd, num=1)

但我也可以修改函数并将其应用于 numpy 数组：

def shift_array(arr, num, fill_value=np.nan):
    if num >= 0:
        return np.concatenate((np.full(num, fill_value), arr[:-num]))
    else:
        return np.concatenate((arr[-num:], np.full(-num, fill_value)))
    
def pct_change_np(series, num):
    idx = series.index

    arr = series.values.flatten()
    arr_out = arr / shift_array(arr, num=num) - 1
    return pd.Series(arr_out, index=idx)

out_np = df.sort_values(['group', 'time']).groupby(["group"]).apply(pct_change_np, num=1)
out_np = out_np.reset_index(level=2, drop=True)

从我的测试来看，numpy 方法似乎更快，即使在np.array 和pd.Series 之间转换的额外开销也是如此。

熊猫：

%%timeit
out_pd = df.sort_values(['group', 'time']).groupby(["group"]).apply(pct_change_pd, num=1)

113 ms ± 548 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

麻木：

%%timeit
out_np = df.sort_values(['group', 'time']).groupby(["group"]).apply(pct_change_np, num=1)
out_np = out_np.reset_index(level=2, drop=True)

94.7 ms ± 642 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

随着索引的增长和用户定义的函数变得越来越复杂，Numpy 的实现将越来越胜过 Pandas 的实现。但是，我想知道是否有其他方法可以更快地获得类似结果。 我特别关注另一种（更有效的）groupby-apply 方法，它允许我使用任意用户定义的函数，而不仅仅是计算百分比变化的示例。会很高兴听听他们是否存在！

Answer 1

游戏的名称通常是尝试使用工具箱中的任何函数（通常经过优化和 C 编译），而不是应用您自己的纯 Python 函数。例如，一种替代方法是：

def f1(df, num=1):
    grb_kwargs = dict(sort=False, group_keys=False)  # avoid redundant ops
    z = df.sort_values(['group', 'time'])
    return z / z.groupby('group', **grb_kwargs).transform(pd.Series.shift, num) - 1

这比.groupby('group').apply(pct_change_pd, num=1) 快​​大约 32%。在您的系统上，它会产生大约 85 毫秒。

然后，有一个技巧是对整个 df 进行“昂贵”的计算，但屏蔽掉其他组溢出的部分：

def f2(df, num=1):
    grb_kwargs = dict(sort=False, group_keys=False)  # avoid redundant ops
    z = df.sort_values(['group', 'time'])
    z2 = z.shift(num)
    gid = z.groupby('group', **grb_kwargs).ngroup()
    z2.loc[gid != gid.shift(num)] = np.nan
    return z / z2 - 1

那个速度要快 2.1 倍（在您的系统上大约是 52.8 毫秒）。

最后，当无法找到直接使用的向量化函数时，您可以使用 numba 来加速您的代码（然后可以用循环编写代码，让您随心所欲）... 一个经典的例子是累积的加上大写，如this SO post 和this one。

Answer 2

你的第一个函数和使用 .apply() 给了我这个结果：

In [42]: %%timeit
    ...: out_pd = df.sort_values(['group', 'time']).groupby(["group"]).apply(pct_change_pd, num=1)
155 ms ± 887 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

使用组，时间达到 56 毫秒。

%%timeit
num=1
outpd_list = []
for g in dfg.groups.keys():
    gc = dfg.get_group(g)
    outpd_list.append(gc['value'] / gc['value'].shift(num) - 1)
out_pd = pd.concat(outpd_list, axis=0)

56 ms ± 821 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

如果您将上述代码中的这一行更改为使用内置函数，您将节省更多时间

outpd_list.append(gc['value'].pct_change(num))
41.2 ms ± 283 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)