Pandas DataFrame列求和时出现索引相关KeyError报错

问题复现 #

执行如下透视表代码生成统计DataFrame：

CA_HousingTrend = CA_HousingTrend_temp.pivot_table(index='YEAR',columns='UNITSSTR', aggfunc='size')

需要统计多户住宅单元总量，指定待求和列后执行按行求和：

cols = ['05', '06']
CA_HousingTrend['sum_stats'] = CA_HousingTrend[cols].sum(axis=1)

运行后抛出KeyError：

KeyError: "None of [Index(['05', '06'], dtype='object', name='UNITSSTR')] are in the [columns]"

完整报错栈：

Traceback (most recent call last):
File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/code.py", line 90, in runcode
exec(code, self.locals)
File "", line 5, in <module>
File "/Users/alexandramaxim/Documents/Py/lib/python3.10/site-packages/pandas/core/frame.py", line 3511, in __getitem__
indexer = self.columns._get_indexer_strict(key, "columns")
File "/Users/alexandramaxim/Documents/Py/lib/python3.10/site-packages/pandas/core/indexes/base.py", line 5782, in _get_indexer_strict
self._raise_if_missing(keyarr, indexer, axis_name)
File "/Users/alexandramaxim/Documents/Py/lib/python3.10/site-packages/pandas/core/indexes/base.py", line 5842, in _raise_if_missing
raise KeyError(f"None of [{key}] are in the [{axis_name}]")

根因 #

传入的列名和DataFrame实际存储的列名类型或值不匹配：
pivot_table生成列时会保留UNITSSTR原字段的类型和值，传入的是带前导零的字符串'05'、'06'，但实际列名要么是整数类型的5、6，要么是不带前导零的字符串'5'、'6'，无法匹配到对应列。

修复方案 #

第一步先确认实际列名的类型和值，执行：

print(CA_HousingTrend.columns)
print(CA_HousingTrend.columns.dtype)

根据输出结果二选一调整：

• 如果输出列是整数5、6，修改cols定义为整数：

cols = [5, 6]
CA_HousingTrend['sum_stats'] = CA_HousingTrend[cols].sum(axis=1)

• 如果输出列是字符串'5'、'6'（无前置零），修改cols为对应字符串：

cols = ['5', '6']
CA_HousingTrend['sum_stats'] = CA_HousingTrend[cols].sum(axis=1)

*通用避坑技巧：选列前可以统一把列名转为字符串，避免类型不匹配问题：

CA_HousingTrend.columns = CA_HousingTrend.columns.astype(str)

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