Azure Databricks按ID拆分写入S3 Parquet文件过慢求助

你的核心瓶颈在于将分布式Spark DataFrame转为单节点pandas DataFrame后，通过串行循环写入50k个文件——这种单线程IO操作完全浪费了Spark的分布式并行能力，导致速度骤降。以下是针对性的优化方案：

优化方案1：Spark原生分布式写入（推荐） #

直接利用Spark的分布式特性按ID分区写入，彻底规避单节点循环操作：

# 读取数据部分保持不变
df_from_query = spark.read \
               .format("com.microsoft.sqlserver.jdbc.spark") \
               .option("url", url) \
               .option("dbtable", table_name) \
               .option("schemaCheckEnabled", False) \
               .option("accessToken", accesstoken) \
               .option("encrypt", "true") \
               .option("hostNameInCertificate", "*.database.windows.net") \
               .load()

# 按ID重新分区+分布式写入，每个ID对应一个文件
df_from_query.repartition("id") \
             .write \
             .mode("overwrite") \
             .partitionBy("id") \
             .parquet(f"/dbfs/mnt/{targetMountName}/")

关键说明： #

• repartition("id")：将数据按ID均匀拆分到分布式节点，每个ID的所有数据归为一个分区，确保后续每个分区生成独立文件。
• partitionBy("id")：写入时自动按ID创建子目录，每个子目录下存放对应ID的Parquet文件（默认每个分区对应一个文件）。
• 整个过程由Spark集群并行执行，50k个文件的写入操作会被分配到多个节点同时处理，速度能提升数十倍甚至上百倍。

优化方案2：针对性调整Spark配置 #

如果写入效率仍有优化空间，可添加以下配置：

• 匹配ID数量设置分区数：spark.conf.set("spark.sql.shuffle.partitions", "50000")
• 启用Parquet压缩减少IO量：spark.conf.set("spark.sql.parquet.compression.codec", "snappy")
• 确保单ID对应单文件：spark.conf.set("spark.sql.files.maxRecordsPerFile", "0")（不限制单文件记录数）

优化方案3：清理冗余操作 #

原代码中的df_from_query.cache()完全无用——因为后续直接转成了pandas DataFrame，Spark缓存的数据并未被利用，反而占用集群内存。如果使用Spark原生写入，仅在需要重复使用数据时才考虑缓存。

原代码慢的根本原因 #

1. 丢失分布式能力：将Spark分布式DataFrame转为单节点pandas DataFrame，所有后续操作只能在单个节点串行执行。
2. 串行IO开销：50k次文件写入是单线程依次执行，每个IO请求都需要等待前一个完成，时间成本呈线性累加。
3. 单节点资源限制：单节点的CPU、内存、网络带宽无法支撑大量数据的循环写入操作。

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