PySpark循环反复更新DataFrame速度过慢，求优化方案

问题背景 #

需要实现逻辑：基于过去14天的更新后score值，对次日的score和Amount进行更新（当当前score与过去14天更新后score的中位数差距超过3时，将score和Amount置为空）。原通过循环遍历日期的实现运行速度极慢，无法高效执行。

原代码核心问题 #

• 循环遍历日期，每次执行过滤、join、union操作，频繁触发Spark作业，产生大量不必要的Shuffle和计算开销
• 每次循环修改整个score_data数据集，导致数据重复读写，性能损耗严重
• 使用collect()将日期列表拉取到Driver端，数据量大时易引发内存溢出，同时破坏Spark的分布式计算特性

优化方案：基于分组状态迭代处理 #

利用Spark的mapGroupsWithState API，按id分组后按日期顺序处理每个用户的记录，维护每个id过去14天的更新后score状态，实时计算中位数并更新当前记录。这种方式仅需一次分组Shuffle，全程分布式处理，避免循环带来的性能损耗。

优化后代码 #

from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, DateType, DoubleType, IntegerType, ArrayType
import datetime

# 定义状态数据结构：存储过去14天的score列表及对应日期
state_schema = StructType([
    StructField("score_history", ArrayType(DoubleType()), nullable=True),
    StructField("latest_dates", ArrayType(DateType()), nullable=True)
])

def update_user_state(id_key, records, state):
    # 初始化状态
    if state is None:
        state = {"score_history": [], "latest_dates": []}
    
    # 按日期排序当前id的所有记录
    sorted_records = sorted(records, key=lambda x: x["AsOfDate"])
    
    updated_records = []
    for record in sorted_records:
        current_date = record["AsOfDate"]
        # 清理超过14天的历史数据
        cutoff_date = current_date - datetime.timedelta(days=14)
        valid_indices = [i for i, dt in enumerate(state["latest_dates"]) if dt > cutoff_date]
        state["score_history"] = [state["score_history"][i] for i in valid_indices]
        state["latest_dates"] = [state["latest_dates"][i] for i in valid_indices]
        
        # 计算过去14天更新后score的中位数（和原代码逻辑对齐）
        rolling_median = None
        if len(state["score_history"]) > 0:
            sorted_scores = sorted(state["score_history"])
            n = len(sorted_scores)
            if n % 2 == 1:
                rolling_median = sorted_scores[n//2]
            else:
                rolling_median = (sorted_scores[n//2 -1] + sorted_scores[n//2])/2
        
        # 更新当前记录的score和Amount
        current_score = record["score"]
        rolling_gap = abs(current_score - rolling_median) if rolling_median is not None and current_score is not None else None
        updated_score = None if (rolling_gap is not None and rolling_gap >3) else current_score
        updated_amount = None if (rolling_gap is not None and rolling_gap >3) else record["Amount"]
        
        # 生成更新后的记录
        updated_record = (
            record["id"],
            record["AsOfDate"],
            updated_score,
            updated_amount,
            record["score_original"],
            record["Amount_original"],
            rolling_median,
            rolling_gap
        )
        updated_records.append(updated_record)
        
        # 更新状态：仅保留更新后的非空score
        if updated_score is not None:
            state["score_history"].append(updated_score)
            state["latest_dates"].append(current_date)
    
    return iter(updated_records)

# 预处理数据：添加原始值列，按id分组聚合记录
score_data = score_data.withColumn('score_original', F.col('score')).withColumn('Amount_original', F.col('Amount'))
grouped_data = score_data.groupBy("id").agg(F.collect_list(F.struct(
    "AsOfDate", "score", "Amount", "score_original", "Amount_original"
)).alias("records"))

# 定义输出结果的Schema
result_schema = StructType([
    StructField("id", IntegerType(), nullable=True),
    StructField("AsOfDate", DateType(), nullable=True),
    StructField("score", DoubleType(), nullable=True),
    StructField("Amount", DoubleType(), nullable=True),
    StructField("score_original", DoubleType(), nullable=True),
    StructField("Amount_original", DoubleType(), nullable=True),
    StructField("rolling_median", DoubleType(), nullable=True),
    StructField("rolling_gap", DoubleType(), nullable=True)
])

# 使用mapGroupsWithState处理每个id的状态更新
updated_score_data = grouped_data.mapGroupsWithState(
    update_user_state,
    state_schema=state_schema,
    output_schema=result_schema,
    outputMode="append"
)

# 查看结果或保存
updated_score_data.show()

优化说明 #

• 分布式分组处理：按id分组后，每个分组独立在Executor上处理，充分利用Spark的分布式计算能力
• 状态复用：每个id维护过去14天的更新后score状态，避免重复查询历史数据，减少计算量
• 避免循环开销：无需遍历日期并反复修改整个数据集，仅一次分组Shuffle，大幅降低作业触发次数
• 内存高效：状态仅保留必要的历史数据，定期清理过期记录，避免内存占用过高

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