Pandas中groupby聚合无法保留原始dtype，求解决方法

问题场景 #

创建int8类型的DataFrame：

import pandas as pd
import numpy as np

d = pd.DataFrame({c: np.random.choice(range(10), size=20) 
                  for c in list('abcd')}).astype('int8')
d.groupby(d.columns.tolist(), as_index=False).size().info()

执行后输出显示所有原int8列（a、b、c、d）和新增的size列都被转换为int64类型，内存占用大幅增加。该问题出现在Pandas 1.3.5和1.5.3版本，2.0.2及以上版本已修复。

解决方法 #

方法1：手动还原原始数据类型 #

先保存原始DataFrame的列类型，聚合后将分组键列还原为原类型：

# 保存原始列的类型映射
original_dtypes = d.dtypes

# 执行groupby聚合
grouped_result = d.groupby(d.columns.tolist(), as_index=False).size()

# 还原分组键列的原始类型
grouped_result[original_dtypes.index] = grouped_result[original_dtypes.index].astype(original_dtypes)

# 验证结果
grouped_result.info()

此方法直接可控，适合无法升级Pandas的场景。

方法2：升级Pandas到2.0+版本 #

Pandas 2.0及以上版本已修复该类型提升的bug，升级后无需额外处理即可保留原始int8类型，直接运行原代码即可：

d.groupby(d.columns.tolist(), as_index=False).size().info()

补充优化 #

size列默认会是int64类型，若分组数量在小整数范围内（比如不超过127），可进一步压缩内存：

grouped_result['size'] = grouped_result['size'].astype('int8')

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