如何加快对多列执行求和的 pandas groupby？

Question

问

我想加快 pandas groupby 的速度，它也在两列上应用求和并返回结果数据框。

代码

df = df.groupby(['key','code','name','period','agg_metric'], sort=False, observed=True, dropna=False)[['metricA','metricB']]\
.sum().reset_index()

（该方法目前需要 2 分钟来处理我最大用例的数据。）

数据

总体而言，最大尺寸的数据框有大约 150 万行应用了 groupby。 Period 和 agg_metric 可以相互推断，其中只有 2 个 period 值（因此 2 个 agg_metric 值）。名称值也可以从代码中推断出来。

在 groupby 之后，我剩下 70 万条记录。如果我理解正确，减速是由于处理的结果组的数量。有没有一种可能的方法来对该方法进行矢量化并将总和一次应用于每个组，而不是我假设当前正在迭代的方法。

备注

我尝试过使用groupby().agg({...}) 和groupby().apply(lambda)，两者花费的时间大致相同。我还尝试删除一些 groupby 列，然后稍后再将它们添加回来，但它并没有加快计算速度，因此不保证将它们从 groupby 中删除。 sn-p 也有 sort=False 和observed=True，但两者都没有改善处理时间。

我已经彻底浏览了尽可能多的资源（尤其是这个很好的参考资料：General Groupby in Python Pandas: Fast way）。我对矢量化相当陌生，并且正在这样做，因为我正在从我们的数据库中卸载几个查询。

Answer 1

你有什么类型的数据？ 看起来metricA / metricB 列的类型为object，pandas 对 Python 对象执行缓慢求和，而不是对 numpy 数组进行快速求和。尝试将度量列转换为float64 或integer 类型。

您可以使用df.info() 方法检查数据类型。

证明：

from string import ascii_letters
from time import time

import numpy as np
import pandas as pd
from numpy.random import choice

N = 1_500_000
np.random.seed(123)
letters = list(ascii_letters)
words = ["".join(choice(letters, 5)) for i in range(30)]

df = pd.DataFrame(
    {
        "key": choice(words, N),
        "code": choice(words, N),
        "name": choice(words, N),
        "period": np.random.randint(0, 10, N),
        "agg_metric": choice(["mean", "sum", "count"], N),
        "metricA": np.random.rand(N),
        "metricB": np.random.rand(N),
    }
)

def aggregate(df):
    return (
        df.groupby(
            ["key", "code", "name", "period", "agg_metric"],
            sort=False,
            observed=True,
            dropna=False,
        )[["metricA", "metricB"]]
        .sum()
        .reset_index()
    )

start = time()
df2 = aggregate(df)
print(f"sum floats took {time() - start}")

start = time()
df3 = aggregate(df.astype({"metricA": object, "metricB": object}))
print(f"sum objects took {time() - start}")

assert df2.equals(df3)

输出：

sum floats took 0.2983248233795166
sum objects took 81.04267287254333