在 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 用于扩展求和一样？如果没有，我能否使用 numpy 来完成这个任务？

- CloseToC

Answer 1

我曾经使用过 .rolling()，虽然很好用，但只适用于小型数据集或者你应用的函数比较非标准。对于 sum()，我建议使用 cumsum() 并减去 cumsum().shift(5)。

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

- Mark

---

我刚刚检查了一下，令人惊讶的是 .rolling() 比我的方法更快， %timeit 为 242 微秒，而我的方法为 371 微秒。我的经验是，对于我的数据集来说，它要快大约10倍，不知道为什么。 - Mark

这是一个很好的解决方案，我应该想到了！对于单个组或少量组，cumsum 比 expanding().sum()（或 rolling()）没有更快的速度。但是对于大量组，它变得明显更快。必须优化 cumsum 与 groupby 的执行方式有关。 - CloseToC

我不确定这个答案是否按预期工作。df.groupby(level=0).cumsum().shift(5)难道不会将所有行都移位并混合不同组的累加和吗？也就是说，下一组的第一个结果被移回到了上一组？我认为你需要在apply中包含shift。像这样：df.groupby(level=0).cumsum() - df.groupby(level=0).apply(lambda x: x.cumsum().shift(10).fillna(0))我的基准测试显示，这比pandas rolling快约两倍。（与上面的答案相比，速度相当慢，而且输出不同）。 - user2175850

Answer 2

为了提供最新信息，如果升级pandas，则groupby rolling的性能得到了显着提高。与0.24或1.0.0相比，在1.1.0中快了约4-5倍，在> 1.2.0中快了x12。

我认为最大的性能改进来自于此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)

虽然这种方法更快，但输出结果不正确。此移位操作会在所有行上执行，并混合不同组的cumsum。即下一组的第一个结果被移回到前一组。

要获得与上述相同的行为，需要使用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）中，使用rolling直接计算速度比较慢。因此，对于groupby rolling sums，我认为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

- user2175850