排查基于特定前缀生成DataFrame新列的Python代码错误

问题场景 #

现有如下Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({
    'ID': [300, 1000, 450, 111, 550],
    'Folder Name': ['ABC 12345', pd.NA, 'AML 2233', 'ABC 2234', 'AML 3312'],
    'Country': ['CANADA', 'USA', 'USA', 'USA', 'AFRICA']
})

需要新增Folder Name - ABC和Folder Name - AML两列：仅保留对应前缀的原Folder Name值，其余情况（包括原列值为NaN）设为NaN，最终输出如下：

ID Folder Name Country Folder Name - ABC Folder Name - AML
0    300   ABC 12345  CANADA        ABC 12345               NaN
1   1000         NaN     USA               NaN               NaN
2    450    AML 2233     USA               NaN        AML 2233
3    111    ABC 2234     USA        ABC 2234               NaN
4    550    AML 3312  AFRICA               NaN        AML 3312

错误代码分析 #

你尝试的代码存在3个问题：

df_['Folder Name - ABC'] = df['Folder Name'].apply(lambda x: x.str.startswith('ABC',na = False))

1. 未定义变量df_：直接对df_赋值会触发NameError，应该直接在原df上操作，或者先复制df给df_（比如df_ = df.copy()）。
2. 冗余的apply调用：str.startswith是Series原生方法，无需用apply包裹，直接调用即可。
3. 返回值不符合需求：这段代码返回的是布尔值（True/False），而不是原字符串或NaN，和预期的输出格式不符。

正确解决方案 #

方法一：使用where方法（推荐） #

where方法会保留满足条件的原数值，不满足条件的自动替换为NaN，完美匹配需求：

# 新增ABC前缀匹配列
df['Folder Name - ABC'] = df['Folder Name'].where(df['Folder Name'].str.startswith('ABC', na=False))
# 新增AML前缀匹配列
df['Folder Name - AML'] = df['Folder Name'].where(df['Folder Name'].str.startswith('AML', na=False))

方法二：自定义函数+apply #

如果需要更灵活的逻辑，可以自定义函数后用apply调用：

def filter_prefix(value, target_prefix):
    if pd.notna(value) and value.startswith(target_prefix):
        return value
    return pd.NA

df['Folder Name - ABC'] = df['Folder Name'].apply(filter_prefix, target_prefix='ABC')
df['Folder Name - AML'] = df['Folder Name'].apply(filter_prefix, target_prefix='AML')

两种方法执行后都能得到你期望的输出结果。

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