解决大型Pandas DataFrame上传至S3时的MemoryError问题

核心问题 #

1. 分块逻辑完全错误：chunked_dataframe函数中的分组表达式np.arange(len(df)) // len(df) * chunk_size会将所有数据行分到同一个组中，等于未做任何分块，直接处理整个DataFrame，必然触发MemoryError。
2. 内存冗余开销：chunk.to_csv(index=False).encode('utf-8')会先把整个分块的CSV内容转换成字符串，再编码为字节串，两步操作都会在内存中生成完整的分块副本，大幅放大内存占用。

修复步骤 #

1. 修正分块逻辑 #

放弃按字节数预估分块（CSV每行大小不固定，精准控制难度高），改为固定行数分块，逻辑简单且可控。若需要接近指定字节大小，可先采样估算每行平均字节数，再反推分块行数。

2. 改用流式输出减少内存占用 #

使用io.BytesIO直接写入CSV内容，避免生成中间字符串，写完后将文件指针移到开头，直接传给S3的upload_part，减少内存拷贝环节。

完整修复代码 #

import io
import boto3
import concurrent.futures
import pandas as pd

def send_df_to_s3(df, bucket_name, key_name, chunk_rows=100000, num_threads=4):
    s3 = boto3.client('s3')
    
    # 初始化分片上传
    response = s3.create_multipart_upload(Bucket=bucket_name, Key=key_name)
    upload_id = response['UploadId']
    
    parts = []
    part_futures = {}
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor:
        for part_num, chunk in enumerate(chunked_dataframe(df, chunk_rows), start=1):
            # 提交分片上传任务
            future = executor.submit(
                s3.upload_part,
                Body=chunk,
                Bucket=bucket_name,
                Key=key_name,
                UploadId=upload_id,
                PartNumber=part_num
            )
            part_futures[future] = part_num
            
            # 控制并发数，避免同时加载过多分片到内存
            if len(part_futures) == num_threads:
                for completed_future in concurrent.futures.as_completed(part_futures):
                    parts.append({
                        'ETag': completed_future.result()['ETag'],
                        'PartNumber': part_futures[completed_future]
                    })
                part_futures = {}
        
        # 处理剩余分片
        for completed_future in concurrent.futures.as_completed(part_futures):
            parts.append({
                'ETag': completed_future.result()['ETag'],
                'PartNumber': part_futures[completed_future]
            })
    
    # 按分片号排序，完成上传
    parts.sort(key=lambda x: x['PartNumber'])
    response = s3.complete_multipart_upload(
        Bucket=bucket_name,
        Key=key_name,
        UploadId=upload_id,
        MultipartUpload={'Parts': parts}
    )
    
    return response

def chunked_dataframe(df, chunk_rows):
    """按固定行数分块DataFrame，返回可直接上传的BytesIO对象"""
    total_rows = len(df)
    # 生成每个分块的起始和结束索引
    for start in range(0, total_rows, chunk_rows):
        end = min(start + chunk_rows, total_rows)
        chunk = df.iloc[start:end]
        
        # 使用BytesIO直接写入CSV，避免生成中间字符串
        buffer = io.BytesIO()
        # 直接写入字节流，减少编码环节的内存开销
        chunk.to_csv(buffer, index=False, encoding='utf-8')
        # 将指针移到开头，以便S3读取
        buffer.seek(0)
        yield buffer

额外优化建议 #

• 若需要精准控制分片字节大小，可先采样估算：

  # 取前1000行计算平均每行字节数
  sample_buffer = io.BytesIO()
  df.head(1000).to_csv(sample_buffer, index=False, encoding='utf-8')
  avg_row_size = sample_buffer.tell() / 1000
  # 按100MB分片估算行数
  chunk_rows = int(100 * 1024 * 1024 / avg_row_size)
  

• 提前清理DataFrame冗余数据：删除不需要的列、将object类型转为category类型，降低原DataFrame的内存占用。

内容的提问来源于stack exchange，提问作者carlos de la morena