Python中基于CUI分组生成新列遇IndexError报错求助

原始数据集 #

term            code            source  term_normlz     CUI         CODE        SAB     TTY     STR
0   B-cell lymphoma meddra:10003899 meddra  b-cell lymphoma C0079731    MTHU019696  OMIM    PTCS    b-cell lymphoma
1   B-cell lymphoma meddra:10003899 meddra  b-cell lymphoma C0079731    10003899    MDR     PT      b-cell lymphoma
2   Astrocytoma     meddra:10003571 meddra  astrocytoma     C0004114    10003571    MDR     PT      astrocytoma
3   Astrocytoma     meddra:10003571 meddra  astrocytoma     C0004114    D001254     MSH     MH      astrocytoma

期望输出 #

term            code            source  term_normlz     CUI         OMIM_CODE       OMIM_TTY        OMIM_STR          MDR_CODE   MDR_TTY MDR_STR          MSH_CODE   MSH_TTY MSH_STR
0   B-cell lymphoma meddra:10003899 meddra  b-cell lymphoma C0079731    MTHU019696      PTCS            b-cell lymphoma  10003899   PT      b-cell lymphoma  NA        NA      NA
2   Astrocytoma     meddra:10003571 meddra  astrocytoma     C0004114    NA              NA              NA               10003571   PT      astrocytoma      D001254    MH      astrocytoma

现有代码及报错 #

执行代码 #

HC_subset_umls['OMIM_CODE'] = (
    HC_subset_umls['CUI']
    .map(
        HC_subset_umls
        .groupby('CUI')
        .apply(lambda x: x.loc[x['SAB'].isin(['OMIM']), 'CODE'].values[0])
    )
)

HC_subset_umls['OMIM_TERM'] = (
    HC_subset_umls['CUI']
    .map(
        HC_subset_umls
        .groupby('CUI')
        .apply(lambda x: x.loc[x['SAB'].isin(['OMIM']), 'STR'].values[0])
    )
)

HC_subset_umls['OMIM_TTY'] = (
    HC_subset_umls['CUI']
    .map(
        HC_subset_umls
        .groupby('CUI')
        .apply(lambda x: x.loc[x['SAB'].isin(['OMIM']), 'TTY'].values[0])
    )
)

HC_subset_umls = HC_subset_umls[~(HC_subset_umls['SAB'].isin(['OMIM']))]

报错信息 #

IndexError: index 0 is out of bounds for axis 0 with size 0

问题原因 #

报错源于部分CUI对应的分组中没有SAB='OMIM'的行（比如示例中的C0004114），此时x.loc[x['SAB'].isin(['OMIM']), 'CODE'].values返回空数组，强行取索引0就会触发越界错误。

解决方案 #

方案1：使用pivot_table（推荐，更简洁） #

利用透视表直接完成分组扩展，自动处理空值：

# 定义需要保留的基础列
base_cols = ['term', 'code', 'source', 'term_normlz', 'CUI']
# 按基础列分组，以SAB为列透视CODE/TTY/STR字段
pivoted = HC_subset_umls.pivot_table(
    index=base_cols,
    columns='SAB',
    values=['CODE', 'TTY', 'STR'],
    aggfunc='first'  # 每个CUI+SAB组合取第一个值，避免重复
).reset_index()

# 重命名多层列索引为期望格式（如OMIM_CODE）
pivoted.columns = [f'{col[1]}_{col[0]}' if col[1] else col[0] for col in pivoted.columns]

# 将空值替换为NA
pivoted = pivoted.fillna('NA')

方案2：修复原有代码逻辑 #

在lambda中增加空值判断，避免索引越界：

# 定义通用取值函数
def get_group_value(group, target_sab, target_col):
    subset = group.loc[group['SAB'] == target_sab, target_col]
    return subset.values[0] if not subset.empty else 'NA'

# 生成OMIM相关列
HC_subset_umls['OMIM_CODE'] = HC_subset_umls['CUI'].map(
    HC_subset_umls.groupby('CUI').apply(lambda x: get_group_value(x, 'OMIM', 'CODE'))
)
HC_subset_umls['OMIM_TTY'] = HC_subset_umls['CUI'].map(
    HC_subset_umls.groupby('CUI').apply(lambda x: get_group_value(x, 'OMIM', 'TTY'))
)
HC_subset_umls['OMIM_STR'] = HC_subset_umls['CUI'].map(
    HC_subset_umls.groupby('CUI').apply(lambda x: get_group_value(x, 'OMIM', 'STR'))
)

# 生成MDR相关列
HC_subset_umls['MDR_CODE'] = HC_subset_umls['CUI'].map(
    HC_subset_umls.groupby('CUI').apply(lambda x: get_group_value(x, 'MDR', 'CODE'))
)
HC_subset_umls['MDR_TTY'] = HC_subset_umls['CUI'].map(
    HC_subset_umls.groupby('CUI').apply(lambda x: get_group_value(x, 'MDR', 'TTY'))
)
HC_subset_umls['MDR_STR'] = HC_subset_umls['CUI'].map(
    HC_subset_umls.groupby('CUI').apply(lambda x: get_group_value(x, 'MDR', 'STR'))
)

# 生成MSH相关列
HC_subset_umls['MSH_CODE'] = HC_subset_umls['CUI'].map(
    HC_subset_umls.groupby('CUI').apply(lambda x: get_group_value(x, 'MSH', 'CODE'))
)
HC_subset_umls['MSH_TTY'] = HC_subset_umls['CUI'].map(
    HC_subset_umls.groupby('CUI').apply(lambda x: get_group_value(x, 'MSH', 'TTY'))
)
HC_subset_umls['MSH_STR'] = HC_subset_umls['CUI'].map(
    HC_subset_umls.groupby('CUI').apply(lambda x: get_group_value(x, 'MSH', 'STR'))
)

# 去重，保留每个CUI的唯一行
HC_subset_umls = HC_subset_umls.drop_duplicates(subset='CUI').reset_index(drop=True)

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