Pandas Groupby求和效率低下，求大数据量优化方案

给定条件 #

使用以下辅助函数生成pandas DataFrame格式的「用户vs令牌」合成数据集：

import numpy as np
import pandas as pd
import string
import random

def get_df(nROWs:int=10, nCOLs:int=100, MIN=0.0, MAX=199.0):
    my_strings = string.printable
    df = pd.DataFrame(np.random.uniform(low=MIN, high=MAX, size=(nROWs, nCOLs)).astype("float16"), columns=list(map(lambda orig_string: "tk_"+orig_string, random.sample(my_strings, nCOLs))) )
    df["user_ip"] = [f"u{random.randint(0, nROWs)}" for r in range(nROWs)]
    return df

目标 #

按user_ip分组后，对每组的各列数值求和。

现有低效方案问题 #

针对小数据集，先通过pd.concat([df1, df2, df3], axis=0)合并，再循环遍历groupby("user_ip")结果，将每组非user_ip列的求和结果存入字典后转为DataFrame。该方案在小数据下可行，但处理150万行级大数据时，循环单次执行耗时约10~15秒，效率极低。

高效解决方案 #

在内存充足的超算环境下，直接利用pandas原生的矢量化聚合能力即可解决，完全不需要手动循环：

1. 基础高效实现 #

合并数据集后，直接调用groupby+sum方法：

# 合并数据集（若有多个分块）
combined_df = pd.concat([df1, df2, df3], axis=0)
# 按user_ip分组，对所有数值列求和
result_df = combined_df.groupby("user_ip").sum(numeric_only=True)

• numeric_only=True确保只处理数值列，避免非数值列干扰；
• pandas底层对groupby.sum做了高度优化，是C级别的矢量化操作，性能比Python循环提升数个数量级。

2. 进一步优化细节 #

如果数据量极大，还可以通过以下方式压缩开销：

• 维持精简数据类型：生成数据时已用float16，合并后不要随意转换为更高精度类型，减少内存占用和计算量；
• 合并时锁定类型：使用pd.concat时指定dtype参数，避免自动类型升级：

  combined_df = pd.concat([df1, df2, df3], axis=0, dtype="float16")
  

• 提前筛选目标列：若仅需部分令牌列求和，先筛选再分组，减少计算范围：

  target_cols = [col for col in combined_df.columns if col.startswith("tk_")]
  result_df = combined_df.groupby("user_ip")[target_cols].sum()
  

3. 超大规模数据备选方案（内存仍有压力时） #

如果数据量超出单节点内存（题目说明内存充足，此为备选），可使用Dask实现并行计算，语法与pandas兼容：

import dask.dataframe as dd

# 转换为Dask DataFrame，分区数根据CPU核心数调整
dask_df = dd.from_pandas(combined_df, npartitions=8)
# 分组求和并计算出结果
result_df = dask_df.groupby("user_ip").sum().compute()

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