新版本pandas执行fillna/ffill/bfill后如何保留分组列与索引？

pandas 2.1及之后版本调整了分组变换的默认行为，调用ffill/bfill/fillna等方法时默认不会返回分组所用的列，可通过以下三种常用方案解决：

方法1：groupby时指定group_keys=True #

开启分组键保留参数，填充后重置索引即可拿到完整列：

import pandas as pd
import io

data = """one;two;three
1;1;10
1;1;nan
1;1;nan
1;2;nan
1;2;20
1;2;nan
1;3;nan
1;3;nan"""

df = pd.read_csv(io.StringIO(data), sep=";")
# 开启group_keys保留分组键，重置索引时把分组键转为普通列
res = df.groupby(['one','two'], group_keys=True).ffill().reset_index(names=['原索引', 'one', 'two']).drop('原索引', axis=1)
print(res)

输出结果：

one  two  three
0    1    1   10.0
1    1    1   10.0
2    1    1   10.0
3    1    2    NaN
4    1    2   20.0
5    1    2   20.0
6    1    3    NaN
7    1    3    NaN

方法2：与原表分组列直接拼接 #

无需修改分组逻辑，直接拼接原表的分组列与填充结果，兼容性最好：

# 只对需要填充的列做分组操作，再和原表的分组列合并
res = pd.concat([
    df[['one', 'two']],
    df.groupby(['one','two'])[['three']].ffill()
], axis=1)
print(res)

输出结果与方法1完全一致。

方法3：用transform直接修改原列 #

如果不需要保留原始填充前的列，可直接通过transform方法赋值到原列：

df['three'] = df.groupby(['one','two'])['three'].transform(lambda x: x.ffill())
print(df)

如果是bfill或者指定值填充，替换ffill()为对应方法即可，比如按分组均值填充：

df['three'] = df.groupby(['one','two'])['three'].transform(lambda x: x.fillna(x.mean()))

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