如何在Amazon SageMaker中高效读取S3大文本文件？

问题分析 #

当前代码采用单线程循环读取每个S3文件后逐个拼接DataFrame，存在明显瓶颈：

• 单线程无法利用实例多核CPU资源
• 每次pd.read_csv单独发起S3请求，重复网络开销累积
• ml.p3.2xlarge是GPU优化实例，CPU与网络IO配置并非针对大规模数据读取优化

核心优化方案 #

1. 更换适配的实例类型 #

放弃GPU实例，选择高CPU、高网络IO的实例：

• 推荐实例：c5.12xlarge（多核CPU+高带宽）、r5.12xlarge（大内存+高IO），按需可选用更大规格（如c5.24xlarge）
• 若数据量持续增长，可使用AWS EMR集群，通过分布式框架实现并行读取

2. 并行读取与拼接 #

方案A：用Dask实现分布式并行读取

Dask可自动拆分文件并利用多核并行读取，适配超大规模数据集：

import dask.dataframe as dd

def get_data_dask(bucket_name, prefix, cols):
    # 匹配S3路径下的所有目标文件
    s3_path = f"s3://{bucket_name}/{prefix}/*.txt"  # 根据实际文件后缀调整
    # 并行读取所有文件
    ddf = dd.read_csv(s3_path, sep="|", usecols=cols)
    # 转换为pandas DataFrame（内存足够时执行，否则保留Dask格式处理）
    return ddf.compute()

方案B：用多进程加速pandas读取

借助Pythonmultiprocessing库，并行处理多文件读取：

import pandas as pd
from multiprocessing import Pool

def read_single_file(file_path, cols):
    return pd.read_csv(file_path, sep="|", usecols=cols)

def get_data_multiprocess(bucket, prefix_objs, cols):
    # 生成所有文件的完整S3路径
    file_paths = [f"s3://{bucket.name}/{obj.key}" for obj in prefix_objs]
    # 根据CPU核心数创建进程池
    with Pool() as pool:
        # 并行读取所有文件
        df_list = pool.starmap(read_single_file, [(path, cols) for path in file_paths])
    # 拼接所有DataFrame
    return pd.concat(df_list, ignore_index=True)

3. 预优化数据存储格式 #

若需频繁读取该数据集，建议将文本文件转换为Parquet/ORC格式：

• 这类列式存储支持压缩、谓词下推，读取速度比CSV快数倍
• 转换示例（单次读取后执行）：

# 假设df为已读取的目标DataFrame
df.to_parquet("s3://your-bucket/path/to/output/*.parquet", engine="pyarrow", compression="snappy")
# 后续读取直接使用Parquet格式
df = pd.read_parquet("s3://your-bucket/path/to/output/*.parquet")

4. 细节优化 #

• 用s3fs优化S3连接：提前创建文件系统对象，避免重复建立连接

import s3fs
fs = s3fs.S3FileSystem()
# 读取时直接传入文件对象
with fs.open(file_path) as f:
    temp = pd.read_csv(f, sep="|", usecols=cols)

• 移除不必要的打印操作：print("File No: {}".format(i))会增加IO开销，可改为日志记录或直接移除

原代码中文注释版 #

def get_data(prefix_objs=prefix_objs, bucket=bucket, i=0, prefix_df=[], cols=col):
    # 遍历S3前缀下的所有文件对象
    for obj in prefix_objs:
        i += 1
        key = obj.key
        # 拼接完整的S3文件路径
        file_path = 's3://' + bucket.name + '/' + str(key)
        # 读取单个CSV文件，指定分隔符与需加载的列
        temp = pd.read_csv(file_path,
                          sep="|",
                          usecols=cols)
        print("已读取第{}个文件".format(i))
        prefix_df.append(temp)
    # 拼接所有读取的DataFrame并返回
    return pd.concat(prefix_df)

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