Databricks集群Driver内存溢出问题及分布式处理方案咨询

你遇到的Driver内存溢出，本质是统计计算阶段可能将大量中间结果或原始数据拉到了Driver节点，后续创建pandas DataFrame时，Driver内存已被占满，触发OOM导致集群重启。虽然最终的pandas DataFrame只有6行数据，但之前的统计逻辑如果是通过collect()、toPandas()这类将数据拉到Driver的操作获取统计值，就会导致Driver内存过载。

一、优化统计计算逻辑，全程在Executor执行 #

核心思路是避免将大数据集或中间结果拉到Driver，所有统计计算都在Spark集群的Executor节点分布式执行，最后只把聚合好的小结果集拉到Driver。

1. 用Spark原生API完成所有统计计算 #

替代手动计算单个统计值（如avg_fuelling_mean、avg_fuelling_Median等），直接用Spark的agg()方法一次性计算所有需要的统计量，结果保留为Spark DataFrame，而非单个变量。

示例代码：

# 假设原Spark DataFrame名为df，包含需要统计的列：fuelling, def_consumption, bsfc, diesel_cost, def_cost, fuel_cost
from pyspark.sql import functions as F

# 定义需要计算的统计函数
stats_expr = [
    # 柴油消耗相关统计
    F.avg("fuelling").alias("avg_fuelling_mean"),
    F.stddev("fuelling").alias("avg_fuelling_StdDev"),
    F.percentile_approx("fuelling", 0.5).alias("avg_fuelling_Median"),
    F.max("fuelling").alias("avg_fuelling_Max"),
    F.min("fuelling").alias("avg_fuelling_Min"),
    F.sum("fuelling").alias("avg_fuelling_Total"),
    F.percentile_approx("fuelling", 0.2).alias("avg_fuelling_20_Percentile"),
    F.percentile_approx("fuelling", 0.8).alias("avg_fuelling_80_Percentile"),
    # DEF消耗相关统计
    F.avg("def_consumption").alias("avg_def_mean"),
    F.stddev("def_consumption").alias("avg_def_StdDev"),
    F.percentile_approx("def_consumption", 0.5).alias("avg_def_Median"),
    F.max("def_consumption").alias("avg_def_Max"),
    F.min("def_consumption").alias("avg_def_Min"),
    F.sum("def_consumption").alias("avg_def_Total"),
    F.percentile_approx("def_consumption", 0.2).alias("avg_def_20_Percentile"),
    F.percentile_approx("def_consumption", 0.8).alias("avg_def_80_Percentile"),
    # BSFC、各类成本的统计同理补充...
]

# 执行分布式聚合计算，结果是只有1行的Spark DataFrame
aggregated_df = df.agg(*stats_expr)

# 只将聚合后的小结果集拉到Driver（此时数据量极小，不会OOM）
aggregated_dict = aggregated_df.collect()[0].asDict()

2. 基于聚合结果创建pandas DataFrame #

利用拉到Driver的小字典数据，构造目标pandas DataFrame，此时Driver内存压力极小：

import pandas as pd

FluidConsumptionValues_Table = pd.DataFrame(
{
    'Units': [fluid_units, fluid_units, bsfc_units,'$/hr', '$/hr', '$/hr'],
     'Mean': [aggregated_dict['avg_fuelling_mean'], aggregated_dict['avg_def_mean'], aggregated_dict['bsfc_mean'], aggregated_dict['DieselCost_mean'], aggregated_dict['DefCost_mean'], aggregated_dict['FuelCost_mean']],
     'StdDev': [aggregated_dict['avg_fuelling_StdDev'], aggregated_dict['avg_def_StdDev'], aggregated_dict['bsfc_StdDev'], aggregated_dict['DieselCost_StdDev'], aggregated_dict['DefCost_StdDev'], aggregated_dict['FuelCost_StdDev']],
     'Median': [aggregated_dict['avg_fuelling_Median'], aggregated_dict['avg_def_Median'], aggregated_dict['bsfc_Median'], aggregated_dict['DieselCost_Median'], aggregated_dict['DefCost_Median'], aggregated_dict['FuelCost_Median']],
     'Max': [aggregated_dict['avg_fuelling_Max'], aggregated_dict['avg_def_Max'], aggregated_dict['bsfc_Max'], aggregated_dict['DieselCost_Max'], aggregated_dict['DefCost_Max'], aggregated_dict['FuelCost_Max']], 
     'Min': [aggregated_dict['avg_fuelling_Min'], aggregated_dict['avg_def_Min'], aggregated_dict['bsfc_Min'], aggregated_dict['DieselCost_Min'], aggregated_dict['DefCost_Min'], aggregated_dict['FuelCost_Min']],
     'Total': [aggregated_dict['avg_fuelling_Total'], aggregated_dict['avg_def_Total'], aggregated_dict['bsfc_Total'], aggregated_dict['DieselCost_Total'], aggregated_dict['DefCost_Total'], aggregated_dict['FuelCost_Total']],
     '20th Percentile': [aggregated_dict['avg_fuelling_20_Percentile'], aggregated_dict['avg_def_20_Percentile'], aggregated_dict['bsfc_20_Percentile'], aggregated_dict['DieselCost_20_Percentile'], aggregated_dict['DefCost_20_Percentile'], aggregated_dict['FuelCost_20_Percentile']],
     '80th Percentile': [aggregated_dict['avg_fuelling_80_Percentile'], aggregated_dict['avg_def_80_Percentile'], aggregated_dict['bsfc_80_Percentile'], aggregated_dict['DieselCost_80_Percentile'], aggregated_dict['DefCost_80_Percentile'], aggregated_dict['FuelCost_80_Percentile']],
     'GDC': [aggregated_dict['avg_fuelling_GDC'], aggregated_dict['avg_def_GDC'], aggregated_dict['bsfc_GDC'], aggregated_dict['DieselCost_GDC'], aggregated_dict['DefCost_GDC'], aggregated_dict['FuelCost_GDC']]
 },
index=['Diesel Consumption', 'DEF Consumption', 'BSFC', 'Diesel Cost*', 'DEF Cost*', 'Total Cost*'])

二、临时缓解：调整Driver内存配置 #

如果暂时无法修改代码，可以通过调整Databricks集群的Driver内存参数缓解：

• 集群创建/编辑页面，找到Driver type，选择内存更大的实例类型（如从m5.xlarge改为m5.2xlarge）
• 也可以在集群的Advanced options -> Spark -> Spark config中添加：spark.driver.memory 32g（根据实例内存调整值）

但这是临时方案，核心还是要优化计算逻辑，避免Driver处理大数据。

三、额外注意点 #

• 绝对避免使用collect()、toPandas()直接操作8亿条的原始DataFrame，这类操作会将全量数据拉到Driver，必然触发OOM
• 对于百分位数计算，优先用Spark的percentile_approx()而非精确的percentile()，前者性能更高，占用内存更少，适合大数据集
• 如果统计逻辑复杂，可以考虑将中间聚合结果写入Delta Lake，再读取小结果集进行后续处理

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