如何高效替换Pandas中loc[[idx]]以提升循环索引查询性能？

问题背景 #

使用Pandas 1.5.3版本处理DataFrame时，需要按索引筛选记录并循环遍历。最初采用df.loc[[idx]]获取对应行的DataFrame，但因循环执行次数极多，导致整体耗时严重。通过cProfile分析发现，耗时主要集中在pandas._libs.index.IndexEngine的get_indexer_non_unique方法。尝试改用df.loc[idx]后速度有所提升，但返回结果类型不稳定——单条记录时返回Series，多条记录时返回DataFrame，而业务需求要求结果始终为DataFrame。

原实现代码 #

import pandas as pd
import cProfile
from tqdm import tqdm

def iterate_through_df():
    indexes = df.index.unique()
    for idx in tqdm(indexes):            
        df_info = df.loc[[idx]]
        # 后续业务代码...

df = pd.read_csv('random_data.csv', index_col='id')
cProfile.run('iterate_through_df()', sort='cumulative')

测试数据生成代码 #

用于生成测试用的CSV文件：

import pandas as pd
import numpy as np

size = 100000
num_columns = 100

data = {}
for i in range(1, num_columns + 1):
    key = f'name{i}:'
    data[key] = np.random.choice(['Alice', 'Bob', 'Charlie', 'David', 'Eva'], size=size)
random_indexes = np.random.randint(1, 100, size=size)
df = pd.DataFrame(data, index=random_indexes)
df.to_csv('random_data.csv', index_label='id')

性能分析结果 #

cProfile统计显示，耗时核心在索引查找方法：

ncalls  tottime  percall  cumtime  percall filename:lineno(function)
99    1.046    0.011    1.046    0.011 {method 'get_indexer_non_unique' of 'pandas._libs.index.IndexEngine' objects}

解决方案 #

方案1：使用groupby批量分组（推荐，性能最优） #

通过groupby按索引提前分组，遍历分组后的结果即可直接得到对应索引的DataFrame，无需每次调用loc进行索引查找，彻底解决性能问题，同时保证结果始终为DataFrame。

优化后代码：

import pandas as pd
from tqdm import tqdm

def iterate_through_df_optimized():
    # 按索引分组，一次性完成所有索引的分组逻辑
    grouped = df.groupby(df.index)
    for idx, df_info in tqdm(grouped):
        # df_info 始终是对应索引的DataFrame，单条/多条记录均适用
        # 后续业务代码...

df = pd.read_csv('random_data.csv', index_col='id')
# 可再次用cProfile验证性能提升
# cProfile.run('iterate_through_df_optimized()', sort='cumulative')

方案2：强制转换结果为DataFrame（兼容原有逻辑） #

如果必须保留循环调用loc的逻辑，可在获取结果后判断类型，若为Series则转换为DataFrame：

import pandas as pd
from tqdm import tqdm

def iterate_through_df_fixed():
    indexes = df.index.unique()
    for idx in tqdm(indexes):            
        df_info = df.loc[idx]
        # 若返回Series则转为单行DataFrame
        if isinstance(df_info, pd.Series):
            df_info = df_info.to_frame().T
        # 后续业务代码...

df = pd.read_csv('random_data.csv', index_col='id')

原理说明 #

• df.loc[[idx]]每次调用都会触发Pandas的非唯一索引查找逻辑（get_indexer_non_unique），循环次数过多时会累积大量耗时；
• groupby通过一次性遍历索引完成分组，后续遍历仅需直接读取分组数据，避免了重复的索引查找开销；
• df.loc[idx]虽省略了列表包装的索引查找开销，但返回类型随记录数变化，需额外处理以保证格式统一。

内容的提问来源于stack exchange，提问作者Carlos Eduardo de Schuller Ban