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

Posted 2023-03-11

【中文标题】加快pandas groupby中的滚动总和计算

【英文标题】：Speeding up rolling sum calculation in pandas groupby

【发布时间】：2019-11-15 00:29:51

【问题描述】：

我想为大量组计算滚动总和，但我无法以可接受的速度快速完成。

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 来完成这个吗？

【参考方案1】：

为了提供这方面的最新信息，如果升级 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 来提高性能，我认为必须小心（无论熊猫版本如何）。例如，使用如下内容：

# 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) 中，它比直接使用滚动要慢。 因此，对于 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

【参考方案2】：

我对@9​​87654321@ 也有同样的体验，它很好，但仅适用于小型数据集，或者如果您应用的函数是非标准的，使用sum() 我建议使用cumsum() 并减去cumsum().shift(5)

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

【讨论】：

我刚刚检查过，令人惊讶的是 .rolling() 在 242 µs 时比 %timeit 快一点，而我的方法是 %timeit 371 µs，我的数据集的体验不同，它快了大约 10 倍，不知道为什么。

这是一个很好的解决方案，应该想到这一点！对于单个组或少量组，Cumsum 并不比 expand().sum()（或 rolling()）快。但是对于大量组而言，它变得更快。必须有 cumsum 的优化，这与 groupby 的完成方式有关

我不确定这个答案是否按预期工作。 df.groupby(level=0).cumsum().shift(5) 不是在所有行上移动并混合不同组的 cumsum 吗？即下一组的第一个结果移回上一组？我认为您需要在申请中包含转变。像这样：df.groupby(level=0).cumsum() - df.groupby(level=0).apply(lambda x: x.cumsum().shift(10).fillna(0)) 我的基准测试表明这比 pandas 滚动快大约 2 倍。 （与上述没有给出相同输出的答案的时间相比，速度相当慢）。