优化Pandas读取大量文本数据并按类别计算指定变量均值的方法

问题背景 #

需要读取大量结构一致的模拟文本文件，按类别分组计算指定变量的时间步均值，当前实现效率偏低，尤其是最后分组计算均值的循环部分。

数据特征 #

• 文件总数、每个类别对应文件数、类别名称均已知
• 所有文件的列结构完全一致

当前实现的核心问题 #

• 读取文件时加载了不必要的列（如Q、M），浪费IO和内存
• 手动循环拼接DataFrame再计算均值的方式冗余，内存占用高且计算效率低

优化方案 #

一、优化文件读取环节 #

直接在读取时只加载目标变量列，避免无关数据占用资源：

import pandas as pd
import numpy as np

sim_nr = np.arange(1,13)
categories = ['a','b','c']
variables = ['S', 'P', 'L']
# 修复原代码拼写错误：mutlicolumn → multi_column
multi_column = pd.MultiIndex.from_tuples(list(zip(categories, sim_nr)))

all_data = {}
# 假设number_of_files是文件索引列表，与multi_column的索引一一对应
for idx, file_num in enumerate(number_of_files):
    # 仅读取需要的变量列，跳过其他列
    data = pd.read_csv(f'data_{file_num}.txt', usecols=variables)
    all_data[multi_column[idx]] = data

二、优化分组均值计算环节 #

利用Pandas的concat+MultiIndex分组功能，一次性完成所有类别的均值计算，彻底替代多层循环：

# 将字典中的所有DataFrame按列拼接，用原字典的键作为列的MultiIndex
combined_df = pd.concat(all_data.values(), axis=1, keys=all_data.keys())

# 按类别（MultiIndex第一层）分组，计算每个类别下各变量的时间步均值
category_mean = combined_df.groupby(level=0, axis=1).mean()

# 若需拆分出单个变量的DataFrame（匹配原需求的variableS/variableP/variableL）：
variableS = category_mean.xs('S', level=1, axis=1)
variableP = category_mean.xs('P', level=1, axis=1)
variableL = category_mean.xs('L', level=1, axis=1)

三、超大规模数据的内存友好方案 #

如果文件数量极大，内存不足以存储所有DataFrame，可采用边读取边累加的方式，避免存储全部数据：

# 初始化每个类别的累加器和计数（假设时间步固定为4行，可根据实际情况调整）
category_accum = {cat: pd.DataFrame(0.0, index=range(4), columns=variables) for cat in categories}
category_count = {cat: 0 for cat in categories}

for idx, file_num in enumerate(number_of_files):
    data = pd.read_csv(f'data_{file_num}.txt', usecols=variables)
    current_cat = multi_column[idx][0]
    category_accum[current_cat] += data
    category_count[current_cat] += 1

# 计算每个类别的均值
final_mean = {cat: category_accum[cat] / category_count[cat] for cat in categories}

# 可选：转换为统一的DataFrame结构
final_mean_df = pd.concat(final_mean.values(), axis=1, keys=final_mean.keys())

关键优化点 #

• 按需加载列：通过usecols参数过滤无关列，减少IO时间和内存占用
• 内置分组替代手动循环：利用Pandas高效的分组计算逻辑，避免冗余的DataFrame拼接操作
• 内存友好的累加模式：针对超大规模数据，仅保留累加结果和计数，内存占用仅为单个文件大小级别

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