Pandas中groupby设置observed=False为何引发巨量内存占用？

问题描述 #

使用pandas的groupby操作时，设置observed=False遇到严重内存问题：输入DataFrame规模仅为79860行×790列，但该操作试图分配40GB以上内存。

• observed=True：耗时约10秒，内存占用合理
• observed=False：耗时极长，内存占用远超预期

可复现代码 #

import pandas as pd
import numpy as np
import time

np.random.seed(42)
n_rows = 79860
n_cols = 790

# Create sample data
data = {
    'chrom': ['chr1'] * n_rows,
    'start': range(n_rows),
    'end': range(1, n_rows + 1)
}

# Add sample columns with random values
for i in range(n_cols - 3): 
    data[f'sample_{i}'] = np.random.random(n_rows)

df = pd.DataFrame(data)

# Melt the dataframe
df_melted = pd.melt(df, 
                    id_vars=["chrom","start","end"],
                    var_name="SampleId",
                    value_name="beta")

# Create beta groups
bins = [0, 0.3, 0.7, 1.0]
labels = ["UU", "MU", "MM"]
df_melted['beta_group'] = pd.cut(df_melted['beta'], 
                                bins=bins, 
                                labels=labels, 
                                right=True, 
                                include_lowest=True)

# Test with observed=True (fast, memory efficient)
start = time.time()
result = df_melted.groupby(
    ['chrom','start','end','beta_group'],
    observed=True
).size().reset_index(name='count')
end = time.time()
print(f"Time with observed=True: {(end-start) * 1000:.2f} ms")

# Test with observed=False (very slow, memory intensive)
start = time.time()
result = df_melted.groupby(
    ['chrom','start','end','beta_group'],
    observed=False
).size().reset_index(name='count')
end = time.time()
print(f"Time with observed=False: {(end-start) * 1000:.2f} ms")

环境信息：Python 3.13，Pandas >=2.2.3

原因解析 #

这是Pandas的预期行为，核心差异源于observed参数的逻辑和分组键的组合：

1. observed参数的核心区别 #

• observed=True：仅统计数据中实际存在的分组键组合，忽略那些从未出现过的组合，计算和内存开销仅针对真实存在的分组。
• observed=False：会生成所有分组键取值的笛卡尔积组合——哪怕某些组合在数据中完全不存在。该逻辑仅在分组键包含分类变量时生效（你的场景中beta_group是pd.cut生成的分类类型）。

2. 分组数爆炸的关键原因 #

你的分组键包含start和end两个高基数列：

• start和end各有79860个唯一值（数据中二者一一对应，但Pandas不会自动识别这种关联）
• beta_group有3个固定类别，chrom仅1个取值

当observed=False时，Pandas会生成所有可能的组合：

总可能分组数 = 1（chrom） × 79860（start） × 79860（end） × 3（beta_group） ≈ 1.91×10¹⁰

这样的量级直接导致内存需求飙升——仅存储这些分组的计数就需要数百GB内存，远超常规系统的资源上限，因此出现内存分配失败的情况。

3. observed=True高效的原因 #

observed=True仅保留数据中真实存在的组合：每个start对应唯一的end，再加上该组合下实际出现的beta_group类别，总分组数远小于笛卡尔积的量级，因此内存和计算都保持高效。

解决方案 #

• 若需保留beta_group的所有类别，但不需要start和end的笛卡尔积：可将start和end合并为单个分组键（比如用元组(start, end)或字符串拼接），让Pandas将其视为一个整体维度。
• 优先使用observed=True（Pandas 2.0+对SeriesGroupBy默认启用该设置，DataFrameGroupBy默认是False），除非你明确需要包含分类变量的未出现类别。

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