如何基于迭代器在不加载全量DataFrame到内存的情况下生成合并CSV？

要实现不加载全部数据到内存的情况下生成合并CSV，核心思路是逐行读取小DataFrame的行，结合生成器取出对应chunk数据，实时写入CSV文件，避免一次性存储所有合并后的数据。以下是具体实现步骤和优化建议：

1. 优化Chunk查找效率（可选但关键） #

原代码中每次查找chunk都遍历所有文件名，效率极低。先预处理文件名映射，把start索引和对应的文件名、end索引存在字典里：

import os
import re
import pickle
import gc
import csv
import pandas as pd

# 预处理chunk映射：key=chunk起始索引，value=(文件名, chunk结束索引)
chunk_map = {}
pattern = re.compile(r'^(?:embeddingsssss_|featuress_)(\d+)_(\d+)$')
data_dir = "/content/drive/MyDrive/ProteiNNGO_Data/"

for filename in os.listdir(data_dir):
    match = pattern.match(filename)
    if match:
        start_idx = int(match.group(1))
        end_idx = int(match.group(2))
        chunk_map[start_idx] = (filename, end_idx)

# 优化后的chunk获取函数
def get_embedding_chunk(index, old_chunk_start):
    # 检查是否在当前已加载的chunk范围内
    if old_chunk_start in chunk_map:
        _, curr_end = chunk_map[old_chunk_start]
        if old_chunk_start <= index < curr_end:
            return None, old_chunk_start
    
    # 找到包含目标索引的chunk
    candidate_starts = [s for s in chunk_map.keys() if s <= index]
    if not candidate_starts:
        raise ValueError(f"No chunk found for index {index}")
    
    target_start = max(candidate_starts)
    filename, target_end = chunk_map[target_start]
    if not (target_start <= index < target_end):
        raise ValueError(f"Index {index} out of chunk range {target_start}-{target_end}")
    
    # 加载目标chunk
    with open(os.path.join(data_dir, filename), "rb") as f:
        return pickle.load(f), target_start

2. 生成CSV文件 #

第一步：确定CSV表头 #

需要合并小DataFrame的列和chunk中嵌入数据的列（假设嵌入是一维数组，自动生成列名）：

# 假设你的小DataFrame名为small_df
small_cols = small_df.columns.tolist()

# 获取嵌入数据的维度，用于生成列名
first_chunk, _ = get_embedding_chunk(0, -1)
embedding_dim = len(first_chunk[0])
emb_cols = [f"emb_{i}" for i in range(embedding_dim)]

# 合并后的完整表头
csv_headers = small_cols + emb_cols

第二步：逐行写入CSV #

方式一：使用你编写的embedding_generator

output_csv = "merged_output.csv"

# 先写入表头
with open(output_csv, 'w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    writer.writerow(csv_headers)

# 遍历生成器，逐行写入
with open(output_csv, 'a', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    for idx, embedding_data in embedding_generator():
        # 获取小DataFrame对应行的数据
        small_row = small_df.iloc[idx].tolist()
        # 合并数据（注意嵌入数据转成列表，numpy数组需转list）
        merged_row = small_row + list(embedding_data)
        writer.writerow(merged_row)
        
        # 可选：打印进度
        if idx % 1000 == 0:
            print(f"Processed {idx}/{len(small_df)} rows")

方式二：直接循环索引（避免生成器的额外封装）

output_csv = "merged_output.csv"

# 写入表头
with open(output_csv, 'w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    writer.writerow(csv_headers)

current_chunk = None
current_start = -1

with open(output_csv, 'a', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    for idx in range(len(small_df)):
        # 获取chunk数据
        new_chunk, new_start = get_embedding_chunk(idx, current_start)
        current_start = new_start
        
        if new_chunk is not None:
            # 释放旧chunk内存
            del current_chunk
            current_chunk = new_chunk
            gc.collect()
        
        # 取出当前索引对应的嵌入数据
        offset = idx - current_start
        embedding_data = current_chunk[offset]
        
        # 合并并写入
        small_row = small_df.iloc[idx].tolist()
        merged_row = small_row + list(embedding_data)
        writer.writerow(merged_row)
        
        if idx % 1000 == 0:
            print(f"Processed {idx}/{len(small_df)} rows")

关键注意事项 #

• 内存控制：切换chunk时删除旧chunk并调用gc.collect()，避免内存堆积；csv模块逐行写入不会缓存大量数据。
• 效率优化：预处理chunk_map把查找chunk的时间从O(n)降到O(1)，对50万行的场景提升非常明显。
• 数据格式：如果嵌入数据是numpy数组或其他非列表类型，必须转成list才能写入CSV。
• 编码问题：写入CSV时指定encoding='utf-8'，避免中文或特殊字符乱码。

内容的提问来源于stack exchange，提问作者dev-mirzabicer