Python Dataframe列存在性检查及缺失列创建报错问题解决

需求 #

希望DataFrame包含两列，列名通过变量指定，若某列缺失则创建该列。

给定代码 #

# Columns dataframe or series. It contains names of the actual columns
# I get below information from another source
cols_df = pd.Series(index=['main_col'],data=['A'])
# This also the dataframe I get from another source
df = pd.DataFrame(data={'A':[10,20,30]})
# My job is see if the given dataframe has two columns
if (cols_df['main_col'] in df.columns) and (cols_df['aux_col'] not in df.columns):
  print("Aux col not in df")
  cols_df['aux_col'] = 'B'
  df[cols_df['aux_col']] = df[cols_df['main_col']]-10

当前报错 #

_check_indexing_error will raise

KeyError: 'aux_col'

期望输出 #

df = 
    A   B
0  10   0
1  20  10
2  30  20

cols_df = 
main_col    A
aux_col     B

报错原因 #

原代码中直接访问cols_df['aux_col']来判断列是否存在，但初始的cols_df并没有aux_col这个索引，因此触发KeyError。正确逻辑应该是先处理aux_col未定义的情况，再判断对应的列是否在DataFrame中。

修复方案 #

方案一：用get()方法避免KeyError #

通过cols_df.get('aux_col')获取辅助列名，不存在时返回None，这样就能安全判断：

import pandas as pd

cols_df = pd.Series(index=['main_col'], data=['A'])
df = pd.DataFrame(data={'A': [10, 20, 30]})

aux_col_name = cols_df.get('aux_col')
if cols_df['main_col'] in df.columns and (aux_col_name is None or aux_col_name not in df.columns):
    print("Aux col not in df")
    cols_df['aux_col'] = 'B'
    df[cols_df['aux_col']] = df[cols_df['main_col']] - 10

方案二：先确保aux_col存在于cols_df中 #

先给cols_df添加aux_col索引（如果不存在），再判断列是否在DataFrame里：

import pandas as pd

cols_df = pd.Series(index=['main_col'], data=['A'])
df = pd.DataFrame(data={'A': [10, 20, 30]})

if 'aux_col' not in cols_df.index:
    cols_df['aux_col'] = 'B'

if cols_df['main_col'] in df.columns and cols_df['aux_col'] not in df.columns:
    print("Aux col not in df")
    df[cols_df['aux_col']] = df[cols_df['main_col']] - 10

验证结果 #

运行任意一种修复方案后，都能得到期望的输出：

• 生成的df：

  A   B
  0  10   0
  1  20  10
  2  30  20
  

• 更新后的cols_df：

  main_col    A
  aux_col     B
  dtype: object
  

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