百万行场景下，如何用Pandas高效查找列表列指定位置的值？

问题场景 #

现有如下结构的Pandas DataFrame，data列存储长度不一（最大50个元素）的数值列表，lookup_val列是需要提取的元素位置（1-based索引）。需要快速生成output列存储对应位置的元素，且数据量达数百万行，常规apply或循环效率极低。

示例输入：

data  lookup_val
0       [1.3, 4.5, 6.4, 7.3, 8.9, 10.3]           5
1  [2.5, 4.7, 6.4, 6.6, 8.5, 9.3, 17.4]           3
2       [3.3, 4.2, 5.1, 7.8, 9.2, 11.5]           6

预期输出：

data  lookup_val  output
0       [1.3, 4.5, 6.4, 7.3, 8.9, 10.3]           5     8.9
1  [2.5, 4.7, 6.4, 6.6, 8.5, 9.3, 17.4]           3     6.4
2       [3.3, 4.2, 5.1, 7.8, 9.2, 11.5]           6    11.5

示例数据生成代码：

import pandas as pd

df = pd.DataFrame(
    [
        {'data': [1.3, 4.5, 6.4, 7.3, 8.9, 10.3], 'lookup_val': 5},
        {'data': [2.5, 4.7, 6.4, 6.6, 8.5, 9.3, 17.4], 'lookup_val': 3},
        {'data': [3.3, 4.2, 5.1, 7.8, 9.2, 11.5], 'lookup_val': 6},
    ]
)

高效实现方案 #

使用Pandas的lookup向量化方法，底层基于numpy实现，避免逐行循环，性能远超apply：

# 将data列的列表转换为二维DataFrame，短列表自动补NaN
data_matrix = pd.DataFrame(df['data'].tolist())

# 利用lookup提取对应位置元素（注意将1-based的lookup_val转为0-based索引）
df['output'] = data_matrix.lookup(df.index, df['lookup_val'] - 1)

方案说明 #

1. 向量化优势：lookup是Pandas原生的向量化操作，直接对整个数据块进行处理，无需逐行遍历，处理百万级数据的速度比apply快10-100倍。
2. 索引转换：示例中lookup_val是1-based位置（第5个元素对应索引4），所以需要减1转为0-based索引匹配DataMatrix的列索引。
3. 异常处理：如果某行的lookup_val超过对应列表的长度，会返回NaN，可根据业务需求添加过滤或填充逻辑（比如df = df[df['lookup_val'] <= df['data'].str.len()]提前过滤无效行）。

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