加速 agg 并加入一个十亿记录的 pandas 表

Posted 2023-03-11

【中文标题】加速 agg 并加入一个十亿记录的 pandas 表

【英文标题】：speeding up agg and join on a billion-record pandas table

【发布时间】：2019-07-01 03:24:57

【问题描述】：

[python 3.5.2, pandas 0.24.1, numpy 1.16.1, scipy 1.2.0]

我有以下熊猫数据框

data_pd
    nrows: 1,032,749,584
    cols: ['mem_id':np.uint32, 'offset':np.uint16 , 'ctype':string, 'code':string]

obsmap_pd
    nrows: 10,887,542
    cols: ['mem_id':np.uint32, 'obs_id':np.uint32]    
             (obs_id has consecutive integers between 0 and obsmap_pd nrows)

varmap_pd
    nrows: 4,596
    cols: ['ctype':string, 'code': string, 'var_id':np.uint16]   
             (var_id has consecutive integers between 0 and varmap_pd nrows)

这些是我正在运行的步骤

***
sparse_pd = data_pd.groupby(['mem_id','ctype','code'])['offset'].nunique().reset_index(name='value')
sparse_pd['value'] = sparse_pd['value'].astype(np.uint16)
sparse_pd = pd.merge(pd.merge(sparse_pd, obsmap_pd, on='mem_id', sort=False),
                  varmap_pd, on=['ctype','code'], sort=False)[['obs_id','var_id','value']]
***

这样做的目的是在下一步创建一个scipy csc_matrix

mat_csc = csc_matrix((sparse_pd['value'].values*1., (sparse_pd['obs_id'].values,sparse_pd['var_id'].values)), 
                     shape=(obsmap_pd.shape[0],varmap_pd.shape[0]))

csc_matrix 的创建速度非常快，但是带有 pandas 代码的三行（*** 之间）需要 25.7 分钟。关于如何加快速度的任何想法？

【问题讨论】：

一个好的策略是使用连接而不是合并，并在on 列上设置索引。您需要设置索引、连接、设置索引连接，而不是全部放在一行上。您可以在这里看到一个示例：***.com/questions/40860457/… 也许这甚至是重复的？

Improve Pandas Merge performance的可能重复

都是合并吗？ groupby step 和 astype 需要多长时间？您会惊讶于 .groupby.nunique 对于可能有数千万或数亿组的 DataFrame 的速度有多慢。

你考虑过使用Dask

【参考方案1】：

加速合并的一个好方法是使用 join 代替：

sparse_pd = sparse_pd.\
    .set_index(['mem_id']).join(obsmap_pd.set_index('mem_id']))
    
sparse_pd = sparse_pd.\
    .reset_index().set_index(['ctype','code']).join(varmap_pd.set_index(['ctype','code'])).loc['obs_id','var_id','value']

另外，考虑在 pandas 中使用 pipe 模块（documentation here）； 你失去了一些可读性，但它可能更优化了。