Pandas中GROUP BY HAVING的等价操作

Question

在pandas中，使用groupby和并行应用过滤器的最有效方法是什么？

基本上，我在问这个操作在SQL中相当于什么。

select *
...
group by col_name
having condition

我认为有很多使用场景，比如条件平均值、总和、条件概率等等，这样的命令会非常强大。

我需要非常好的性能，所以理想情况下，这样的命令不应该是通过在python中进行多层操作得到的结果。

Answer 1

我按照州和县进行分组，筛选出最大值大于20的记录，然后使用数据框的loc方法对结果进行子查询，找出为真的值。

counties=df.groupby(['state','county'])['field1'].max()>20
counties=counties.loc[counties.values==True]

Answer 2

在pandas中，使用groupby和并行应用过滤的最有效方法是什么？

使用 groupby.transform + 布尔索引

虽然在pandas中有个类似的写法叫 groupby.filter，但它的速度非常慢。如果你在意性能，最好是在进行groupby操作后再过滤数据，而不是在操作过程中进行过滤。因为 groupby.filter 对每个组都要调用Python函数（比如lambda），而 groupby.transform 是对整个数据框调用一个经过Cython优化的函数，所以当组的数量很多时，后者的速度要快得多。

使用 groupby.transform 的好处是，它返回一个和原始数据框有相同索引的新数据框，里面填充的是聚合后的值。由于输出的索引和原始数据框相同，所以可以用来过滤原始数据框。

所以，

SELECT * FROM df GROUP BY colA HAVING COUNT(*) > 1

相当于

df[df.groupby('colA').transform('size') > 1]

而

SELECT * FROM df GROUP BY colA HAVING SUM(colB) > 5

相当于

df[df.groupby('colA')['colB'].transform('sum') > 5]

无论如何，正如下面的性能图所示，随着组的数量增加， groupby.transform + 布尔索引的速度远远快于 groupby.filter；例如，当组数达到1万时，它的速度快了1000倍。实际上，如果你的数据框有数百万个组， groupby.filter 可能根本无法运行，而 groupby.transform + 布尔索引会在合理的时间内完成工作。

用于生成上述图表的代码

import perfplot
import pandas as pd
import numpy as np

def groupby_filter(df):
    g = df.groupby('A')
    return g.filter(lambda x: x['B'].sum() > 5)

def groupby_transform(df):
    g = df.groupby('A')
    return df[g['B'].transform('sum') > 5]


perfplot.plot(
    kernels=[groupby_filter, groupby_transform],
    n_range=[2**k for k in range(16)],
    setup=lambda n: pd.DataFrame({
        'A': np.random.choice(n, size=n, replace=False), 
        'B': np.random.randint(n, size=n)}),
    xlabel='Number of groups'
)

Answer 3

正如unutbu的评论中提到的，groupby的过滤功能相当于SQL中的HAVING语句：

In [11]: df = pd.DataFrame([[1, 2], [1, 3], [5, 6]], columns=['A', 'B'])

In [12]: df
Out[12]:
   A  B
0  1  2
1  1  3
2  5  6

In [13]: g = df.groupby('A')  #  GROUP BY A

In [14]: g.filter(lambda x: len(x) > 1)  #  HAVING COUNT(*) > 1
Out[14]:
   A  B
0  1  2
1  1  3

你可以写更复杂的函数（这些函数会应用到每个组上），只要它们返回一个简单的布尔值（真或假）：

In [15]: g.filter(lambda x: x['B'].sum() == 5)
Out[15]:
   A  B
0  1  2
1  1  3

注意：可能存在一个bug，就是你不能在用于分组的列上直接写函数... 一个解决办法是手动对这些列进行分组，比如 g = df.groupby(df['A']))。