大CSR稀疏矩阵写入S3遇EntityTooLarge报错，求读写函数方案

问题原因 #

出现EntityTooLarge错误是因为S3的PutObject接口单次上传最大支持5GB，500万行的CSR矩阵生成的npz文件超出了这个限制。解决核心是使用分段上传，boto3的upload_fileobj方法会自动处理这一逻辑，无需手动拆分文件。

改进后的写入函数 #

该函数自动处理大文件分段上传，同时保留原有的header功能：

import boto3
import io
import numpy as np
from scipy.sparse import csr_matrix

def write_large_sparse_matrix_to_s3(matrix, bucket_name, key, header=None):
    """
    将大型CSR格式稀疏矩阵写入S3桶，自动处理分段上传。
    
    Args:
        matrix (scipy.sparse.csr_matrix): 要写入的稀疏矩阵
        bucket_name (str): S3桶名称
        key (str): S3对象的键（文件名）
        header (list, optional): 要保存的头部信息列表
    """
    s3 = boto3.client('s3')
    output_buffer = io.BytesIO()
    
    # 先写入header（如果有）
    if header is not None:
        header_str = '\n'.join(header) + '\n'
        output_buffer.write(header_str.encode('utf-8'))
    
    # 将CSR矩阵数据压缩写入buffer，节省存储和上传时间
    np.savez_compressed(output_buffer, 
                        data=matrix.data, 
                        indices=matrix.indices, 
                        indptr=matrix.indptr, 
                        shape=matrix.shape)
    
    # 重置buffer指针到开头，准备上传
    output_buffer.seek(0)
    
    # 自动处理分段上传，无需手动拆分
    s3.upload_fileobj(output_buffer, bucket_name, key)
    
    print(f"大型稀疏矩阵已保存至 s3://{bucket_name}/{key}")

优化说明 #

• 用np.savez_compressed替代np.savez：对稀疏矩阵数据压缩，大幅降低文件体积，提升上传效率并减少存储成本。
• 改用upload_fileobj：自动处理超5GB文件的分段上传，同时采用流式写入，降低内存占用。
• 调整写入顺序：先写header再写矩阵数据，避免原代码中全量读取buffer拼接的内存开销。

对应的读取函数 #

从S3读取带header的大型CSR矩阵：

def read_large_sparse_matrix_from_s3(bucket_name, key, has_header=False):
    """
    从S3桶读取CSR格式稀疏矩阵，支持读取带header的文件。
    
    Args:
        bucket_name (str): S3桶名称
        key (str): S3对象的键（文件名）
        has_header (bool): 是否存在头部信息
    
    Returns:
        tuple: 如果has_header为True，返回(header_list, csr_matrix)；否则返回csr_matrix
    """
    s3 = boto3.client('s3')
    input_buffer = io.BytesIO()
    
    # 下载文件到内存buffer
    s3.download_fileobj(bucket_name, key, input_buffer)
    input_buffer.seek(0)
    
    header = None
    if has_header:
        # 读取header直到遇到空行
        header_lines = []
        while True:
            line = input_buffer.readline()
            if not line.strip():
                break
            header_lines.append(line.decode('utf-8').strip())
        header = header_lines
    
    # 读取npz数据重建CSR矩阵
    npz_data = np.load(input_buffer)
    matrix = csr_matrix(
        (npz_data['data'], npz_data['indices'], npz_data['indptr']),
        shape=npz_data['shape']
    )
    
    return (header, matrix) if has_header else matrix

使用示例 #

# 写入模拟大型稀疏矩阵
large_csr_matrix = csr_matrix(np.random.rand(5_000_000, 1000) > 0.99)
write_large_sparse_matrix_to_s3(large_csr_matrix, "my-bucket", "large_matrix.npz", header=["version:1.0", "date:2024-05-20"])

# 读取矩阵
header, loaded_matrix = read_large_sparse_matrix_from_s3("my-bucket", "large_matrix.npz", has_header=True)
print("Header:", header)
print("Matrix shape:", loaded_matrix.shape)

额外注意事项 #

• 内存优化：若矩阵大到无法完全放入内存，可先写入本地临时文件，再用upload_file方法上传，避免内存溢出。
• 分段阈值调整：可通过TransferConfig自定义分段大小，示例：

  from boto3.s3.transfer import TransferConfig
  config = TransferConfig(multipart_threshold=1024*1024*50)  # 设置50MB分段阈值
  s3.upload_fileobj(output_buffer, bucket_name, key, Config=config)
  

• 权限配置：确保boto3客户端拥有S3读写权限，可通过环境变量、IAM角色或配置文件完成认证。

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