Pandas Groupby + Agg to List 非常慢

Question

我有一个包含 200,000 行的 DataFrame，由以下代码生成：

In [5]: def create_df():
   ...:     df = {}
   ...:     for col_name in ['a', 'b', 'c', 'd', 'e', 'f']:
   ...:         df[col_name] = np.tile(list(range(100000)), 2)
   ...:     return pd.DataFrame(df)

我正在尝试执行一些需要 groupby 的下游应用程序，然后将 groupby 行聚合到列表中，如下所示：

In [7]: df.groupby(['a', 'b', 'c', 'd']).agg(list)
Out[7]: 
                                      e               f
a     b     c     d                                    
0     0     0     0              [0, 0]          [0, 0]
1     1     1     1              [1, 1]          [1, 1]
2     2     2     2              [2, 2]          [2, 2]
3     3     3     3              [3, 3]          [3, 3]
4     4     4     4              [4, 4]          [4, 4]
...                                 ...             ...
99995 99995 99995 99995  [99995, 99995]  [99995, 99995]
99996 99996 99996 99996  [99996, 99996]  [99996, 99996]
99997 99997 99997 99997  [99997, 99997]  [99997, 99997]
99998 99998 99998 99998  [99998, 99998]  [99998, 99998]
99999 99999 99999 99999  [99999, 99999]  [99999, 99999]

但是，鉴于 DataFrame 只有 200,000 行，此操作比预期慢得多（平均约 15 秒）：

In [8]: %timeit df.groupby(['a', 'b', 'c', 'd']).agg(list)
14.7 s ± 1.13 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

相比之下，对 groupby 应用计数平均只需要大约 45 毫秒，这表明减速似乎只是按列表聚合：

In [10]: %timeit df.groupby(['a', 'b', 'c', 'd']).agg('count')
44.7 ms ± 2.23 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

我可以在这里做些什么来加快按列表操作的聚合？现在它是我代码中的真正瓶颈。

Answer 1

尝试df.groupby(['a', 'b', 'c', 'd']).agg(lambda x: x.tolist())。

Answer 2

一种方法是使用元组而不是列表：

df.groupby(['a', 'b', 'c', 'd']).agg(tuple)

这在我的环境中提供了 2x 加速：

df = create_df()
%timeit df.groupby(['a', 'b', 'c', 'd']).agg(tuple)
2.38 s ± 52.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit df.groupby(['a', 'b', 'c', 'd']).agg(list)
5.41 s ± 28.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Pandas Groupby + Agg to List 非常慢

Pandas Groupby + Agg to List Very Slow

python

aggregate-functions

pandas