基于条件提取DataFrame记录并生成新列的Python代码修正问题

原始问题 #

需要从Excel读取的DataFrame中基于条件提取记录并生成新列，原始代码如下：

import pandas as pd
df = pd.read_excel("C:/Users/0403 - Copy (3).xlsx")

def add_action(row): 
    if row["FILE TYPE"] == 'NaN': 
        return row['Row Labels']   

df = df.assign(Label_header=df.apply(add_action, axis=1)) 

输入DataFrame：

|FILE TYPE |Row Labels               |TOTAL COUNT|    
0      | NaN      | CHANGE                  | 668       |
1      | 87I      | ei.cfc.h87.n5.20211     | 98        |   
2      | 87P      | ep.cfc.m87.n5.2023      | 570       |   
3      | NaN      | CHN                     | 5642      |   
4      | 87P      | NMMCMS.AC.2021_1_R.txt  | 1         |   
5      | 87P      | NS.AC.201_1_R.txt       | 1         |

期望输出：

|FILE TYPE |Row Labels               |TOTAL COUNT| Label_header |
0      | 87I      | ei.cfc.h87.n5.20211     | 98        | CHANGE       |
1      | 87P      | ep.cfc.m87.n5.2023      | 570       | CHANGE       |
2      | 87P      | NMMCMS.AC.2021_1_R.txt  | 1         | CHN          |
3      | 87P      | NS.AC.201_1_R.txt       | 1         | CHN          |

问题分析 #

1. 空值判断错误：pandas中的空值是pd.NA/np.nan，不是字符串'NaN'，用== 'NaN'无法匹配真实空值。
2. 逻辑不完整：原始代码仅为空值行赋值Label_header，未将分组头向下填充到后续非空行，也未过滤掉原始空值行。

修正后的代码 #

import pandas as pd

# 读取Excel数据
df = pd.read_excel("C:/Users/0403 - Copy (3).xlsx")

# 为FILE TYPE为空的行设置Label_header为Row Labels，其余行留空
df['Label_header'] = df.apply(
    lambda row: row['Row Labels'] if pd.isna(row['FILE TYPE']) else None,
    axis=1
)

# 向前填充Label_header，将分组头传递给后续行
df['Label_header'] = df['Label_header'].ffill()

# 过滤掉FILE TYPE为空的行，得到最终结果
result_df = df[~pd.isna(df['FILE TYPE'])].reset_index(drop=True)

print(result_df)

输出结果 #

FILE TYPE               Row Labels  TOTAL COUNT Label_header
0       87I      ei.cfc.h87.n5.20211           98       CHANGE
1       87P       ep.cfc.m87.n5.2023          570       CHANGE
2       87P  NMMCMS.AC.2021_1_R.txt            1          CHN
3       87P     NS.AC.201_1_R.txt             1          CHN

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