DataFrame.stack()内存占用过高原因及优化方案咨询

问题描述 #

使用memray分析代码时发现，仅约5MB的DataFrame调用df.stack()时，竟分配了22MB内存。代码示例如下：

import numpy as np
import pandas as pd

columns = list('abcdefghijklmnopqrstuvwxyz')
df = pd.DataFrame(np.random.randint(0,100,size=(1000, 26*26)), columns=pd.MultiIndex.from_product([columns, columns]))
print(df.memory_usage().sum()) # 5408128, ~5MB
df.stack() # 重塑维度: (1000,26*26) -> (1000*26,26)

经分析，内存消耗中30%来自dropna操作，剩余70%源于三次底层数组的重新分配以完成重塑。现需明确该现象的原因，以及是否应改用原生numpy.reshape或其他轻量化优化方法。

原因解析 #

• 默认dropna的冗余开销：stack()方法默认开启dropna=True，即便数据中没有缺失值，pandas仍会执行缺失值检查与相关逻辑，这部分操作会产生额外的内存分配。
• 多层索引重塑的多次数组拷贝：处理MultiIndex列时，stack()需要重新组织数据结构，涉及多次数组拆分与重组。pandas不会直接修改原数组，而是创建多个中间数组来完成层级转换，多次内存分配累加后，总占用就远超过原DataFrame的大小。此外，重塑后的数据会保留完整的索引元数据，这部分也会占用额外内存。

优化方案 #

1. 关闭自动dropna #

如果确认数据中没有缺失值，直接调用stack(dropna=False)，可以省去dropna相关的内存开销，直接减少约30%的内存占用。

2. 用numpy.reshape替代（适合规则重塑场景） #

你的重塑逻辑是完全规则的维度转换，完全可以用numpy的reshape来实现，避开pandas索引转换带来的额外开销：

# 提取原数据的numpy数组做reshape
reshaped_values = df.values.reshape(1000*26, 26)
# 构造新的索引和列名
new_index = pd.MultiIndex.from_product([df.index, columns], names=[df.index.name, columns[0]])
new_df = pd.DataFrame(reshaped_values, index=new_index, columns=columns)

numpy的reshape在满足连续数组条件时会返回视图（而非拷贝），即便需要拷贝，也仅需一次内存分配，整体内存占用远低于stack()。

3. 其他轻量化技巧 #

• 如果不需要保留原索引结构，可以直接用简单的RangeIndex构造新DataFrame，进一步降低元数据的内存开销。
• 若处理超大规模数据，可考虑用dask.dataframe做延迟计算，避免一次性分配大量内存，但小数据场景下没必要。

内容的提问来源于stack exchange，提问作者C. Claudio