PySpark DataFrame写入过慢问题排查与优化求助

核心问题 #

你当前的性能瓶颈是循环遍历每个ID单独执行FPGrowth的模式，导致Spark构建了包含N个独立FPGrowth任务的巨型执行计划（N为唯一ID数量）。Spark的惰性评估会把所有循环内的逻辑延迟到最后触发action（write/count/display）时才执行，此时执行计划的复杂度随ID数量指数级上升——哪怕最终结果只有5万行，调度和执行的开销也会被拉满。

具体优化措施 #

1. 放弃单ID循环，改为按ID分组批量处理 #

不要逐个ID遍历，而是把交易数据按ID分组，在每个分组内执行FPGrowth逻辑。这种方式会让Spark将每个分组任务拆分为独立小作业，避免生成超级大的执行计划。

Scala示例：

// 按ID分组，生成每个ID对应的去重购物篮
val groupedData = transactionData
  .select("id", "product")
  .distinct()
  .groupBy("id")
  .agg(collect_set("product").alias("basket"))

// 每个分组内执行FPGrowth并生成规则
val allRules = groupedData.flatMapGroups { (id, basketsIter) =>
  val baskets = basketsIter.toList.map(_.getAs[Seq[String]]("basket"))
  val df = spark.createDataFrame(baskets).toDF("items")
  
  val fpGrowth = new FPGrowth()
    .setItemsCol("items")
    .setMinSupport(0.1) // 根据业务调整参数
    .setMinConfidence(0.5)
  val model = fpGrowth.fit(df)
  
  // 给规则打上ID标识
  model.associationRules.collect().map(rule => 
    (id, rule.antecedent, rule.consequent, rule.confidence, rule.lift)
  )
}

val rulesDF = allRules.toDF("id", "antecedent", "consequent", "confidence", "lift")

Python示例（用applyInPandas更高效）：

from pyspark.ml.fpm import FPGrowth
import pyspark.sql.functions as F

def process_id_baskets(pdf):
    # pdf是单个ID对应的DataFrame，包含basket列
    baskets = pdf['basket'].tolist()
    fpGrowth = FPGrowth(minSupport=0.1, minConfidence=0.5, itemsCol="items")
    model = fpGrowth.fit(spark.createDataFrame(baskets, "array<string>").toDF("items"))
    
    # 给规则添加ID字段
    rules = model.associationRules
    rules['id'] = pdf['id'].iloc[0]
    return rules

# 生成按ID分组的购物篮数据
grouped_data = transaction_data.select("id", "product").distinct() \
    .groupBy("id").agg(F.collect_set("product").alias("basket"))

# 分组执行FPGrowth
all_rules = grouped_data.groupBy("id").applyInPandas(
    process_id_baskets,
    schema="antecedent array<string>, consequent array<string>, confidence double, lift double, id string"
)

2. 若必须保留循环，提前物化中间结果 #

如果业务逻辑要求必须遍历每个ID，就在每个ID处理完后立即持久化规则DataFrame，强制Spark执行当前任务并缓存结果，避免执行计划无限堆积：

// Scala示例
val rulesList = new ListBuffer[DataFrame]()
for (id <- uniqueIds) {
  val filteredDF = transactionData.filter(col("id") === id).select("product").distinct()
  val basketDF = filteredDF.withColumn("basket", collect_set("product").over(Window.partitionBy("id")))
    .select("basket").distinct()
  
  val fpGrowth = new FPGrowth().setItemsCol("basket").setMinSupport(0.1).setMinConfidence(0.5)
  val model = fpGrowth.fit(basketDF)
  val rules = model.associationRules.withColumn("id", lit(id))
  
  // 持久化到磁盘，避免执行计划堆积
  rules.persist(StorageLevel.DISK_ONLY)
  rules.count() // 触发action，强制计算并缓存
  rulesList.append(rules)
}

val finalRulesDF = rulesList.reduce(_ union _)

3. 优化FPGrowth参数与数据格式 #

• 调大minSupport和minConfidence，减少每个ID生成的规则数量，从根源降低计算量。
• 确保basket列是去重后的商品集合，避免重复商品干扰计算，同时减少数据体积。

4. 检查执行计划的关键问题 #

查看物理执行计划时重点关注：

• 是否存在CartesianProduct（笛卡尔积）操作，这是执行计划爆炸的典型标志。
• Shuffle阶段的分区数是否合理，建议将spark.sql.shuffle.partitions设置为集群核心数的2-3倍。

验证方式 #

优化后先执行rulesDF.explain(true)对比执行计划复杂度，再用count()测试耗时，确认性能提升。

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