百万行DataFrame分组处理：仅删除分组内存在非NaN行的NaN行

问题场景 #

我有一个百万行规模的数据集，想避免用循环暴力处理。数据集示例如下：

import pandas as pd
import numpy as np

data = {'A': ['I', 'I', 'I', 'I', 'I', 'I', 'S'],
'B': ['J', 'J', 'J', 'J', 'J', 'J', 'J'],
'C': ['K', 'K', 'K', 'K', 'K', 'K', 'K'],
'D': ['L', 'L', 'L', 'L', 'L', 'L', 'T'],
'E': [np.nan, 'O', 'O', np.nan, 'Q', 'Q', np.nan],
'F': ['M', 'M', 'M', 'M', 'M', 'M', 'M'],
'G': ['N', 'N', 'N', np.nan, 'N', 'N', 'N'],
'H': [np.nan, np.nan, 'P', np.nan, np.nan, 'R', np.nan]}
df = pd.DataFrame(data)

期望得到的输出结果：

A  B  C  D    E  F  G     H
0  I  J  K  L    O  M  N     P
1  I  J  K  L    Q  M  N     R
2  S  J  K  T  NaN  M  N  None

尝试过的方法及问题 #

1. 按['A', 'B', 'C', 'D']分组并使用last()聚合：

result_df = df.groupby(['A', 'B', 'C', 'D'], as_index=False, dropna = False).last()

问题：丢失了df['E'] == 'O'的行，同一A/B/C/D组内的O和Q会被合并成最后一个Q的行。

2. 将E加入分组键：

result_df = df.groupby(['A', 'B', 'C', 'D', 'E'], as_index=False, dropna = False).last()

问题：保留了不需要的含NaN的冗余行。

高效解决方案 #

核心思路是分两类处理：有有效E值的组和全为NaN的E组，全程用矢量化操作，避免循环，适合百万行数据。

方案代码

# 1. 标记每个A/B/C/D组是否存在非NaN的E值
df['has_valid_e'] = df.groupby(['A', 'B', 'C', 'D'])['E'].transform(lambda x: x.notna().any())

# 2. 处理有有效E值的组：按A/B/C/D/E分组取last，保留每个E有效值对应的唯一行
valid_e_groups = df[df['has_valid_e'] & df['E'].notna()].groupby(['A', 'B', 'C', 'D', 'E'], as_index=False).last()

# 3. 处理全为NaN的E组：按A/B/C/D分组取last，保留该组唯一行
nan_only_groups = df[~df['has_valid_e']].groupby(['A', 'B', 'C', 'D'], as_index=False).last()

# 4. 合并两部分结果，删除辅助列
final_result = pd.concat([valid_e_groups, nan_only_groups], ignore_index=True).drop('has_valid_e', axis=1)

方案说明

• 辅助列标记：用transform快速给每行标记所在组是否有有效E值，这一步是矢量化操作，效率极高。
• 分治处理：
  • 对于有有效E的组，按A/B/C/D/E分组，确保每个E的有效值对应独立行，同时用last()获取其他列的有效数据（比如H列的P和R）。
  • 对于全NaN的E组，直接按A/B/C/D分组取唯一行，避免冗余。
• 合并结果：将两部分数据合并后得到预期输出，全程无循环，处理百万级数据毫无压力。

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