超大型100GB+CSV文件子组均值计算脚本性能优化求助

有一个大小超100GB的分号分隔文本文件，包含约1800万行数据、772列。列包括：sc16（整数）、cpid（整数）、type（字符串）、pubyr（整数），以及768个均为整数类型的dim_0至dim_767列。文件已按sc16和pubyr升序排列，sc16+pubyr的每个组合均为连续分组。

需求是：针对sc16与pubyr的每个唯一组合，计算所有dim_列的均值，将结果存入新DataFrame后保存为新文本文件。

但编写的脚本运行速度逐渐变慢，处理到500万行时已近乎停滞。机器配备96GB内存，首次尝试用itertools处理无法全量加载的超大型文件，脚本效率极低，需要优化。

原脚本代码：

import itertools
import pandas as pd

# Step 1: create an empty dataframe to store the mean values
mean_df = pd.DataFrame(columns=['sc16', 'pubyr'] + [f"dim_{i}" for i in range(768)]) 

# Step 2: open the file and iterate through the rows
with open('C:\Python_scratch\scibert_embeddings_sorted.txt') as f:
    counter = 0
    total_lines = sum(1 for line in f)
    f.seek(0)
    for key, group in itertools.groupby(f, key=lambda x: (x.split(';')[0], x.split(';')[3])): # group by the first (sc16) and fourth (pubyr) column 
        sc16, pubyr = key
        rows = [row.strip().split(';') for row in group]
        columns = rows[0]
        rows = rows[1:]
        # Step 3: convert the group of rows to a dataframe
        group_df = pd.DataFrame(rows, columns=columns)

        # Step 4: calculate the mean for the group
        mean_row = {'sc16': sc16, 'pubyr': pubyr}
        for col in group_df.columns:
            if col.startswith('dim_'):
                mean_row[col] = group_df[col].astype(float).mean()

        # Step 5: append the mean row to the mean dataframe
        mean_df = pd.concat([mean_df, pd.DataFrame([mean_row])], ignore_index=True)
        counter += len(rows)
        print(f"{counter} of {total_lines}")

# Step 6: save the mean dataframe to a new file
mean_df.to_csv('C:\Python_scratch\scibert_embeddings_mean.txt', sep=';', index=False)

核心问题分析 #

原脚本效率低下的关键原因：

• 每次分组都将整组数据加载到内存并转换为DataFrame，内存开销大且类型转换耗时
• 使用pd.concat频繁拼接DataFrame，随着结果集增大，性能呈指数级下降
• itertools.groupby的key函数重复执行split操作，且需将组迭代器转为列表，额外消耗资源

具体优化措施 #

1. 手动跟踪分组，逐行累加统计量
   利用文件已按sc16+pubyr排序的特性，无需依赖itertools.groupby，直接逐行读取并跟踪当前分组的key，实时累加dim列的总和与行数，计算均值后写入结果列表。

2. 避免频繁拼接DataFrame
   将所有结果行存入普通列表，最后一次性转换为DataFrame，彻底消除pd.concat的性能损耗。

3. 直接通过索引访问列，减少字符串操作
   提前确定各列的索引位置：sc16是第0列，pubyr是第3列，dim_0至dim_767对应第4到第771列，直接通过索引取值，避免列名判断和重复split操作。

4. 跳过不必要的全文件遍历
   原代码先遍历整个文件统计总行数，增加了一次IO开销，改为边处理边计数，节省时间。

优化后的代码 #

import pandas as pd

# 提前定义结果列名
result_columns = ['sc16', 'pubyr'] + [f"dim_{i}" for i in range(768)]
# 初始化结果列表
result_rows = []

# 定义列索引常量
SC16_IDX = 0
PUBYR_IDX = 3
DIM_START_IDX = 4
DIM_END_IDX = 4 + 767  # 对应dim_767的索引

with open(r'C:\Python_scratch\scibert_embeddings_sorted.txt', 'r') as f:
    # 读取表头（如果文件无表头，直接删除此句）
    header = f.readline().strip().split(';')
    current_key = None
    dim_sums = [0.0] * 768
    row_count = 0
    processed_rows = 0

    for line in f:
        line = line.strip()
        if not line:
            continue
        parts = line.split(';')
        # 获取当前行的分组key
        key = (parts[SC16_IDX], parts[PUBYR_IDX])

        if key != current_key:
            # 如果不是第一个分组，计算上一个分组的均值并加入结果
            if current_key is not None:
                sc16, pubyr = current_key
                mean_vals = [s / row_count for s in dim_sums]
                result_row = [sc16, pubyr] + mean_vals
                result_rows.append(result_row)
                # 打印进度
                print(f"已处理 {processed_rows} 行")
            
            # 重置当前分组的统计量
            current_key = key
            dim_sums = [0.0] * 768
            row_count = 0
        
        # 累加当前行的dim列数值
        for i in range(768):
            dim_val = float(parts[DIM_START_IDX + i])
            dim_sums[i] += dim_val
        row_count += 1
        processed_rows += 1
    
    # 处理最后一个分组
    if current_key is not None:
        sc16, pubyr = current_key
        mean_vals = [s / row_count for s in dim_sums]
        result_row = [sc16, pubyr] + mean_vals
        result_rows.append(result_row)
        print(f"已处理 {processed_rows} 行")

# 转换为DataFrame并保存
mean_df = pd.DataFrame(result_rows, columns=result_columns)
mean_df.to_csv(r'C:\Python_scratch\scibert_embeddings_mean.txt', sep=';', index=False)

额外优化建议 #

• 如果文件没有表头，直接删除读取表头的代码，从第一行数据开始处理
• 可以使用csv模块的csv.reader代替手动split，处理分号分隔更稳定（适用于列中包含转义分号的场景）
• 若需要更精确的进度显示，可以先获取文件总大小，跟踪已读取的字节数来计算进度百分比

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