在DataFrame中新增收入分类列时触发KeyError: 'Salary'的咨询

问题背景 #

• 现有代码：

from tensorflow import keras
import pandas as pd
df=pd.read_csv('/Salary_Data.csv',header=None)

• 需求：给DataFrame新增highincome列，规则是Salary列数值低于100000时赋值0，大于等于100000时赋值1
• 尝试的代码（来自Google Bard）：

df['highincome'] = (df['Salary'] >= 100000) | (df['Salary'] > 100000)

• 运行后触发错误：

KeyError Traceback (most recent last)
/usr/local/lib/python3.10/dist-packages/pandas/core/indexes/base.py in get_loc(self, key, method, tolerance)
3801 try:
-> 3802 return self._engine.get_loc(casted_key)
3803 except KeyError as err:

4 frames
pandas/_libs/index_class_helper.pxi in pandas._libs.index.Int64Engine._check_type()

KeyError: 'Salary'

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

KeyError Traceback (most recent last)
/usr/local/lib/python3.10/dist-packages/pandas/core/indexes/base.py in get_loc(self, key, method, tolerance)
3802 return self._engine.get_loc(casted_key)
3803 except KeyError as err:
-> 3804 raise KeyError(key) from err
3805 except TypeError:
3806 # If we have a listlike key, _check_indexing_error will raise

KeyError: 'Salary'

• 疑问：不清楚KeyError: 'Salary'是什么意思

问题解析 #

1. KeyError: 'Salary'的含义：当前DataFrame里不存在名为Salary的列，用df['Salary']去访问自然会报错。
2. 无Salary列的原因：读取CSV时加了header=None，这个参数会告诉pandas“CSV没有表头行”，于是把原本是列名的第一行当成了普通数据行，最终DataFrame的列名变成了0、1这类数字索引，自然找不到Salary这个列名。

解决方法 #

1. 修正CSV读取代码 #

去掉header=None，让pandas自动识别CSV第一行为列名：

df = pd.read_csv('/Salary_Data.csv')

不确定列名的话，可以先执行print(df.columns)查看所有列，确认Salary存在。

2. 简化highincome列的生成逻辑 #

Bard给的代码逻辑重复，(df['Salary'] >= 100000) | (df['Salary'] > 100000)完全等价于df['Salary'] >= 100000，直接把布尔值转成整数即可（True对应1，False对应0）：

df['highincome'] = (df['Salary'] >= 100000).astype(int)

完整修正代码 #

import pandas as pd
# 如果暂时用不到tensorflow，可以注释掉下面一行
# from tensorflow import keras

df = pd.read_csv('/Salary_Data.csv')
df['highincome'] = (df['Salary'] >= 100000).astype(int)

内容的提问来源于stack exchange，提问作者오상윤