如何在Databricks中用Spark并行写入JSON文件至挂载目录？

这个问题我太熟悉了——用collect()把5万条记录拉到驱动端单线程写入，速度慢到离谱完全是意料之中的事！而且dbutils确实只能在驱动端用，executor节点上直接调用会报错。下面给你一个分布式并行写入的方案，能把速度提升几个数量级：

核心思路 #

放弃驱动端单线程处理，改用Spark的foreachPartition()让每个executor节点并行处理一部分数据，同时用Hadoop原生的文件系统API替代dbutils（因为Hadoop API可以在executor上直接访问Databricks的挂载目录，不管是Azure的ABFS还是其他兼容存储都适用）。

具体实现代码 #

from py4j.java_gateway import java_import
from pyspark.sql import SparkSession

def write_json_batch(iterator):
    # 获取当前Spark会话的JVM实例，用来调用Hadoop API
    spark = SparkSession.getActiveSession()
    jvm = spark._jvm
    
    # 导入Hadoop文件系统相关的Java类
    java_import(jvm, 'org.apache.hadoop.fs.FileSystem')
    java_import(jvm, 'org.apache.hadoop.fs.Path')
    
    # 获取Hadoop文件系统实例，自动适配Databricks的挂载存储
    fs = jvm.FileSystem.get(spark._jsc.hadoopConfiguration())
    
    # 遍历当前分区的所有记录，并行写入
    for record in iterator:
        file_path = record['path']
        json_content = record['json']
        
        # 创建Hadoop路径对象
        hadoop_path = jvm.Path(file_path)
        
        # 打开输出流，第二个参数设为True表示覆盖已存在的文件
        output_stream = fs.create(hadoop_path, True)
        try:
            # 将JSON内容转为字节写入文件
            output_stream.write(json_content.encode('utf-8'))
        finally:
            # 确保流被关闭，避免资源泄漏
            output_stream.close()

# 对RDD执行分布式写入
my_rdd.foreachPartition(write_json_batch)

为什么这个方案更快？ #

• 并行处理：每个executor节点会处理RDD的一部分分区，5万条记录会被分散到多个节点同时写入，彻底告别驱动端单线程的瓶颈。
• 内存友好：不需要把所有数据拉到驱动端，避免驱动端内存溢出的风险，同时每个executor只处理自己分区的数据，内存压力更小。
• 原生API支持：Hadoop的FileSystem API是Databricks挂载存储的底层访问方式，比通过dbutils转发效率更高。

优化建议 #

1. 调整分区数：如果你的集群资源充足，可以先对RDD重新分区，比如my_rdd.repartition(100)，让分区数和executor的核心数匹配，进一步提升并行度。
2. 避免文件覆盖：如果你的path存在重复，代码里的fs.create(hadoop_path, True)会直接覆盖旧文件，若需要保留旧文件，可在写入前判断文件是否存在：

   if not fs.exists(hadoop_path):
       output_stream = fs.create(hadoop_path)
       # ... 写入逻辑
   

3. 权限检查：确保你的Spark作业有写入目标挂载目录的权限，Databricks默认挂载的存储权限是开放的，但自定义挂载可能需要额外配置。

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