pandas自定义函数计算列长拆分DataFrame报KeyError修复方法

问题描述 #

我是编程初学者，目前正在使用pandas完成小型项目，需要实现一个自定义函数：接收一个DataFrame以及该DataFrame中的2个列名作为入参，最终输出两个DataFrame。

原始实现代码 #

dataframe = df

def string_filter(dataframe, dataframest1, dataframest2):
    dataframe0 = dataframe[dataframe[dataframest2].notnull()]
       
    dataframe0[ dataframest1 + ' refined '] = dataframe0[dataframest1].str.len()
    dataframe0[ dataframest2 + ' refined '] = dataframe0[dataframest2].str.len()
    print(dataframe0)
    
    x == dataframe0[ dataframest1 + ' refined ']
    z == dataframe0[ dataframest2 + ' refined ']
   
    dataframe1 = dataframe0[x | z != 1] 
    dataframe2 = dataframe0[x | z == 1] 
    
    return(dataframe1, dataframe2)

string_filter(dataframe, 'c1', 'c2')

触发报错 #

KeyError: ('c1', 'c2')

The above exception was the direct cause of the following exception:

KeyError                                  Traceback (most recent call last)
Input In [28], in <cell line: 1>()
----> 1 string_filter(dataframe, "c1","c2")

错误原因 #

代码一共有4处核心问题，直接触发KeyError的是第一点：

• 变量赋值误用比较运算符==：x == ...和z == ...是逻辑判断语句，不会给x、z变量赋值，后续用x、z做行筛选时，pandas无法识别未定义的变量，会将传入的两个列名组成元组('c1','c2')当作目标列名查询，自然找不到对应列触发KeyError。
• 链式赋值风险：直接在筛选得到的dataframe0上新增列，会触发pandas的链式赋值警告，可能导致新列值写入失败。
• 空值过滤不全：初始步骤只过滤了第二列的空值，第一列如果存在空值，调用.str.len()时会触发空值错误。
• 逻辑运算优先级错误：x | z !=1会优先计算z!=1，再和x做或运算，和预期的判断逻辑不符；且原逻辑用|判断长度等于1的分支，会导致两个输出DataFrame的行出现重叠、遗漏，不满足互斥拆分的要求。

修正后代码 #

def string_filter(dataframe, dataframest1, dataframest2):
    # 同时过滤两列的空值，显式copy避免链式赋值问题
    dataframe0 = dataframe[dataframe[[dataframest1, dataframest2]].notna().all(axis=1)].copy()

    # 计算两列字符串长度
    col1_refined = dataframest1 + ' refined '
    col2_refined = dataframest2 + ' refined '
    dataframe0[col1_refined] = dataframe0[dataframest1].str.len()
    dataframe0[col2_refined] = dataframe0[dataframest2].str.len()
    print(dataframe0)

    # 用=给变量赋值，替换原来的==
    x = dataframe0[col1_refined]
    z = dataframe0[col2_refined]

    # 加括号明确运算优先级，调整拆分逻辑保证两个结果互斥无遗漏
    # dataframe1存任意一列字符串长度不等于1的行
    dataframe1 = dataframe0[(x != 1) | (z != 1)]
    # dataframe2存两列字符串长度都等于1的行
    dataframe2 = dataframe0[(x == 1) & (z == 1)]

    return dataframe1, dataframe2

# 调用方式
# df1, df2 = string_filter(df, 'c1', 'c2')

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