Pandas分组自定义函数状态赋值异常：结果不符合预期排查

需求说明 #

按Site分组检查指定前缀的出现次数，规则如下：

• 指定的4个前缀['A0L-0', 'A1L-0', 'Z0L-0', 'Z1L-0']每个在每个Site下恰好出现1次 → 状态标记为OK
• 任意前缀出现次数>1 → 状态标记为NOK
• 任意前缀缺失（出现次数=0） → 状态标记为Not exist

数据示例 #

import pandas as pd

data = [
    ['KA7022', 'A0L-0'],
    ['KA7022', 'A1L-0'],
    ['KA7022', 'Z0L-0'],
    ['KA7022', 'Z1L-0'],
    ['TA5251', 'A1L-0'],
    ['TA5251', 'Z0L-0'],
    ['TA5251', 'Z1L-0'],
    ['AL7777', 'A0L-0'],
    ['AL7777', 'A0L-0'],
    ['AL7777', 'A1L-0'],
    ['AL7777', 'Z0L-0'],
    ['AL7777', 'Z1L-0']
]
df_input = pd.DataFrame(data, columns=['Site', 'Name pref'])

现有实现代码 #

def my_func(x, name_pref_list):
    count = 0
    for name_pref in name_pref_list:
        x = x[x['Name pref'] == name_pref]
        # 若出现次数为0则返回Not exist
        if len(x) == 0:
            return 'Not exist'
        # 若出现次数>1则返回NOK
        if len(x) > 1:
            return 'NOK'
        count = count + 1
    # 若所有前缀都符合条件则返回OK
    if len(name_pref_list) == count:
        return 'OK'
    else:
        return 'NOK'

df_result = df_input.groupby('Site').\
                     apply(lambda x: my_func(x, ['A0L-0', 'A1L-0', 'Z0L-0', 'Z1L-0'])).\
                     reset_index(name='Status')

结果对比 #

实际运行结果 #

Site    Status
0   AL7777  NOK
1   KA7022  Not exist
2   TA5251  Not exist

预期结果 #

Site    Status
0   AL7777  NOK
1   KA7022  OK
2   TA5251  Not exist

错误原因 #

函数my_func中存在逻辑错误：循环时每次执行x = x[x['Name pref'] == name_pref]，会覆盖原始的分组数据x。以KA7022为例：

1. 第一次循环处理A0L-0，x被过滤为仅包含A0L-0的行
2. 第二次循环处理A1L-0时，当前x中已经没有A1L-0的数据，判断len(x)==0，直接返回Not exist，导致后续逻辑无法执行

修正后的代码 #

修改函数逻辑，不再覆盖原始分组数据，而是基于原始数据统计每个前缀的出现次数：

def my_func(x, name_pref_list):
    for name_pref in name_pref_list:
        # 统计当前前缀在分组中的出现次数
        cnt = (x['Name pref'] == name_pref).sum()
        if cnt == 0:
            return 'Not exist'
        if cnt > 1:
            return 'NOK'
    # 所有前缀都满足恰好出现1次的条件
    return 'OK'

df_result = df_input.groupby('Site').\
                     apply(lambda x: my_func(x, ['A0L-0', 'A1L-0', 'Z0L-0', 'Z1L-0'])).\
                     reset_index(name='Status')

修正后运行结果 #

Site    Status
0   AL7777  NOK
1   KA7022  OK
2   TA5251  Not exist

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