Spark驱动内存溢出问题：循环复用前日Flag的逻辑优化求助

错误信息：Spark driver stopped unexpectedly due to memory

当前代码采用循环每日数据的方式复用前一日Flag列表计算当日状态，但循环执行1000次后，对data_wt_flag1的操作耗时极长，最终触发Driver内存溢出。原代码逻辑需先获取前一日Flag才能计算当日结果，误以为无法使用分布式处理，实际上可通过分组迭代优化。

原代码 #

DateList=data.select("Date").distinct().orderBy('AsOfDate').rdd.flatMap(lambda x: x).collect()
Flag_list=[]

data_wt_flag1=spark.createDataFrame(data = [],schema = StructType([]))

for daily_date in DateList:
  print(daily_date)  
  Temp_data_daily=data.filter(col("Date").isin(daily_date))
  Temp_data_daily=Temp_data_daily.withColumn('lag_1',when(col("identifier").isin(Flag_list),1).otherwise(0))
    
  Temp_data_daily=Temp_data_daily.withColumn("condition_1", when(((col('col_1')==1) & ((col('col_2')==1) | (col('lag_1')==1))),1).otherwise(0))
  Flag_list=Temp_data_daily.filter(col('condition_1')==1).select("identifier").distinct().rdd.flatMap(lambda x: x).collect() 
  data_wt_flag1=data_wt_flag1.unionByName(Temp_data_daily,allowMissingColumns=True)

代码逻辑说明 #

当日condition_1为1的条件：col_1 == 1 且 (col_2 == 1 或 前日该identifier的condition_1 == 1)，否则为0。首个日期所有identifier的lag_condition_1为0，后续每日用前一日符合条件的identifier列表标记当日的lag_condition_1。

样本输入数据 #

期望输出数据 #

优化方案 #

原代码问题分析 #

1. 数据集中到Driver：DateList.collect()和Flag_list.collect()将大量数据加载到Driver内存，随着日期和Identifier数量增加，内存占用急剧上升。
2. Lineage过度累加：循环中多次unionByName会不断扩展DataFrame的操作历史，Spark需维护所有步骤的元数据，最终导致Driver内存溢出。
3. 串行低效处理：每日循环过滤数据、更新列表的操作完全依赖Driver串行执行，无法利用Spark的分布式计算能力。

优化实现代码 #

利用Spark分组迭代处理，将每个Identifier的时间序列独立计算，所有操作在Executor分布式执行，避免Driver内存瓶颈：

from pyspark.sql.types import StructType, StructField, StringType, DateType, IntegerType

# 定义分组处理函数
def process_identifier_group(iterator):
    # 将迭代器转为列表并按日期排序
    rows = list(iterator)
    rows.sort(key=lambda x: x.Date)
    
    prev_condition = 0
    result_rows = []
    
    for row in rows:
        # 当日的lag_condition_1为前一日的condition_1
        lag_condition = prev_condition
        
        # 计算当日condition_1
        if row.col_1 == 1 and (row.col_2 == 1 or lag_condition == 1):
            current_condition = 1
        else:
            current_condition = 0
        
        # 构造包含新列的结果行
        result_rows.append(
            (row.Identifier, row.Date, row.col_1, row.col_2, lag_condition, current_condition)
        )
        
        # 更新前一日状态为当日状态
        prev_condition = current_condition
    
    return result_rows

# 定义输出Schema
output_schema = StructType([
    StructField("Identifier", StringType(), nullable=False),
    StructField("Date", DateType(), nullable=False),
    StructField("col_1", IntegerType(), nullable=False),
    StructField("col_2", IntegerType(), nullable=False),
    StructField("lag_condition_1", IntegerType(), nullable=False),
    StructField("condition_1", IntegerType(), nullable=False)
])

# 执行分组处理
optimized_result = data.groupBy("Identifier").mapGroups(process_identifier_group, output_schema)

# 查看结果
optimized_result.show()

优化优势 #

1. 分布式计算：所有业务逻辑在Executor节点完成，Driver仅负责任务调度，不存储大量业务数据。
2. Lineage简洁：避免了多次unionByName导致的元数据膨胀，内存占用可控。
3. 状态独立维护：按Identifier分组后，每个分组的状态独立传递，无需全局Flag列表，逻辑更清晰。

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