Python Pandas中的通用分组: 快速方法

Question

终极问题

是否有一种通用且高效的分组操作方法，不依赖于pd.groupby函数？

输入

pd.DataFrame([[1, '2020-02-01', 'a'], [1, '2020-02-10', 'b'], [1, '2020-02-17', 'c'], [2, '2020-02-02', 'd'], [2, '2020-03-06', 'b'], [2, '2020-04-17', 'c']], columns=['id', 'begin_date', 'status'])`

   id  begin_date status
0   1  2020-02-01      a
1   1  2020-02-10      b
2   1  2020-02-17      c
3   2  2020-02-02      d
4   2  2020-03-06      b

期望输出

   id status  count  uniquecount
0   1      a      1            1
1   1      b      1            1
2   1      c      1            1
3   2      b      1            1
4   2      c      1            1

问题

现在，有一种使用Pandas的Python简单方法可以做到这一点。

df = df.groupby(["id", "status"]).agg(count=("begin_date", "count"), uniquecount=("begin_date", lambda x: x.nunique())).reset_index()
# As commented, omitting the lambda and replacing it with "begin_date", "nunique" will be faster. Thanks!

对于较大的数据集，这个操作速度较慢，我猜测时间复杂度为O(n²)。

现有解决方案缺乏所需的普适性

经过一些谷歌搜索，StackOverflow 上有一些替代方案，使用 numpy、iterrows 或其他不同的方式。

更快的 Pandas 分组操作替代方法

Pandas 快速加权随机选择（Weighted Random Choice）

还有一个很好的答案：

Python Pandas 中的 Groupby：快速方法

这些解决方案通常旨在创建“count”或“uniquecount”，即聚合值，就像我例子中的一样。但是，很遗憾，它们总是只有一种聚合方式，并且没有多个 groupby 列。

另外，它们很可惜从未解释如何将它们合并到分组的 dataframe 中。

是否有一种方法可以使用 itertools（例如这个答案：更快的 Pandas 分组操作替代方法，或者更好的这个答案：Python Pandas 中的 Groupby：快速方法），不仅返回系列“count”，而且还要以分组形式返回整个 dataframe？

终极问题

是否有一种通用、高效的 groupby 操作方式，不依赖于 pd.groupby？

它可能看起来像这样：

from typing import List
def fastGroupby(df, groupbyColumns: List[str], aggregateColumns):
    # numpy / iterrow magic
    return df_grouped

df = fastGroupby(df, ["id", "status"], {'status': 'count',
                             'status': 'count'}

并返回期望的输出。

- Dustin

---

lambda函数正在拖垮你。它强制将其转换为组之间的缓慢循环。uniquecount=('status', 'nunique') 可能会将速度提高数倍。然后，可以在 groupby 调用中进一步添加 sort=False，这将使输出无序，但将显著提高多个组的速度。 - ALollz

感谢您的评论，@ALollz提到的lambda表达式是完全正确的。实际上，在示例中我犯了一个错误，我聚合的特定列并不重要，但我会进行修正。谢谢！ - Dustin

Answer 1

在放弃使用 groupby 之前，建议先评估一下您是否真正利用了 groupby 所提供的功能。

使用内置的 pd.DataFrameGroupBy 方法来替换 lambda 函数。

许多 Series 和 DataFrame 方法都是实现为 pd.DataFrameGroupBy 方法。您应该直接使用这些方法，而不是通过调用 groupby + apply（lambda x：...） 来调用它们。

此外，对于许多计算，您可以将问题重新构造为对整个 DataFrame 的某个向量化操作，然后使用在 cython 中实现的 groupby 方法。这样会很快。

一个常见的例子是在组内找到答案为 'Y' 的比例。一个简单的方法是在每个组内检查条件，然后获取比例：

N = 10**6
df = pd.DataFrame({'grp': np.random.choice(range(10000), N),
                   'answer': np.random.choice(['Y', 'N'], N)})

df.groupby('grp')['answer'].apply(lambda x: x.eq('Y').mean())

以这种方式考虑问题需要使用lambda，因为我们在groupby中进行了两个操作：检查条件然后求平均值。可以将完全相同的计算视为首先在整个DataFrame上检查条件，然后在组内计算平均值：

df['answer'].eq('Y').groupby(df['grp']).mean()

这是一个非常微小的变化，但后果却是巨大的。随着团队数量的增加，收益也将越来越大。

%timeit df.groupby('grp')['answer'].apply(lambda x: x.eq('Y').mean())
#2.32 s ± 99.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit df['answer'].eq('Y').groupby(df['grp']).mean()
#82.8 ms ± 995 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

添加sort=False作为参数

默认情况下，groupby会按键对输出进行排序。如果没有必要进行排序，则可以通过指定sort=False来获得轻微的提速。

---

添加observed=True作为参数

如果分组键是分类变量，它将重新索引到所有可能的组合，即使这些组在您的数据框中从未出现过。如果这些组不重要，则从输出中删除它们将极大地提高速度。

---

针对您的示例，我们可以比较一下两者之间的差异。使用pd.DataFrameGroupBy.nunique可以获得巨大提速，同时去除排序可以额外提高速度。两者结合在一起形成了一个“相同”的解决方案（除排序外），且在许多分组时快近100倍。

import perfplot
import pandas as pd
import numpy

def agg_lambda(df):
    return df.groupby(['id', 'status']).agg(uniquecount=('Col4', lambda x: x.nunique()))
    
def agg_nunique(df):
    return df.groupby(['id', 'status']).agg(uniquecount=('Col4', 'nunique'))

def agg_nunique_nosort(df):
    return df.groupby(['id', 'status'], sort=False).agg(uniquecount=('Col4', 'nunique'))

perfplot.show(
    setup=lambda N: pd.DataFrame({'Col1': range(N),
                       'status': np.random.choice(np.arange(N), N),
                       'id': np.random.choice(np.arange(N), N),
                       'Col4': np.random.choice(np.arange(N), N)}),
    kernels=[
        lambda df: agg_lambda(df),
        lambda df: agg_nunique(df),
        lambda df: agg_nunique_nosort(df),
    ],
    labels=['Agg Lambda', 'Agg Nunique', 'Agg Nunique, No sort'],
    n_range=[2 ** k for k in range(20)],
    # Equality check same data, just allow for different sorting
    equality_check=lambda x,y: x.sort_index().compare(y.sort_index()).empty,
    xlabel="~ Number of Groups"
)

- ALollz

---

非常棒的答案，非常感谢。我采纳了你的反馈，确实运行速度更快了。但因为仍然需要花费相当长的时间，我将保持问题开放，以便讨论其他解决方案，例如使用numpy。 - Dustin

@Dustin同意，可能有更快的方法。 - ALollz