Azure Databricks单节点集群中位数插补内存溢出问题求助

问题定位 #

你的代码存在几个关键问题导致内存占满：

1. 循环中反复调用withColumn叠加列操作，会让DataFrame的血统（Lineage）无限拉长，Spark需要维护所有中间步骤的元数据，内存消耗剧增。
2. 批量处理逻辑存在变量未定义的错误（代码中batch变量未声明），且每次小批量处理后写盘再读回的操作效率低下，反而额外增加IO和内存负担。
3. 单机集群（local[*,4]）的内存资源有限，30万行300多列的数据若分区不合理，会导致单分区数据量过大，触发内存溢出。

具体优化方案 #

1. 优化中位数计算逻辑 #

approxQuantile支持一次性计算多列的近似中位数，避免循环单列计算带来的重复数据扫描：

from pyspark.sql import functions as F

# 筛选所有数值型列
num_cols = [col for col, dtype in df_new.dtypes if dtype in ("int", "double", "float", "long")]
# 一次性计算所有数值列的近似中位数（误差0.01）
median_results = df_new.approxQuantile(num_cols, [0.5], 0.01)
# 转换为列名-中位数的字典
median_dict = dict(zip(num_cols, [res[0] for res in median_results]))

2. 避免Lineage过长：单步批量替换缺失值 #

不要循环调用withColumn，改用select一次性生成所有处理后的列，减少中间元数据维护：

# 构建所有列的处理逻辑：非数值列直接保留，数值列用中位数填充缺失值
processed_cols = []
for col in df_new.columns:
    if col in median_dict:
        # 填充缺失值并保留原列名
        filled_col = F.when(df_new[col].isNull(), median_dict[col]).otherwise(df_new[col]).alias(col)
    else:
        filled_col = df_new[col]
    processed_cols.append(filled_col)

# 一次性生成处理后的DataFrame
df_processed = df_new.select(*processed_cols)

3. 优化集群内存与分区配置 #

针对单节点集群的资源限制，调整以下参数：

• 在集群配置中添加内存参数（根据机器实际内存调整，比如16G内存的机器可设置）：

  spark.executor.memory 8g
  spark.driver.memory 8g
  spark.sql.shuffle.partitions 16
  

• 提前对DataFrame做合理分区，避免单分区数据过大：

  # 分区数建议为CPU核心数的2-4倍，你的集群是4核，设置8-16分区
  df_new = df_new.repartition(16)
  

4. 移除不必要的中间写盘操作 #

原代码中每次小批量处理后写盘再读回的操作完全冗余，直接在处理完成后一次性写入最终路径即可：

df_processed.write.mode("overwrite").parquet("/dbfs/tmp/filled_data")

5. 利用Delta Lake优化特性（可选） #

既然开启了Delta Lake预览，可通过以下操作减少内存压力：

• 若原数据是Delta表，处理前执行优化合并小文件：

  OPTIMIZE delta.`/path/to/your/original/data`
  

• 处理后的结果也可执行OPTIMIZE，提升后续读取性能

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