Spark如何保证Dataset按唯一分区拆分且同键值数据同分区

问题背景 #

• 定义样例类：case class Data(col: String, ...)
• 现有Dataset[Data]类型数据集ds，其中col字段部分取值为'a'，其余为'b'等其他值
• 需求：分别全量处理col='a'、col='b'的数据，要求所有col取值为'a'的数据必须归属同一个分区
• 初始尝试方案1：

  ds.repartition(col("col")).mapPartition(data => ???)
  

• 初始尝试方案2（可实现需求但会触发额外Action，开销较高）：

  val nbDistinct = ds.select("col").distinct.count
  ds.repartition(nbDistinct , col("col")).mapPartition(data => ???)
  

问题1解答：默认repartition(col("col"))是否能满足分区要求？ #

不能完全满足。
ds.repartition(col("col"))默认使用Hash分区规则，不指定分区数时，会采用spark.sql.shuffle.partitions配置的默认分区数（默认值200，生产环境通常设置为200~2000）。这个实现仅能保证同一个col值的所有数据一定落在同一个分区，不会拆分到多个分区，但完全无法保证单个分区内只有一个col值：只要两个不同col值的哈希值对分区数取模结果相同，就会被分到同一个分区，出现'a'和'b'等其他值混存的情况。如果col的去重值总数超过了默认分区数，这种混存是必然发生的。

问题2解答：无额外Action开销的强保证分区方案 #

根据核心诉求不同，可选择两类低开销方案：

场景1：核心需求仅为「同col值全量数据在同一分区，处理时不需要跨分区拉取」 #

这种场景根本不需要强制要求单分区只存一个col值，直接用ds.repartition(col("col"))就足够。
因为同col值的数据已经100%集中在唯一的分区里，你只需要在mapPartitions逻辑内，把当前分区的所有数据按col字段做一次内存分组，每个分组就是对应col值的全量数据，直接处理即可。哪怕一个分区里有多个col值，也不会出现某个col值的数据散落在多个分区、拿不到全量的问题，完全没有额外的Shuffle或者Action开销。

场景2：硬要求「单分区内只能有同一个col值的数据，不能混存」 #

这种场景不需要提前触发distinct.count计算去重值数量，可选方案如下：

• 如果业务上明确知道col字段的所有枚举取值（比如固定只有'a'/'b'/'c'三类值），直接在repartition时指定分区数为枚举值总数即可：

  // 以col固定为a、b两个值为例
  ds.repartition(2, col("col")).mapPartition(data => ???)
  

  只要指定的分区数大于等于实际col的去重值总数，Hash分区下不同值碰撞到同一分区的概率极低；生产环境如果担心极小概率的哈希碰撞，可以把分区数设置为枚举值数量的1.5~2倍，空分区的计算开销可以忽略，远低于提前算distinct的代价。
• 如果col取值范围未知，可以转用RDD层的自定义分区器实现，不需要提前统计去重值数量，也不会触发额外Job：

  import org.apache.spark.Partitioner
  
  // 自定义分区器，为每个出现的key分配独立分区ID
  class ExclusiveKeyPartitioner(expectMaxKeyNum: Int) extends Partitioner {
    private val key2PartId = new java.util.concurrent.ConcurrentHashMap[String, Int]()
    private var autoIncrementId = 0
    override def numPartitions: Int = expectMaxKeyNum
    override def getPartition(key: Any): Int = {
      val realKey = key.asInstanceOf[String]
      key2PartId.computeIfAbsent(realKey, _ => {
        val currentId = autoIncrementId
        autoIncrementId += 1
        currentId % numPartitions
      })
    }
  }
  
  // 用自定义分区器重分区
  val partitionedDs = ds.rdd
    .map(row => (row.col, row))
    .partitionBy(new ExclusiveKeyPartitioner(1000)) // 分区数设置为预估的col最大去重数量即可，建议留2倍左右冗余
    .values
    .toDS()
  
  partitionedDs.mapPartition(data => ???)
  

  这种方式会在Shuffle的Map阶段自动为每个新出现的col值分配独立的分区ID，只要设置的预估最大分区数大于实际col去重值数量，就可以完全避免不同col值混存的问题，开销远低于额外触发一次全表distinct计算。

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