单节点Databricks集群多线程Zstd压缩时Spark Driver异常终止

核心问题分析 #

你遇到的崩溃本质是Python ThreadPool与Spark执行模型冲突，加上Zstd多线程压缩在单节点下耗尽资源导致的。Spark本身是基于JVM的分布式计算框架，Python线程池会额外占用资源，且和Spark的任务调度机制冲突，同时Zstd默认多线程压缩会进一步加剧单节点的CPU/内存负载。

具体解决办法 #

• 放弃Python ThreadPool，改用Spark原生分区处理
  Spark DataFrame本身支持按字段分区并行写入，完全不需要手动开ThreadPool。直接按country字段分区后写入，Spark会自动调度并行任务：

  # 按country字段重新分区，自动并行处理每个国家的数据
  df.repartition("country") \
    .write \
    .option("compression", "zstd") \
    .mode("overwrite") \
    .csv("/dbfs/path/to/your/output")
  

  如果需要按国家生成独立的目录结构，用partitionBy：

  df.write \
    .option("compression", "zstd") \
    .partitionBy("country") \
    .mode("overwrite") \
    .csv("/dbfs/path/to/your/output")
  

• 禁用Zstd多线程压缩
  单节点资源有限，强制Zstd用单线程压缩，避免资源过载：

  df.repartition("country") \
    .write \
    .option("compression", "zstd") \
    .option("zstd.use.multithreading", "false") \
    .option("zstd.compression.level", 3)  # 可选：调整压缩级别平衡速度和压缩率
    .mode("overwrite") \
    .csv("/dbfs/path/to/your/output")
  

• 优化分区与资源匹配
  避免coalesce的不当使用：如果过滤后数据量不均，coalesce会导致部分分区数据过大，触发内存溢出。改用repartition根据单节点CPU核数设置合理分区数（比如核数的2-4倍）：

  # 假设单节点有4核，设置8个分区
  df.filter(df.country == "XX") \
    .repartition(8) \
    .write \
    .option("compression", "zstd") \
    .mode("overwrite") \
    .csv("/dbfs/path/to/your/output/country_XX")
  

  同时检查集群配置：调高Driver的CPU核数和内存（单节点集群Driver即Worker），比如选择更大的实例类型，确保有足够资源处理压缩和写入。

• 增强日志排查
  调高Spark日志级别，或者在代码中添加关键节点日志：

  import logging
  logging.basicConfig(level=logging.INFO)
  
  # 打印每个国家的数据量
  country_list = df.select("country").distinct().rdd.flatMap(lambda x: x).collect()
  for country in country_list:
      cnt = df.filter(df.country == country).count()
      logging.info(f"Processing country {country}, total records: {cnt}")
      # 写入逻辑...
  

  同时在Databricks集群的Driver日志中查看详细报错（之前可能没找对日志位置，Driver日志在集群页面的"Logs"标签下）。

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