如何优化超大规模稀疏数据的Pandas DataFrame拼接方法？

我用自定义辅助函数get_df_concat(dfs)处理多个Pandas稀疏DataFrame的拼接聚合操作，逻辑是：通过pd.concat纵向拼接DataFrame，按索引分组后执行sum()聚合，最后对索引排序。小样本测试时功能正常，但处理超大规模稀疏数据（比如行数7e+5、列数2e+8的DataFrame）时，sum()方法耗时极长，要花数小时才能完成。

相关代码 #

生成随机稀疏DataFrame的函数 #

import pandas as pd
import numpy as np

def get_rnd_df(row:int=10, col:int=7): # 生成随机稀疏Pandas DataFrame
    np.random.seed(0)
    d=np.random.randint(low=0, high=10, size=(row,col)).astype(np.float32)
    d[d < 3] = np.nan
    df=pd.DataFrame(data=d,
                    index=[f"ip{i}" for i in np.random.choice(range(max(row, 10)), row, replace=False) ],
                    columns=[f"col_{c}" for c in np.random.choice(range(max(col, 10)), col, replace=False) ],
                    dtype=pd.SparseDtype(dtype=np.float32), # 稀疏格式：内存高效但速度极慢
                )
    df.index.name='usr'
    return df

原拼接聚合函数 #

def get_df_concat(dfs):
    t=time.time()
    dfc=pd.concat(dfs, axis=0, sort=True) # 纵向拼接dfs=[df1, df2,..., dfN], sort=True：对列排序
    print(f"elapsed_time [concat]{time.time()-t:>{12}.{4}f} sec")

    t=time.time()
    dfc=dfc.groupby(level=0) # 按索引分组
    print(f"elapsed_time [groupby]{time.time()-t:>{11}.{4}f} sec")

    t=time.time()
    dfc=dfc.sum() # <<<========== 耗时瓶颈 ==========>>> 
    print(f"elapsed_time [sum]{time.time()-t:>{15}.{4}f} sec")

    t=time.time()
    dfc=dfc.sort_index(key=lambda x: ( x.to_series().str[2:].astype(int) ))
    print(f"elapsed_time [sort idx]{time.time()-t:>{10}.{4}f} sec")

    return dfc

测试情况 #

小样本数据测试正常，但处理如下大规模数据时sum()方法速度极慢：

df1=get_rnd_df(row=int(7e+5), col=int(2e+8)) # 模拟真实数据规模
df2=get_rnd_df(row=int(9e+6), col=int(1e+9)) # 模拟真实数据规模

%%time
df_concat=get_df_concat(dfs=[df1, df2]) # 速度极慢！耗时严重！

需求询问 #

有没有更高效的替代方案实现该拼接聚合操作？能否利用SciPy csr_matrix提升处理速度？

更新优化方案 #

参考@RomanPerekhrest的建议，使用Pandas的numba引擎优化sum()方法，优化后的函数如下：

def get_df_concat_optimized(dfs):
    dfc=pd.concat(dfs, axis=0, sort=True).astype(pd.SparseDtype(dtype=np.float32)) # 纵向拼接dfs=[df1, df2,..., dfN], sort=True：对列排序
    dfc=dfc.groupby(level=0) # 按索引分组
    dfc=dfc.sum(engine="numba", # <<<=== 使用NUMBA引擎提速！
                engine_kwargs={'nopython': True, 'parallel': True, 'nogil': False},
                ).astype(pd.SparseDtype(dtype=np.float32, fill_value=0.0,))
    dfc=dfc.sort_index(key=lambda x: ( x.to_series().str[2:].astype(int) )).astype(pd.SparseDtype(dtype=np.float32, fill_value=0.0))
    return dfc

时间对比显示，优化后处理效率大幅提升：

df1=get_rnd_df(row=int(6e2), col=int(9e2))
df2=get_rnd_df(row=int(2e2), col=int(7e2))

%%time
df_concat_opt=get_df_concat_optimized(dfs=[df1, df2])
CPU times: user 2.32 s, sys: 11.6 ms, total: 2.33 s
Wall time: 2.47 s

%%time
df_concat=get_df_concat(dfs=[df1, df2])
elapsed_time [concat]      0.2443 sec
elapsed_time [groupby]     0.0008 sec
elapsed_time [sum]        67.2486 sec <<< 耗时严重 >>>
elapsed_time [sort idx]    0.2136 sec
CPU times: user 1min 6s, sys: 721 ms, total: 1min 7s
Wall time: 1min 7s

内容的提问来源于stack exchange，提问作者farid