使用pandas处理含NA值的分组取首条最大值行报错求助

问题场景 #

原DataFrame定义 #

import pandas as pd
df = pd.DataFrame(
    {
        1: ['a', 'a', 'a', 'b', 'b', 'c', 'd'],
        2: ['x', 'y', 'y', 'x', 'y', pd.NA, 'x'],
        3: [10, 8, 8, 10, 7, pd.NA, 0],
        4: [4, 4, 5, 7, 8, pd.NA, 0]
    },
)

原始数据展示 #

1     2     3     4
0  a     x    10     4
1  a     y     8     4
2  a     y     8     5
3  b     x    10     7
4  b     y     7     8
5  c  <NA>  <NA>  <NA>
6  d     x     0     0

需求说明 #

按列1和列2分组，保留每组中列3为最大值的首条行，同时完整保留含NA值的行，预期结果如下：

1     2     3     4
0  a     x    10     4
1  a     y     8     4
3  b     x    10     7
4  b     y     7     8
5  c  <NA>  <NA>  <NA>
6  d     x     0     0

尝试的方法及报错 #

已执行列类型转换：

df[3] = pd.to_numeric(df[3])

尝试用groupby([1, 2], dropna=False)分组，结合loc和idxmax实现需求：

df2: pd.DataFrame = df.loc[
    df.groupby([1, 2], dropna=False)[3].idxmax()
]

触发异常：KeyError: '[nan] not in index'

问题原因 #

这个错误确实和pandas issue #52234直接相关——在pandas 2.0.x版本中，当分组包含NA值且设置dropna=False时，idxmax()会返回NaN作为该分组的索引值，但df.loc[]无法识别NaN作为行索引，因此抛出KeyError。另外，分组后直接取max()会丢失其他列数据，不符合需求。

解决方案 #

以下两种方法均可解决问题：

方法1：用transform标记最大值首行 #

先给每个分组标记出列3为最大值的行，再筛选并保留每组首条，同时保留NA行：

# 标记每组中列3为最大值的行，全NA分组直接标记为True
df['is_max'] = df.groupby([1, 2], dropna=False)[3].transform(
    lambda x: x == x.max() if not x.isna().all() else True
)
# 筛选标记行，每个分组只取首条
result = df[df['is_max']].groupby([1, 2], dropna=False).first().reset_index(drop=True)
# 删除临时标记列
result = result.drop(columns='is_max')

方法2：分离NA与非NA分组处理 #

先处理非NA分组的最大值首行，再单独提取NA行合并：

# 分离非NA行和全NA行
non_na_df = df.dropna(subset=[3])
na_df = df[df[3].isna()]

# 非NA分组取最大值首行
non_na_result = non_na_df.loc[non_na_df.groupby([1, 2])[3].idxmax()]

# 合并结果并按原索引排序
result = pd.concat([non_na_result, na_df]).sort_index()

两种方法都能得到预期结果，且规避了idxmax()返回NaN导致的KeyError问题。

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