优化大型数据集上Pandas GroupBy与自定义聚合的性能咨询

当前方法（简化版本） #

import pandas as pd
 
df = pd.DataFrame({
    'category': ['A', 'B', 'A', 'B'] * 10**6,
    'subcategory': ['X', 'Y', 'X', 'Y'] * 10**6,
    'value': [1, 2, 3, 4] * 10**6,
    'quantity': [10, 20, 30, 40] * 10**6
})
 
agg_functions = {
    'value': ['sum', 'mean'],
    'quantity': [lambda x: x.sum(), lambda x: (x > 20).mean()]
}

result = df.groupby(['category', 'subcategory']).agg(agg_functions)

存在问题 #

• 数据集规模达3050万行，真实数据处理时出现内存不足
• 自定义聚合函数（如quantity列的lambda函数）运行缓慢，无法利用向量化操作优势
• 希望在不采用分块处理的前提下，实现更内存高效的聚合

咨询问题 #

1. 如何优化GroupBy与聚合操作，高效处理大型数据集，尤其是应用自定义函数时？
2. Numba、Cython或并行处理等高级技术能否加速Pandas GroupBy中的自定义聚合？
3. 转换为Polars、Dask或PySpark等结构能否带来显著性能提升？还是可以在Pandas内实现足够优化？

寻求兼顾内存管理、速度优化且保留自定义聚合灵活性的方案，期待最佳实践或高级技术指导。

已尝试方案 #

• 基础Pandas GroupBy：使用groupby()结合自定义lambda函数，百万级数据集下速度过慢，条件聚合的lambda函数拖慢性能尤为明显
• 内存优化：将列转换为高效数据类型（分类列用category，整数用int32），缓解了内存问题，但GroupBy与聚合的执行时间仍过长
• 数据分块：用pd.read_csv()的chunksize参数分块处理，但跨分块GroupBy操作管理难度大，合并结果额外增加开销
• joblib多线程：并行化自定义聚合函数，但性能提升极小，推测是lambda函数特性与线程管理开销导致
• Dask DataFrame：尝试分配工作负载至多核心，但自定义函数管理与分布式操作效率问题导致复杂度提升，未获显著性能改善

解决方案 #

一、Pandas内的GroupBy与聚合优化 #

1. 替换自定义lambda为内置向量化函数
   对于(x > 20).mean()这类条件聚合，无需用lambda，可预先计算布尔列后直接调用内置聚合函数，完全利用Pandas向量化计算优势：

   # 预先计算布尔列
   df['quantity_gt20'] = df['quantity'] > 20
   agg_functions = {
       'value': ['sum', 'mean'],
       'quantity': ['sum'],
       'quantity_gt20': ['mean']
   }
   result = df.groupby(['category', 'subcategory']).agg(agg_functions)
   

2. 使用命名聚合替代匿名lambda
   明确命名聚合函数，避免匿名lambda的额外开销，同时提升代码可读性：

   def count_gt20(x):
       return (x > 20).mean()
   
   agg_functions = {
       'value': {'value_sum': 'sum', 'value_mean': 'mean'},
       'quantity': {'qty_sum': 'sum', 'qty_gt20_mean': count_gt20}
   }
   result = df.groupby(['category', 'subcategory']).agg(agg_functions)
   

3. 内存优化进阶

   • 分类列设置ordered=True，加速分组哈希计算
   • 使用df.groupby(..., observed=True)（针对分类列），只聚合实际存在的分组，减少无效计算
   • 避免创建不必要的中间列，优先复用已有列完成计算

二、高级技术加速自定义聚合 #

1. Numba JIT编译加速
   Pandas支持用Numba装饰自定义聚合函数，将Python循环编译为机器码，大幅提升速度：

   from numba import jit
   
   @jit(nopython=True)
   def numba_gt20_mean(arr):
       count = 0
       for val in arr:
           if val > 20:
               count +=1
       return count / len(arr)
   
   # 针对单列的自定义聚合用apply更高效
   result = df.groupby(['category', 'subcategory'])['quantity'].apply(numba_gt20_mean)
   

   复杂自定义逻辑下，Numba能将速度提升10-100倍，接近C语言水平。

2. 并行处理优化

   • Pandas 2.0+支持groupby.agg(..., engine='numba', parallel=True)，结合Numba实现自动并行化
   • 使用swifter库自动适配函数的向量化/并行化逻辑，简化代码：

     import swifter
     
     def custom_agg(x):
         return (x >20).mean()
     
     df.groupby(['category', 'subcategory'])['quantity'].swifter.apply(custom_agg)
     

   并行处理适合CPU密集型自定义聚合，IO密集型场景提升有限。

3. Cython优化（极致性能需求）
   编写Cython代码直接操作内存数据，能达到接近原生C的性能，但学习成本高，适合长期维护的核心逻辑。

三、其他数据结构的性能对比 #

1. Polars
   基于Rust的DataFrame库，天生支持向量化和并行计算，内存效率是Pandas的1/3-1/2，聚合速度快2-5倍。自定义聚合用表达式语法即可实现，代码改动量小：

   import polars as pl
   
   pl_df = pl.from_pandas(df)
   result = pl_df.groupby(['category', 'subcategory']).agg(
       pl.col('value').sum().alias('value_sum'),
       pl.col('value').mean().alias('value_mean'),
       pl.col('quantity').sum().alias('qty_sum'),
       (pl.col('quantity') >20).mean().alias('qty_gt20_mean')
   )
   

   适合单机器处理3000万级别的数据集，是当前性价比最高的替代方案。

2. Dask
   适合超大规模（远超内存）数据集，但单机器内存可容纳的场景下，分布式调度开销大于性能收益，自定义函数适配成本高，不优先推荐。

3. PySpark
   针对TB级分布式场景，单机器处理3000万行数据时，集群启动和调度开销远大于性能提升，自定义UDF性能不如Polars或Numba加速的Pandas。

最终建议 #

• 若需保留Pandas生态：优先用向量化操作替换lambda，配合Numba加速复杂自定义逻辑，叠加内存优化
• 若追求极致性能与内存效率：直接迁移到Polars，代码改动小，性能提升显著
• 仅当数据集远超单机器内存时，再考虑Dask或PySpark

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