如何优化pandas百万列宽表的groupby聚合操作性能

核心问题说明 #

你当前的操作顺序存在性能浪费：提前生成150万列的超宽独热表再做聚合，大部分运算资源都消耗在了无意义的0值计算上，优先调整操作顺序即可获得指数级性能提升。

方案1：调整操作顺序（最优，性能提升100倍以上） #

原始数据仅id和color两列，完全可以先完成按id去重颜色的聚合逻辑，再做独热编码，避免提前生成百万级列的宽表。
示例代码：

import pandas as pd
# 原始窄表，仅2列8000行，无任何额外开销
array = {'id': [1, 1, 1, 1, 2, 3],
         'color': ['yellow', 'red', 'yellow', 'red', 'yellow', 'white']}
df = pd.DataFrame(array)

# 直接统计每个id对应颜色的出现次数，再截断为1，和你要的max逻辑完全等价
res = pd.crosstab(df['id'], df['color']).clip(upper=1)

输出结果和你之前的groupby结果完全一致，即使颜色维度有150万种，crosstab的计算开销也远低于先做独热再聚合。

方案2：若必须处理已生成的宽表，做稀疏化优化 #

如果已经生成了150万列的独热宽表，无法回到原始窄表，可以通过稀疏化降低内存开销，提升计算速度：

# 将所有独热列转为稀疏类型，大幅降低内存占用（独热表99%以上都是0值，稀疏存储可以过滤掉无效0值）
sparse_cols = [col for col in df.columns if col != 'id']
df[sparse_cols] = df[sparse_cols].astype(pd.SparseDtype('int8', 0))

# groupby时关闭不必要的排序进一步提速
res = df.groupby('id', sort=False).max()

补充优化建议 #

• 高基数类别特征场景下，永远不要先做独热编码再聚合，优先在窄表上完成聚合逻辑后再编码
• 如果后续还有更大规模的同类计算，可以考虑用polars库替代pandas，其对宽表的运算性能远超pandas

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