如何基于SAB列将Pandas DataFrame行数据转换为新增列？

原始数据结构 #

现有如下Pandas DataFrame：

CUI      SDUI  SAB  TTY                   STR
325040  C0011405   D003788  MSH   MH  Dental Pulp Diseases
325054  C0011405  10012328  MDR  LLT  Dental pulp disorder
325054  C0011405  10012326  MDR  LLT  Dental pulp disorders
325055  C0011405  10012328  MDR   PT  Dental pulp disorder

预期输出 #

希望基于SAB列生成新列，将不同SAB分组的字段拆分到同一行，预期结果如下：

CUI       SDUI_MSH    TTY_MSH  STR_MSH               SDUI_MDR    TTY_MDR  STR_MDR
325040  C0011405  D003788     MH       Dental Pulp Diseases  10012328    LLT      Dental pulp disorder
325040  C0011405  D003788     MH       Dental Pulp Diseases  10012326    LLT      Dental pulp disorders
325055  C0011405  D003788     MH       Dental Pulp Diseases  10012328    PT       Dental pulp disorder

尝试过的代码 #

之前使用pivot_table的方式无法得到预期结果：

df_umls = df_umls.pivot_table(index=['CUI'], columns='SAB', values=['SDUI', 'TTY', 'STR']).reset_index()
df_umls.columns = ['_'.join(col).strip() for col in df_umls.columns.values]

解决方案 #

pivot_table会对重复的索引+列组合进行聚合，不符合当前需求。正确做法是拆分不同SAB的数据集，再按CUI合并：

# 拆分MSH分组并重命名列
df_msh = df_umls[df_umls['SAB'] == 'MSH'].rename(
    columns={'SDUI': 'SDUI_MSH', 'TTY': 'TTY_MSH', 'STR': 'STR_MSH'}
).drop('SAB', axis=1)

# 拆分MDR分组并重命名列
df_mdr = df_umls[df_umls['SAB'] == 'MDR'].rename(
    columns={'SDUI': 'SDUI_MDR', 'TTY': 'TTY_MDR', 'STR': 'STR_MDR'}
).drop('SAB', axis=1)

# 按CUI左连接，保留所有MDR条目并匹配对应MSH数据
result = pd.merge(df_mdr, df_msh, on='CUI', how='left')

# 调整列顺序为预期格式
result = result[['CUI', 'SDUI_MSH', 'TTY_MSH', 'STR_MSH', 'SDUI_MDR', 'TTY_MDR', 'STR_MDR']]

print(result)

代码说明 #

1. 先将原始数据按SAB拆分为MSH和MDR两个子DataFrame，分别重命名字段以区分来源
2. 使用merge按CUI进行左连接，确保每个MDR条目都能匹配到对应的MSH数据
3. 最后调整列顺序，与预期输出一致

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