Pandas如何不使用iterrows根据colname列值匹配对应列值生成新列

问题场景 #

现有结构如下的pandas DataFrame：

import pandas as pd
df = pd.DataFrame({
    'a': [1, 2, 3],
    'b': [10, 30, 40],
    'colname': ['a', 'b', 'a']
})

需求为：禁止使用iterrows方法，根据colname列中存储的列名字符串，逐行获取对应列的取值，生成新的结果列，要求实现简洁优雅。

实现方案 #

1. 简洁优先：常规场景首选 #

代码最短、可读性最高，不需要额外依赖，绝大多数数据量场景下足够用：

df['result'] = df.apply(lambda row: row[row['colname']], axis=1)

执行后输出结果如下：

a   b colname  result
0  1  10       a       1
1  2  30       b      30
2  3  40       a       3

2. 性能优先：大数据量场景选 #

如果是百万行级以上的大数据集，apply的按行遍历效率偏低，可以用numpy高级索引实现纯向量化运算，性能和pandas原生列运算持平：

import numpy as np
# 将列名映射为DataFrame中的列位置索引
col_pos = df.columns.get_indexer(df['colname'])
# 按行索引+列位置批量取值
df['result'] = df.to_numpy()[np.arange(len(df)), col_pos]

运行结果和上面的方案完全一致。

注：旧版pandas曾提供lookup方法实现该需求，但该方法从pandas 1.2.0版本开始已被标记为弃用，不建议在新代码中使用。

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