基于检查点列多条件提取数据的Pandas代码性能优化求助

问题背景 #

处理包含数万行的Pandas DataFrame，需根据40个检查点的规则（point_*_a和point_*_b）提取对应分组特征列（c组15个、d组10个），并与项目信息整合成统一表格。当前循环拼接的实现耗时20-30秒，需优化至2-3秒。

性能瓶颈诊断 #

当前代码的核心问题在于循环内重复执行高开销操作：

• 40次循环中，每次重复创建特征列的MultiIndex、执行stack和merge，时间复杂度随循环次数累积；
• 每次循环修改原DataFrame（df = df.loc[~df[f'point_{i}'].isnull()]），会丢失部分数据且重复过滤；
• 多次concat小DataFrame，导致内存碎片化和重复的索引重建。

优化方案：向量化+预处理批量操作 #

核心思路是一次性预处理所有特征和检查点，避免循环内重复计算，利用Pandas向量化操作替代循环：

完整优化代码 #

import pandas as pd
import numpy as np

# 读取原始数据
df = pd.read_csv('test_r.csv', sep='   ', low_memory=False)

# 定义列集合（复用原定义）
collist_tmp = [
    "a1_b1_c_foo", "a1_b1_c_foo_bar", "a1_b1_d_foo_bar_baz",
    "a2_b1_c_foo", "a2_b1_c_foo_bar", "a2_b1_d_foo_bar_baz",
    "a1_b2_c_foo", "a1_b2_c_foo_bar", "a1_b2_d_foo_bar_baz",
    "a2_b2_c_foo", "a2_b2_c_foo_bar", "a2_b2_d_foo_bar_baz",
    "a1_b3_c_foo", "a1_b3_c_foo_bar", "a1_b3_d_foo_bar_baz",
    "a2_b3_c_foo", "a2_b3_c_foo_bar", "a2_b3_d_foo_bar_baz",
]
collist_final_df = [
    "name", "country", "c_foo", "c_foo_bar",
    "d_foo_bar_baz", 'point', 'point_description', 'point_a', 'point_b'
]

# --------------------------
# 步骤1：预处理分组特征列（一次性完成）
# --------------------------
tmp = df.reindex(columns=collist_tmp)
# 拆分特征列名为(a_b组合, 特征名)的MultiIndex
tmp.columns = pd.MultiIndex.from_frame(tmp.columns.str.extract(r'(a\d+_b\d+)_(.*)'))
# 转成长格式，方便后续批量匹配
features_long = tmp.stack(level=0).reset_index().rename(columns={'level_1': 'a_b'})
features_long = features_long.melt(id_vars=['index', 'a_b'], var_name='feature', value_name='value')

# --------------------------
# 步骤2：预处理所有检查点列（一次性转长格式）
# --------------------------
# 提取所有检查点ID
point_ids = sorted(list(set(col.split('_')[1] for col in df.columns if col.startswith('point_'))))

point_dfs = []
for i in point_ids:
    # 提取当前检查点的相关列
    point_df = df[['index', 'name', 'country', 
                   f'point_{i}', f'point_{i}_description', 
                   f'point_{i}_a', f'point_{i}_b']].copy()
    # 过滤空检查点
    point_df = point_df.dropna(subset=[f'point_{i}'])
    # 统一列名
    point_df = point_df.rename(columns={
        f'point_{i}': 'point',
        f'point_{i}_description': 'point_description',
        f'point_{i}_a': 'point_a',
        f'point_{i}_b': 'point_b'
    })
    point_dfs.append(point_df)

# 合并所有检查点为长表
points_long = pd.concat(point_dfs, ignore_index=True)
# 转换a/b为字符串，用于构建匹配键
points_long['point_a_str'] = points_long['point_a'].astype('Int16').astype(str)
points_long['point_b_str'] = points_long['point_b'].astype('Int16').astype(str)

# --------------------------
# 步骤3：批量匹配c组和d组特征
# --------------------------
# 构建c组、d组的匹配键
points_long['c_a_b'] = 'a' + points_long['point_a_str'] + '_b' + points_long['point_b_str']
points_long['d_a_b'] = 'a' + points_long['point_a_str'] + '_b' + (points_long['point_b'].astype('Int16') + 1).astype(str)

# 提取c组特征并转宽表
c_features_wide = features_long[features_long['feature'].str.startswith('c_')]\
    .pivot(index=['index', 'a_b'], columns='feature', values='value').reset_index()
# 合并c组特征到检查点表
points_with_c = points_long.merge(c_features_wide, left_on=['index', 'c_a_b'], right_on=['index', 'a_b'], how='left')

# 提取d组特征并转宽表
d_features_wide = features_long[features_long['feature'].str.startswith('d_')]\
    .pivot(index=['index', 'a_b'], columns='feature', values='value').reset_index()
# 合并d组特征到检查点表
points_with_cd = points_with_c.merge(d_features_wide, left_on=['index', 'd_a_b'], right_on=['index', 'a_b'], how='left')

# 整理最终结果
df_final = points_with_cd[collist_final_df].copy()
print(df_final)

优化效果说明 #

通过预处理+向量化操作，将循环内的重复计算合并为一次性操作，避免了40次merge和stack的开销，处理数万行数据+40个检查点的耗时可压缩至2-3秒以内。

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