Spark分区内如何实现顺序处理并更新lookup，同时避免OOM？

你的需求本质是分区内有序的连通分量标签传播，可在保证多分区并行处理的同时，实现单分区内逐行顺序处理逻辑，且避免OOM，具体实现如下：

1 预处理保证分区与排序规则 #

首先给DF1同一业务分区（即你表里的p字段）的行新增顺序号，确保处理顺序符合预期，再按p字段重分区，让同一p的所有数据落到同一个Spark分区：

import org.apache.spark.sql.functions._
// 按业务分区p分组，给每个行分配全局唯一的顺序号，orderBy的字段替换为你实际用来判断行顺序的字段
val df1WithSeq = DF1.withColumn("row_seq", row_number().over(Window.partitionBy("p").orderBy("order_col")))
  .repartition(col("p"))

2 分区内流式迭代处理避免OOM #

用mapPartitions算子处理每个分区，仅在内存中维护当前分区的字符-分组映射，无需加载全量分区数据到内存，从根源避免OOM：

// 提前把DF2的已知分组映射广播到所有Executor，DF2不大的情况下优先用广播变量
val groupMapBroadcast = spark.sparkContext.broadcast(
  DF2.collect().map(row => (row.getAs[String]("char"), row.getAs[String]("group"))).toMap
)

val resultRDD = df1WithSeq.rdd.mapPartitions(iter => {
  // 初始化分组映射，仅加载当前分区需要的映射数据
  val localGroupMap = scala.collection.mutable.Map[String, String]()
  localGroupMap ++= groupMapBroadcast.value

  // 先对当前分区的行按顺序号排序，再转成迭代器流式处理
  val sortedIter = iter.toSeq.sortBy(_.getAs[Int]("row_seq")).iterator

  sortedIter.map(row => {
    val charA = row.getAs[String]("A")
    val charB = row.getAs[String]("B")
    val partitionId = row.getAs[String]("p")

    // 取已有的分组标签：优先取A的分组，没有则取B的分组，都没有则生成临时分组
    val curGroup = localGroupMap.get(charA)
      .orElse(localGroupMap.get(charB))
      .getOrElse(s"temp_${java.util.UUID.randomUUID().toString.take(8)}")

    // 更新两个字符的分组映射
    localGroupMap.put(charA, curGroup)
    localGroupMap.put(charB, curGroup)

    // 可根据需求返回行处理结果，或最终只返回localGroupMap的全量映射
    (charA, charB, partitionId, curGroup)
  })
})

// 如需得到最终全量字符分组映射，可合并所有分区的localGroupMap结果，或单独收集每个分区的映射表
val finalGroupDF = resultRDD.toDF("A", "B", "p", "final_group")

核心优化点 #

• 广播小表DF2的映射，避免每个分区重复拉取数据，降低内存占用
• 单分区仅维护字符-分组映射，处理完的行直接丢弃，内存占用仅和当前分区的唯一字符数成正比，和行数无关
• 若单分区唯一字符数过多，可将大的业务分区p拆分为多个子分区，进一步降低单分区内存压力

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