如何高效将多S3桶文本文件关联至Spark DataFrame？

你的核心问题是将1亿+ S3 URI全部拉取到Driver内存导致的Driver过载，以及Spark处理超大规模路径列表时的Listing瓶颈。以下是两种高效的替代方案，完全避免Driver端的内存压力：

方案1：分布式分区读取并直接关联ID #

利用mapPartitions在Executor端批量处理每个分区的URI，读取文件时直接保留原始id，省去后续通过path关联的开销：

from pyspark.sql import Row
import boto3
from io import StringIO

def read_s3_partition(partition):
    # 每个分区初始化一次S3客户端，减少连接开销
    s3 = boto3.client('s3')
    for row in partition:
        s3_uri = row.s3_uri
        record_id = row.id
        # 解析S3 URI的桶和路径
        bucket = s3_uri.split('//')[1].split('/')[0]
        key = '/'.join(s3_uri.split('//')[1].split('/')[1:])
        try:
            # 读取文件内容
            resp = s3.get_object(Bucket=bucket, Key=key)
            content = resp['Body'].read().decode('utf-8')
            yield Row(id=record_id, content=content, path=s3_uri)
        except Exception as e:
            # 可选：记录错误日志或标记失败的URI
            print(f"Read failed for {s3_uri}: {str(e)}")
            continue

# 先对元数据DF做合理分区，建议按哈希或桶名分区，控制每个分区1w-10w条URI
# 分区数根据集群规模调整，比如10000个分区对应1亿条数据，每个分区约1万条
df_meta_part = df_meta.repartition(10000)

# 分布式读取并转换为DataFrame
df_want = df_meta_part.rdd.mapPartitions(read_s3_partition).toDF()

方案2：基于Spark分区临时文件的批量读取 #

如果不想依赖boto3，可以利用Spark的分布式能力，将URI按分区拆分后批量读取：

from pyspark.sql import functions as F

# 1. 给元数据DF添加分区标识，按URI哈希拆分
df_meta_with_part = df_meta.withColumn("part_id", F.hash("s3_uri") % 10000)

# 2. 将每个分区的URI写入临时文本文件（每个分区对应一个文件）
df_meta_with_part.select("s3_uri").write.partitionBy("part_id").text("/tmp/s3_uri_parts")

# 3. 每个分区读取对应的URI列表
def read_part_uris(partition_file):
    # 读取单个分区的URI文件
    uris = [line.strip() for line in open(partition_file[1]) if line.strip()]
    if not uris:
        return []
    # 读取该分区所有URI对应的文件
    sub_df = spark.read.text(uris, wholetext=True).withColumn("path", F.input_file_name())
    return sub_df.rdd.collect()

# 分布式处理所有分区文件
temp_rdd = spark.sparkContext.wholeTextFiles("/tmp/s3_uri_parts/*/*").flatMap(read_part_uris)
df_temp = temp_rdd.toDF()

# 4. 关联原始ID（path与s3_uri直接匹配）
df_want = df_temp.join(df_meta, df_temp.path == df_meta.s3_uri, "inner").select("id", "content", "path")

关键优化说明 #

• 杜绝Driver内存过载：所有URI处理都在Executor端完成，Driver仅负责任务调度，不会加载亿级数据。
• 消除大表Join开销：方案1在读取时直接绑定id，无需后续通过path进行大规模关联，性能提升显著。
• 分区粒度控制：根据集群CPU、内存调整分区数，确保每个分区的URI数量在合理范围，避免任务过轻或过重。
• 容错处理：添加异常捕获，单个文件读取失败不会导致整个任务终止，可根据需求记录错误或跳过无效文件。

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