Pandas中如何利用np.select一次性生成多列？

原始DataFrame：

A   B
0   1.0 4
1   2.0 5
2   NaN 6

期望转换后的结果：

A   B   val val_source
0   1.0 4   1.0 A
1   2.0 5   2.0 A
2   NaN 6   6.0 B

尝试用*np.select*一次性生成多列但报错：

conds = [df['A'].notna(), True]
choices = [df[['A']].assign(val_source='A'), df[['B']].assign(val_source='B')]
df[['val', 'val_source']] = np.select(conds, choices)

目前只能通过两次独立的*np.select*实现：

conds = [df['A'].notna(), True]
_choices_val_src = [
    (df['A'], 'A'),
    (df['B'], 'B'),
]
choices_val, choices_src = zip(*_choices_val_src)
df['val'] = np.select(conds, choices_val, default=np.nan)
df['val_source'] = np.select(conds, choices_src, default=np.nan)

请问是否有更简洁的实现方式？

方案1：用*np.where*（最适合两分支场景） #

你的逻辑本质是A非空时取A和'A'，否则取B和'B'，np.where天生适合这种二元分支，能一次性处理两列：

import numpy as np
import pandas as pd

df['val'] = np.where(df['A'].notna(), df['A'], df['B'])
df['val_source'] = np.where(df['A'].notna(), 'A', 'B')

这比两次np.select代码更紧凑，可读性也更强。

方案2：用*apply*逐行生成元组拆分列 #

如果后续条件扩展为多个分支，用apply可以更灵活地生成多列：

df[['val', 'val_source']] = df.apply(
    lambda row: (row['A'], 'A') if pd.notna(row['A']) else (row['B'], 'B'),
    axis=1,
    result_type='expand'
)

注意要指定result_type='expand'，让返回的元组自动拆分为多列。

方案3：先处理val列再推导来源 #

先用*mask*或*combine_first*处理数值列，再用np.where生成来源列：

df['val'] = df['A'].mask(df['A'].isna(), df['B'])
# 或 df['val'] = df['A'].combine_first(df['B'])
df['val_source'] = np.where(df['A'].notna(), 'A', 'B')

这种写法把数值和来源逻辑拆分，适合需要单独调整数值处理逻辑的场景。

关于你之前np.select报错的原因 #

np.select的choices参数要求每个选项是与条件同形状的数组/Series，你传入的是DataFrame，所以会触发维度不匹配的错误。如果一定要用np.select一次性生成多列，可以把选项转成二维数组：

conds = [df['A'].notna(), True]
choices = [
    np.column_stack([df['A'], ['A']*len(df)]),
    np.column_stack([df['B'], ['B']*len(df)])
]
df[['val', 'val_source']] = np.select(conds, choices)

但这种写法反而冗余，只适合多分支（3个及以上条件）的场景。

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