Pandas中DataFrame.groupby采样后索引异常的解决方法

问题背景 #

现有如下pandas.DataFrame数据：

labels               texts
0         labelA  Some Text 12345678
1         labelA  Some Text 12345678
2         labelA  Some Text 12345678
3         labelA  Some Text 12345678
4         labelB  Some Text 12345678
5         labelB  Some Text 12345678
6         labelB  Some Text 12345678
7         labelC  Some Text 12345678
8         labelC  Some Text 12345678
9         labelC  Some Text 12345678
10        labelC  Some Text 12345678
11        labelC  Some Text 12345678
12        labelC  Some Text 12345678

目标是按labels分组后每组抽取2条样本，执行以下代码：

grouped = df.groupby('labels')
result = grouped.apply(lambda x: x.sample(n=2))
print(result)

得到的结果出现多层索引且子索引异常：

labels               texts
labels
labelA    0         labelA  Some Text 12345678
          0         labelA  Some Text 12345678
          0         labelB  Some Text 12345678
          0         labelB  Some Text 12345678
          0         labelC  Some Text 12345678
          0         labelC  Some Text 12345678

期望输出为带连续单索引的格式：

labels               texts
0         labelA  Some Text 12345678
1         labelA  Some Text 12345678
2         labelB  Some Text 12345678
3         labelB  Some Text 12345678
4         labelC  Some Text 12345678
5         labelC  Some Text 12345678

尝试使用result.dropout(0).reset_index()处理后，结果多出了index列（注：dropout是拼写错误，正确方法需通过reset_index参数控制）：

index         labels               texts
0        0         labelA  Some Text 12345678
1        0         labelA  Some Text 12345678
2        0         labelB  Some Text 12345678
3        0         labelB  Some Text 12345678
4        0         labelC  Some Text 12345678
5        0         labelC  Some Text 12345678

解决方案 #

方法1：直接使用groupby.sample（推荐，Pandas 1.1.0及以上版本支持） #

Pandas原生提供分组抽样方法，无需手动apply，自动规避索引问题：

result = df.groupby('labels').sample(n=2).reset_index(drop=True)
print(result)

• groupby.sample直接返回每组抽样后的行，不会生成多层索引
• reset_index(drop=True)将索引重置为连续单索引，同时丢弃旧索引列

方法2：修复原有apply后的索引处理 #

若使用旧版Pandas，可修改索引处理逻辑：

grouped = df.groupby('labels')
result = grouped.apply(lambda x: x.sample(n=2)).reset_index(level=0, drop=True).reset_index(drop=True)
print(result)

• reset_index(level=0, drop=True)：移除分组产生的第一层索引（labels层级）
• 再次调用reset_index(drop=True)：将剩余的原DataFrame索引替换为连续单索引

也可在apply阶段直接重置每组索引：

result = grouped.apply(lambda x: x.sample(n=2).reset_index(drop=True)).reset_index(drop=True)
print(result)

错误原因说明 #

• groupby.apply默认保留分组键作为第一层索引，同时保留原DataFrame索引作为第二层，因此产生多层索引
• 之前尝试的reset_index()未指定参数，会将所有层级索引转为普通列，导致多出index列

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