加快 pandas groupby 中的滚动总和计算

Question

我想按组计算大量组的滚动总和，但我无法快速完成。

Pandas 内置滚动和扩展计算方法

这是一个例子：

import pandas as pd
import numpy as np
obs_per_g = 20
g = 10000
obs = g * obs_per_g
k = 20
df = pd.DataFrame(
    data=np.random.normal(size=obs * k).reshape(obs, k),
    index=pd.MultiIndex.from_product(iterables=[range(g), range(obs_per_g)]),
)

为了滚动和扩展总和，我可以使用

df.groupby(level=0).expanding().sum()
df.groupby(level=0).rolling(window=5).sum()

但这对于非常多的组来说需要很长时间。对于扩展总和，使用 pandas 方法 cumsum 几乎快 60 倍（16 秒对上例中的 280 毫秒）并将小时变成分钟。

df.groupby(level=0).cumsum()

在pandas中有没有像cumsum一样快速实现rolling sum的？如果没有，我可以使用 numpy 来完成这个吗？

Answer 1

我对 .rolling() 有同样的体验，它很好，但只适用于小数据集，或者如果您应用的函数不是标准的，sum() 我建议使用 cumsum()并减去 cumsum().shift(5)

df.groupby(level=0).cumsum() - df.groupby(level=0).cumsum().shift(5)

Answer 2

为了提供这方面的最新信息，如果您升级 pandas，groupby 滚动的性能已得到显着改善。与 0.24 或 1.0.0 相比，这在 1.1.0 中快了大约 4-5 倍，在 >1.2.0 中快了 12 倍。

我相信最大的性能提升来自这个 PR 这意味着它可以在 cython 中做更多（在它像 groupby.apply(lambda x: x.rolling()) 那样实现之前）。

我使用以下代码进行基准测试：

import pandas
import numpy

print(pandas.__version__)
print(numpy.__version__)


def stack_overflow_df():
    obs_per_g = 20
    g = 10000
    obs = g * obs_per_g
    k = 2
    df = pandas.DataFrame(
        data=numpy.random.normal(size=obs * k).reshape(obs, k),
        index=pandas.MultiIndex.from_product(iterables=[range(g), range(obs_per_g)]),
    )
    return df


df = stack_overflow_df()

# N.B. droplevel important to make indices match
rolling_result = (
    df.groupby(level=0)[[0, 1]].rolling(10, min_periods=1).sum().droplevel(level=0)
)
df[["value_0_rolling_sum", "value_1_rolling_sum"]] = rolling_result
%%timeit
# results:
# numpy version always 1.19.4
# pandas 0.24 = 12.3 seconds
# pandas 1.0.5 = 12.9 seconds
# pandas 1.1.0 = broken with groupby rolling bug
# pandas 1.1.1 = 2.9 seconds
# pandas 1.1.5 = 2.5 seconds
# pandas 1.2.0 = 1.06 seconds
# pandas 1.2.2 = 1.06 seconds

我认为如果尝试使用 numpy.cumsum 来提高性能（无论 pandas 版本如何），必须小心。例如，使用如下内容：

# Gives different output
df.groupby(level=0)[[0, 1]].cumsum() - df.groupby(level=0)[[0, 1]].cumsum().shift(10)

虽然这要快得多，但输出不正确。此移位在所有行上执行并混合不同组的累积和。即下一组的第一个结果移回上一组。

要获得与上述相同的行为，您需要使用 apply:

df.groupby(level=0)[[0, 1]].cumsum() - df.groupby(level=0)[[0, 1]].apply(
    lambda x: x.cumsum().shift(10).fillna(0)
)

在最新版本 (1.2.2) 中，这比直接使用滚动要慢。因此，对于 groupby 滚动总和，我认为 numpy.cumsum 不是 pandas>=1.1.1

的最佳解决方案

为了完整起见，如果您的组是列而不是索引，您应该使用如下语法：

# N.B. reset_index important to make indices match
rolling_result = (
    df.groupby(["category_0", "category_1"])[["value_0", "value_1"]]
    .rolling(10, min_periods=1)
    .sum()
    .reset_index(drop=True)
)
df[["value_0_rolling_sum", "value_1_rolling_sum"]] = rolling_result

加快 pandas groupby 中的滚动总和计算

Speeding up rolling sum calculation in pandas groupby

python

performance

pandas

rolling-computation

pandas-groupby