Pandas按公共列合并DataFrame行并追加统计字段的效率优化

方案1：Pandas分块处理（无需额外依赖） #

你原有方案的核心问题是多次重复扫描全量文件，IO开销过大，时间复杂度达到了O(n*m)级，处理10GB级文件必然很慢。仅需单次扫描文件即可完成所有处理，时间复杂度降到O(n)级：

• 第一步：分块读取CSV，每次读入chunksize=50万~100万行的小分片，分片大小可根据机器内存调整，避免内存溢出
• 第二步：对每个分片直接按观测公共字段分组，预存每个分组对应的split_on取值、概率、预测结果，同时统计每个观测的概率总和、条目计数，用于后续均值计算
• 第三步：所有分片处理完成后，统一收集所有split_on的唯一值，映射为group1/2/3...N的编号，再展开为要求的groupN_prob、groupN_pred格式列，缺失值自动填充NaN
• 第四步：生成平均预测分类列，可根据业务需求选择概率均值对应阈值分类，或所有分组预测的多数投票结果

示例代码：

import pandas as pd
import numpy as np
from collections import defaultdict

# 同一观测的公共字段，这些字段值在同个观测的多行中完全一致
common_cols = ['datetime', 'bID', 'data1', 'data2', 'data3']
# 初始化结果存储结构
res = defaultdict(lambda: {
    'group_data': dict(), # 存储每个split_on对应的值：(概率, 预测结果)
    'prob_sum': 0,
    'count': 0
})

# 分块读取文件，chunksize可根据内存调整
chunk_size = 500000
for chunk in pd.read_csv('你的文件路径.csv', chunksize=chunk_size, dtype={'split_on': str, 'probability': np.float32}):
    # 按公共字段分组，批量处理同个观测的所有行，比逐行遍历快5~10倍
    for key, group in chunk.groupby(common_cols):
        split_map = dict(zip(group['split_on'], zip(group['probability'], group['prediction'])))
        res[key]['group_data'].update(split_map)
        res[key]['prob_sum'] += group['probability'].sum()
        res[key]['count'] += len(group)

# 生成split_on到group编号的映射
all_split_vals = sorted({s for obs in res.values() for s in obs['group_data'].keys()})
split_to_group = {val: f'group{i+1}' for i, val in enumerate(all_split_vals)}

# 转换为最终输出格式
final_rows = []
for key, obs_info in res.items():
    row = dict(zip(common_cols, key))
    # 计算概率均值
    row['prob_mean'] = obs_info['prob_sum'] / obs_info['count']
    # 计算平均预测分类，此处示例为多数投票，可根据需求修改
    all_preds = [v[1] for v in obs_info['group_data'].values()]
    row['mean_prediction'] = max(set(all_preds), key=all_preds.count)
    # 展开所有分组列
    for split_val, group_name in split_to_group.items():
        prob, pred = obs_info['group_data'].get(split_val, (np.nan, np.nan))
        row[f'{group_name}_prob'] = prob
        row[f'{group_name}_pred'] = pred
    final_rows.append(row)

# 输出结果到CSV
final_df = pd.DataFrame(final_rows)
final_df.to_csv('合并结果.csv', index=False)

方案2：Dask处理（更适合小内存机器、速度更快） #

如果机器内存小于16G，或者想要进一步提升处理速度，可以使用Dask框架，它的API和Pandas完全兼容，会自动实现分块、并行计算，无需手动写分片逻辑：

• 直接用dask.dataframe.read_csv读入10GB文件，自动分块不会爆内存
• 调用内置的pivot_table或groupby+unstack即可一行完成行转列操作
• 最后直接输出到CSV即可，支持单文件输出

示例代码：

import dask.dataframe as dd
import pandas as pd
import numpy as np

common_cols = ['datetime', 'bID', 'data1', 'data2', 'data3']
# 读入文件，自动分块
ddf = dd.read_csv('你的文件路径.csv', dtype={'split_on': str, 'probability': np.float32})

# 1. 计算每个观测的概率均值
ddf_agg = ddf.groupby(common_cols).agg(
    prob_mean=('probability', 'mean')
).reset_index()

# 2. 行转列展开所有分组的概率和预测结果
ddf_pivot = ddf.pivot_table(
    index=common_cols,
    columns='split_on',
    values=['probability', 'prediction'],
    aggfunc='first'
).reset_index()

# 3. 重命名列符合要求的格式
ddf_pivot.columns = [
    f'group{col[1]}_prob' if col[0] == 'probability'
    else f'group{col[1]}_pred' if col[0] == 'prediction'
    else col[0]
    for col in ddf_pivot.columns
]

# 4. 合并聚合结果和展开列，输出到文件
ddf_final = ddf_pivot.merge(ddf_agg, on=common_cols, how='left')
ddf_final.to_csv('合并结果.csv', index=False, single_file=True)

额外优化提示 #

• 读文件时指定dtype参数给每个字段定义类型，避免Pandas自动推断类型浪费时间，同时降低内存占用
• 提前知道split_on的所有取值的话，可以跳过收集唯一值的步骤，直接映射列名，节省处理时间
• 不需要的字段在读入时用usecols指定仅读取需要的列，降低IO开销

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