如何解决睡眠分期数据处理中因列表项缺失导致的pandas KeyError问题

处理睡眠监测文件时，不是所有待分析文件都包含预设sleep_stages列表中的全部睡眠分期项，比如部分文件完全没有Sleep stage N1或其他分期的记录，直接访问不存在的分期索引就会触发KeyError。

• 调用groupby('description')['duration'].sum()得到的tmp是Series类型，它的索引只包含当前文件实际存在的睡眠分期，当你用.loc访问不在索引中的分期项时，就会抛出键错误
• 你需要对缺失的分期项自动填充na，总时长只统计实际存在的分期时长总和

import pandas as pd
import numpy as np

def get_sleep_times(hypno):
    # 注意修正了原代码里分期名换行的语法问题
    sleep_stages=['Sleep stage N1','Sleep stage R','Sleep stage N2', 'Sleep stage N3',  'Sleep stage ?']
    sleep_times = {}
    totsleep_time = 0
    tmp = hypno.groupby('description')['duration'].sum()
    for stage in sleep_stages:
        # 先判断当前分期是否存在，存在则取值统计，不存在填na
        if stage in tmp.index:
            stage_dur = tmp.loc[stage]
            totsleep_time += stage_dur
        else:
            stage_dur = np.nan
        sleep_times['Duration of ' + stage] = stage_dur
    sleep_times['Total Sleep Duration'] = totsleep_time
    return sleep_times

可以直接用pandas的reindex方法自动对齐分期列表，批量填充缺失值，代码更简洁易读：

import pandas as pd
import numpy as np

def get_sleep_times(hypno):
    sleep_stages=['Sleep stage N1','Sleep stage R','Sleep stage N2', 'Sleep stage N3',  'Sleep stage ?']
    tmp = hypno.groupby('description')['duration'].sum()
    # 重索引到预设分期列表，缺失值自动填na
    tmp_aligned = tmp.reindex(sleep_stages, fill_value=np.nan)
    # 批量生成各分期时长的结果字典
    sleep_times = {f'Duration of {stage}': tmp_aligned[stage] for stage in sleep_stages}
    # sum方法会自动跳过na计算总时长
    sleep_times['Total Sleep Duration'] = tmp_aligned.sum()
    return sleep_times

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