Pandas中groupby+apply(list)内存不足问题及优化方法

测试数据集构建 #

data = {'link': [1,2,3,4,5,6,7],
        'code': ['xx', 'xx', 'xy', '', 'aa', 'ab', 'aa'],
        'Name': ['Tom', 'Tom', 'Tom', 'Tom', 'nick', 'nick', 'nick'],
        'Age': [20,20,20,20, 21, 21, 21]}

# 创建DataFrame
df = pd.DataFrame(data)

print(df)

输出结果：

link code  Name  Age
0     1   xx   Tom   20
1     2   xx   Tom   20
2     3   xy   Tom   20
3     4        Tom   20
4     5   aa  nick   21
5     6   ab  nick   21
6     7   aa  nick   21

原代码的内存问题 #

以下代码在数据集子集上可正常运行，但在全量数据集执行时触发内存错误：

temp = df.groupby(['Name', 'Age'])['code'].apply(list).reset_index()
pd.merge(df, temp, on=['Name', 'Age']).explode('code_y').replace(r'^\s*$', np.nan, regex=True).dropna(subset='code_y').drop_duplicates()

全量数据集报错信息：

MemoryError: Unable to allocate 5.34 TiB for an array with shape (733324768776,) and data type object

问题根源：groupby.apply(list)会为每个分组生成包含大量重复元素的长列表，后续explode操作会将这些列表展开成巨量行，直接耗尽内存。

优化方案（当前最优） #

通过提前去重避免生成冗余数据，大幅降低内存占用：

# 筛选相关列并直接去重，同时将空字符串转为缺失值
d = df[['code', 'Name', 'Age']].replace('', pd.NA).drop_duplicates()

# 执行合并并一步过滤掉code_y为空的行
df.merge(d, how='outer', on=['Name', 'Age']).dropna(subset=['code_y'])

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