从Azure Databricks迁移8亿条数据至S3的性能优化求助

问题背景 #

需将8亿条记录从Azure Databricks迁移至AWS S3，测试阶段发现现有方案效率不足：

• 单文件写入100万条记录耗时约60秒，但单文件模式无法适配8亿级数据的存储与操作需求；
• 循环分块写入（拆分为5个20万条文件）耗时长达20-25分钟，且数据量增大后文件间间隔持续变长，完全无法支撑8亿条数据的迁移。

现有代码的核心问题 #

1. 单文件写入方案 #

from boto3.session import Session
from pyspark.sql import SparkSession
import requests
import boto3

file_name = 'full_data_in_one_file.csv'

s3_client = boto3.client('s3', aws_access_key_id=ACCESS_KEY, aws_secret_access_key=SECRET_KEY)
s3_client.put_object(Body=df.toPandas().to_csv(index=False), Bucket=bucket_name, Key=file_name)

• 调用df.toPandas()会将所有Spark分布式数据拉取到Driver节点，极易触发内存溢出，且单线程写入S3完全浪费Spark集群的分布式计算能力。

2. 循环分块写入方案 #

from boto3.session import Session
from pyspark.sql import SparkSession
import requests
import boto3
import json
from pyspark.sql.functions import monotonically_increasing_id
from pyspark.sql import functions as F
from pyspark.sql.window import Window
import random
import time

s3_client = boto3.client('s3', aws_access_key_id=ACCESS_KEY, aws_secret_access_key=SECRET_KEY)

df = spark.read.format("json").options(header='true', inferSchema='true').option("mode","DROPMALFORMED").load("abfss://path/product.db/table_name")
total_records = df.count()
print("Total records : ",total_records)
df_with_index = df.withColumn("index", monotonically_increasing_id())

num_loop=3
data_size_per_loop=200000
start_time = time.time()
for i in range(num_loop):
    start = i * data_size_per_loop
    end = (i + 1) * data_size_per_loop
    paginated_df = df_with_index.filter((df_with_index["index"] >= start) & (df_with_index["index"] < end))
    paginated_records = paginated_df.count()
    print("paginated_records : ",paginated_records)
    paginated_data = paginated_df.filter(paginated_df.index >= start).filter(paginated_df.index < end)
    list_data = paginated_data.toJSON().collect()
    size = len(list_data)
    json_data = json.dumps(list_data)
    
    if size == 0:
            print("No more records for push into AWS.")
            break
    else:
        data_payload = '\n'.join(json_data)
        end_data_range_time = time.time() - start_time
        print("Time taken to prepare aws format json data : ", end_data_range_time, " seconds")
        savefile_name= file_name+str(i)+"_"+str(size)+"_"+str(end)+".json"
        s3_client.put_object(Body=data_payload, Bucket=bucket_name, Key=savefile_name)
        end_AWS = time.time() - end_data_range_time
        print(f"data push into AWS S3. Starting records: {start} and end: {end}, size : {size}, loop nimber : {i}, time to push into AWS : {end_AWS} seconds")

• 每次循环过滤都会重新扫描全量DataFrame，重复计算导致耗时剧增；
• collect()将分区数据拉取到Driver节点，随着循环次数增加，Driver内存压力持续增大；
• 单线程调用boto3写入S3，完全未利用Spark的分布式执行能力，效率极低。

高效迁移方案 #

方案一：Spark原生分布式写入S3（推荐） #

Spark支持直接写入S3，利用集群所有节点并行处理，是效率最高的方案。

写入CSV格式

# 配置AWS凭证（建议通过Databricks Secrets管理，避免硬编码）
spark._jsc.hadoopConfiguration().set("fs.s3a.access.key", ACCESS_KEY)
spark._jsc.hadoopConfiguration().set("fs.s3a.secret.key", SECRET_KEY)
spark._jsc.hadoopConfiguration().set("fs.s3a.endpoint", "s3.amazonaws.com") # 根据AWS区域调整

# 写入S3，自动分块
df.write \
  .mode("overwrite") \
  .option("header", "true") \
  .option("maxRecordsPerFile", 1000000) # 每个文件100万条，可按需调整
  .csv("s3a://bucket_name/path/to/output")

写入JSON格式

df.write \
  .mode("overwrite") \
  .option("maxRecordsPerFile", 1000000)
  .json("s3a://bucket_name/path/to/output")

优势

• 完全分布式执行，集群所有节点同时写入，充分利用计算资源；
• 自动分块，无需手动加索引和循环过滤，避免重复扫描数据；
• 数据无需拉取到Driver节点，彻底规避内存瓶颈；
• 支持多种格式，兼容性强。

方案二：使用Parquet格式优化存储与性能 #

若无需明文格式，推荐使用Parquet列式存储，压缩比高、读写速度快，更适合大数据场景：

df.write \
  .mode("overwrite") \
  .option("maxRecordsPerFile", 1000000) \
  .parquet("s3a://bucket_name/path/to/output")

方案三：基于foreachPartition的分布式boto3写入（特殊场景适配） #

若因业务限制必须使用boto3，需将写入逻辑下放至Executor节点，避免Driver单点瓶颈：

from pyspark.sql import SparkSession
import boto3
import json
import uuid

def write_partition_to_s3(partition):
    # 每个分区初始化一次S3客户端，减少连接开销
    s3_client = boto3.client('s3', aws_access_key_id=ACCESS_KEY, aws_secret_access_key=SECRET_KEY)
    records = list(partition)
    if not records:
        return
    # 转换为JSON Lines格式
    json_lines = '\n'.join([json.dumps(record.asDict()) for record in records])
    # 生成唯一文件名避免冲突
    file_name = f"output_{uuid.uuid4()}.json"
    s3_client.put_object(Body=json_lines, Bucket=bucket_name, Key=file_name)

# 调整分区数，控制每个分区的记录量（示例：每100万条一个分区）
partition_count = df.count() // 1000000 + 1
df.repartition(partition_count).foreachPartition(write_partition_to_s3)

注意事项

• 通过repartition合理设置分区数，避免分区过多或过少；
• 每个分区内初始化一次S3客户端，不要为每条记录创建连接；
• 生成唯一文件名，防止文件覆盖冲突。

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