Pandas为何丢弃分组索引？如何保留分组索引？

测试数据准备 #

df_temp = {
   'id': [0,0,0,1,1,1],
   'time' : [1,2,3,1,2,3],
   'values': [3,6,1,4,0,-1]
}
df_temp = pd.DataFrame(df_temp)
df_temp = df_temp.set_index('time')

两个对比案例 #

案例1 #

df_temp.groupby('id')[['values']].apply(lambda _df: _df)

案例2 #

df_temp.groupby('id')[['values']].apply(lambda _df: _df[:-1])

差异原因 #

Pandas的groupby.apply()内置了自动简化机制：当lambda返回的分组DataFrame和原分组的结构（行数、索引、列数）完全一致时，Pandas会自动“扁平化”结果——它认为既然每个分组的结构都和原分组没区别，就没必要保留分组索引，直接合并成和原数据结构类似的输出；而如果返回的DataFrame和原分组结构不一致（比如案例2中每组少了最后一行），Pandas会保留分组索引并将其作为外层索引，形成(id, time)的多层索引，用来明确区分不同组的结果。

修改案例1的方法 #

要在不删减每组数据的前提下保留分组索引，有三种实用方法：

方法1：修改分组返回的索引结构 #

让lambda返回的DataFrame使用多层索引，把分组的id作为外层索引，这样即使行数不变，结构和原分组也不一样，Pandas就不会扁平化结果：

df_temp.groupby('id')[['values']].apply(
    lambda _df: _df.rename(index=lambda t: (_df.name, t))
)

方法2：手动添加分组标识列 #

给每个分组的DataFrame添加id列，改变原分组的列结构，触发Pandas保留分组索引：

df_temp.groupby('id')[['values']].apply(
    lambda _df: _df.assign(id=_df.name)
)

如果之后需要把id转为索引，再执行df.set_index(['id', 'time'])即可。

方法3：直接用pd.concat拼接分组结果 #

这种方法更直观，直接将每个分组按id作为键拼接成多层索引的DataFrame：

pd.concat(
    [group for _, group in df_temp.groupby('id')[['values']]],
    keys=df_temp['id'].unique(),
    names=['id', 'time']
)

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