python - 最有效的 groupby 方式 => 聚合 pandas 中的大型数据框

Question

我有一个包含大约 150,000,000 行的 pandas 数据框，格式如下:

df.head()

Out[1]:
     ID    TERM    X
0    1     A       0
1    1     A       4
2    1     A       6
3    1     B       0
4    1     B       10
5    2     A       1
6    2     B       1
7    2     F       1

我想按 ID 和 TERM 聚合它，并计算行数。目前，我执行以下操作:

df.groupby(['ID','TERM']).count()

Out[2]:
     ID    TERM    X
0    1     A       3
1    1     B       2
2    2     A       1
3    2     B       1
4    2     F       1

但这大约需要两分钟。使用 R data.tables 的相同操作只需不到 22 秒。在 python 中有更有效的方法吗？

为了比较，R data.table:

system.time({ df[,.(.N), .(ID, TERM)] })
#user: 30.32 system: 2.45 elapsed: 22.88

Answer 1

NumPy 解决方案就像这样 -

def groupby_count(df):
    unq, t = np.unique(df.TERM, return_inverse=1)
    ids = df.ID.values
    sidx = np.lexsort([t,ids])

    ts = t[sidx]
    idss = ids[sidx]

    m0 = (idss[1:] != idss[:-1]) | (ts[1:] != ts[:-1])
    m = np.concatenate(([True], m0, [True]))
    ids_out = idss[m[:-1]]
    t_out = unq[ts[m[:-1]]]
    x_out = np.diff(np.flatnonzero(m)+1)

    out_ar = np.column_stack((ids_out, t_out, x_out))
    return pd.DataFrame(out_ar, columns = [['ID','TERM','X']])

一个更简单的版本-

def groupby_count_v2(df):    
    a = df.values
    sidx = np.lexsort(a[:,:2].T)
    b = a[sidx,:2]
    m = np.concatenate(([True],(b[1:] != b[:-1]).any(1),[True]))
    out_ar = np.column_stack((b[m[:-1],:2], np.diff(np.flatnonzero(m)+1)))
    return pd.DataFrame(out_ar, columns = [['ID','TERM','X']])

sample 运行-

In [332]: df
Out[332]: 
   ID TERM   X
0   1    A   0
1   1    A   4
2   1    A   6
3   1    B   0
4   1    B  10
5   2    A   1
6   2    B   1
7   2    F   1

In [333]: groupby_count(df)
Out[333]: 
  ID TERM  X
0  1    A  3
1  1    B  2
2  2    A  1
3  2    B  1
4  2    F  1

让我们随机打乱行并验证它是否适用于我们的解决方案 -

In [339]: df1 = df.iloc[np.random.permutation(len(df))]

In [340]: df1
Out[340]: 
   ID TERM   X
7   2    F   1
6   2    B   1
0   1    A   0
3   1    B   0
5   2    A   1
2   1    A   6
1   1    A   4
4   1    B  10

In [341]: groupby_count(df1)
Out[341]: 
  ID TERM  X
0  1    A  3
1  1    B  2
2  2    A  1
3  2    B  1
4  2    F  1