如何提高对 pandas 数据框的聚合操作的速度？答案

【问题标题】：How to improve the speed of an aggregate operation on a pandas dataframe?如何提高对 pandas 数据框的聚合操作的速度？
【发布时间】：2018-05-15 07:09:20
【问题描述】：

我有一个名为 data_1 的 pandas 数据框，如下所示：

tenant_id  tenant_class_id tenant_date
100         10001           2004-01-10
100         10001           2004-08-25
101         10001           2016-09-20
102         10002           2013-07-12
103         10003           2011-04-05
103         10004           2002-05-05
103         10003           2003-12-26

我的目标是为每个tenant_id 找到最早/第一个tenant_date。 tenant_id 是 long 格式，tenant_date 是 datetime64 格式

我用这段代码来计算它：

output = data_1.groupby(['tenant_id']).apply(lambda x: x.nsmallest(1, 'tenant_date')).reset_index(drop=1)

这对于较小的数据集非常有效，但我的 data_1 包含大约 2 亿行。这段代码永远不会完成。我有足够的内存和大约 244 GB 的 RAM。有没有更有效的方法可以快速做到这一点？

【问题讨论】：

标签： python pandas pandas-groupby

【解决方案1】：

你可以试试sort_values + drop_duplicates:

output = data_1.sort_values(['tenant_id','tenant_date']).drop_duplicates(['tenant_id'])

或者sort_values + GroupBy.head:

output = data_1.sort_values(['tenant_id','tenant_date']).groupby(['tenant_id']).head(1)

或者通过idxmin查看每组的最小索引：

output = data_1.loc[data_1.groupby(['tenant_id'])['tenant_date'].idxmin()]

或者：

output = data_1.set_index(['tenant_class_id'])
               .groupby(['tenant_id'])['tenant_date'].nsmallest(1).reset_index()

时间安排：

np.random.seed(123)
N = 1000000
L = pd.date_range('2015-01-01', '2018-01-01')

df = pd.DataFrame({'tenant_date': np.random.choice(L, size=N),
                   'tenant_id':np.random.randint(1000,size=N),
                   'tenant_class_id':np.random.randint(1000,size=N)})
print (df)

In [99]: %timeit data_1.sort_values(['tenant_id','tenant_date']).drop_duplicates(['tenant_id'])
1000 loops, best of 3: 1.97 ms per loop

In [100]: %timeit data_1.sort_values(['tenant_id','tenant_date']).groupby(['tenant_id']).head(1)
1000 loops, best of 3: 2.07 ms per loop

In [101]: %timeit data_1.loc[data_1.groupby(['tenant_id'])['tenant_date'].idxmin()]
100 loops, best of 3: 2.04 ms per loop

In [102]: %timeit data_1.set_index(['tenant_class_id']).groupby(['tenant_id'])['tenant_date'].nsmallest(1).reset_index()
100 loops, best of 3: 8.64 ms per loop

In [103]: %timeit data_1.groupby(['tenant_id']).apply(lambda x: x.nsmallest(1, 'tenant_date')).reset_index(drop=1)
100 loops, best of 3: 11.4 ms per loop

警告

考虑到组的数量，结果并未解决性能问题，这将对其中一些解决方案的时间产生很大影响。

【讨论】：

【解决方案2】：

排序然后使用drop_duplicates。但请务必先将日期列转为日期时间。

df['tenant_date'] = pd.to_datetime(df['tenant_date'])
df.sort_values(['tenant_id', 'tenant_date']).drop_duplicates('tenant_id')

   tenant_id  tenant_class_id tenant_date
0        100            10001  2004-01-10
2        101            10001  2016-09-20
3        102            10002  2013-07-12
5        103            10004  2002-05-05

【讨论】：