Spark RDD之Partitioner

概述

Partitioner是shuffle过程中key重分区时的策略，即计算key决定k-v属于哪个分区，Transformation是宽依赖的算子时，父RDD和子RDD之间会进行shuffle操作，shuffle涉及到网络开销，由于父RDD和子RDD中的partition是多对多关系，所以容易造成partition中数据分配不均匀，导致数据的倾斜。

Shuffle

在MapReduce框架中，Shuffle是连接Map和和Reduce之间的桥梁，Map的输出要用到Reduce中必须经过Shuffle这个环节，Spark作为MapReduce框架的一种实现，自然也实现了shuffle的逻辑。

Shuffle描述的是一个过程，表现多对多的依赖关系，是Map和Reduce两个阶段的纽带，是对数据重新分区的过程，将经过mapTask后，key值相同的数据重新划分到同一个partition中。

Shuffle实现分为HashShuffleManager和SortShuffleManager，也可以自定义。在本文中，只讨论Shuffle过程中Partitioner的作用，Partitioner在shuffle的map端发挥作用，根据map端的key值，按照不同Partitioner的逻辑计算出reduce端的partitionId，以实现对相同key值数据的重新聚合。具体Shuffle的介绍可参考之后的文章。

Partitioner定义

抽象类Partitioner定义了两个抽象方法numPartitions和getPartition。getPartition方法根据输入的k-v对的key值返回一个Int型数据。

abstract class Partitioner extends Serializable {
  def numPartitions: Int
  def getPartition(key: Any): Int
}

该抽象类伴生了一个Partitioner对象如下，主要包含defaultPartitioner函数，该函数定义了Partitioner的默认选择策略。如果设置了spark.default.parallelism，则使用该值作为默认partitions，否则使用上游RDD中partitions最大的数作为默认partitions。过滤出上游RDD中包含partitioner的RDD，选择包含有最大partitions并且isEligible的RDD，将该RDD中的partitioner设置为分区策略，否则返回一个带有默认partitions数的HashPartitioner作为Partitioner。Partition个数应该和partition个数最多的上游RDD一致，不然可能会导致OOM异常。

object Partitioner {
  def defaultPartitioner(rdd: RDD[_], others: RDD[_]*): Partitioner = {
    val rdds = (Seq(rdd) ++ others)
    val hasPartitioner = rdds.filter(_.partitioner.exists(_.numPartitions > 0))

    val hasMaxPartitioner: Option[RDD[_]] = if (hasPartitioner.nonEmpty) {
      Some(hasPartitioner.maxBy(_.partitions.length))
    } else {
      None
    }

    val defaultNumPartitions = if (rdd.context.conf.contains("spark.default.parallelism")) {
      rdd.context.defaultParallelism
    } else {
      rdds.map(_.partitions.length).max
    }

    if (hasMaxPartitioner.nonEmpty && (isEligiblePartitioner(hasMaxPartitioner.get, rdds) ||
        defaultNumPartitions < hasMaxPartitioner.get.getNumPartitions)) {
      hasMaxPartitioner.get.partitioner.get
    } else {
      new HashPartitioner(defaultNumPartitions)
    }
  }
  
  private def isEligiblePartitioner(
     hasMaxPartitioner: RDD[_],
     rdds: Seq[RDD[_]]): Boolean = {
    val maxPartitions = rdds.map(_.partitions.length).max
    log10(maxPartitions) - log10(hasMaxPartitioner.getNumPartitions) < 1
  }
}

Partitioner实现类

Partitioner主要有两个实现类：HashPartitioner和RangePartitioner，上面提到过HashPartitioner是默认的重分区方式。

HashPartitioner

numPartitions方法返回传入的分区数，getPartition方法使用key的hashCode值对分区数取模得到PartitionId，写入到对应的bucket中。因为Arrays的hashCodes值并不依赖于arrays的内容，导致hash函数将无法根据key值进行重分区，所以HashPartitioner不支持RDD[Array[_]] 或 RDD[(Array[_], _)]。

class HashPartitioner(partitions: Int) extends Partitioner {
  require(partitions >= 0, s"Number of partitions ($partitions) cannot be negative.")

  def numPartitions: Int = partitions

  def getPartition(key: Any): Int = key match {
    case null => 0
    case _ => Utils.nonNegativeMod(key.hashCode, numPartitions)
  }

  override def equals(other: Any): Boolean = other match {
    case h: HashPartitioner =>
      h.numPartitions == numPartitions
    case _ =>
      false
  }

  override def hashCode: Int = numPartitions
}

RangePartitioner

RangePartitioner继承了Partitoner，Partitioner中定义了两个抽象方法numPartitions和getPartition，前者是重分区后的partition数，后者是获取某个key值重新分区后的partitionId，我们着重关注RangePartitioner类中这两个函数的实现。但我们先对rangeBounds这一函数进行分析，通过分析源码可以发现，该函数首先设置了一个样本大小(默认为每个partition包含20条数据)，再对父RDD中的每个partition需要抽取的样本数进行计算，调用工具类中的reservoir sampling算法对每个partition进行分别抽样。每个分区的记录数*fraction如果大于该partition中设定的样本数(这是由于不同的分区中包含的数据量不同，数据量较大的分区中抽样数将会大于平均值)，则用imbalancedPartitions存储，并重新抽样以确保每个分区中都有足够数量的样本。最后计算权重——分区记录总数/分区样本数，并调用determineBounds方法求分区分隔符。rangeBounds函数返回了一个数组，该数组中存储的是重分区的分隔符，简而言之，该方法的作用是尽可能均衡地将原有的数据进行重分区，并返回用于重分区的分隔符。例如数据为"a c", “a b”, “b c”, “b d”, “c d”，通过以上函数可能会将b、c、d作为分隔符，所有key为a的将会被重分区到同一个partition中(c和d类似)。

class RangePartitioner[K : Ordering : ClassTag, V](
    partitions: Int,
    rdd: RDD[_ <: Product2[K, V]],
    private var ascending: Boolean = true,
    val samplePointsPerPartitionHint: Int = 20)
  extends Partitioner {

  def this(partitions: Int, rdd: RDD[_ <: Product2[K, V]], ascending: Boolean) = {
    this(partitions, rdd, ascending, samplePointsPerPartitionHint = 20)
  }
  
  private var rangeBounds: Array[K] = {
    if (partitions <= 1) {
      Array.empty
    } else {
      val sampleSize = math.min(samplePointsPerPartitionHint.toDouble * partitions, 1e6)
      val sampleSizePerPartition = math.ceil(3.0 * sampleSize / rdd.partitions.length).toInt
      val (numItems, sketched) = RangePartitioner.sketch(rdd.map(_._1), sampleSizePerPartition)
      if (numItems == 0L) {
        Array.empty
      } else {
        val fraction = math.min(sampleSize / math.max(numItems, 1L), 1.0)
        val candidates = ArrayBuffer.empty[(K, Float)]
        val imbalancedPartitions = mutable.Set.empty[Int]
        sketched.foreach { case (idx, n, sample) =>
          if (fraction * n > sampleSizePerPartition) {
            imbalancedPartitions += idx
          } else {
            val weight = (n.toDouble / sample.length).toFloat
            for (key <- sample) {
              candidates += ((key, weight))
            }
          }
        }
        if (imbalancedPartitions.nonEmpty) {
          val imbalanced = new PartitionPruningRDD(rdd.map(_._1), imbalancedPartitions.contains)
          val seed = byteswap32(-rdd.id - 1)
          val reSampled = imbalanced.sample(withReplacement = false, fraction, seed).collect()
          val weight = (1.0 / fraction).toFloat
          candidates ++= reSampled.map(x => (x, weight))
        }
        RangePartitioner.determineBounds(candidates, math.min(partitions, candidates.size))
      }
    }
  }
}

接下来关注RangePartitioner类中两个重要函数——numPartitions和getPartition的实现。首先我们观察numPartitions，它的返回值为数组rangeBounds的长度加上1，这很容易理解，如果我们重分区的分隔符有n个，那么我们重分区后的partition便是n+1个。getPartition先判断间隔符的个数，如果小于128则直接遍历比较key和分隔符得到PartitionId，否则使用二分查找，并对边界条件进行了判断，最后根据构造RangePartitioner时传入的ascending参数确定是升序或降序返回PartitionId。

  def numPartitions: Int = rangeBounds.length + 1

  def getPartition(key: Any): Int = {
    val k = key.asInstanceOf[K]
    var partition = 0
    if (rangeBounds.length <= 128) {
      while (partition < rangeBounds.length && ordering.gt(k, rangeBounds(partition))) {
        partition += 1
      }
    } else {
      partition = binarySearch(rangeBounds, k)
      if (partition < 0) {
        partition = -partition-1
      }
      if (partition > rangeBounds.length) {
        partition = rangeBounds.length
      }
    }
    if (ascending) {
      partition
    } else {
      rangeBounds.length - partition
    }
  }

总结RangePartitioner的过程为以下几步：

• 使用reservoir Sample方法对每个Partition进行分别抽样
• 对数据量大(大于sampleSizePerPartition)的分区进行重新抽样
• 由权重信息计算出分区分隔符rangeBounds
• 由rangeBounds计算分区数和key的所属分区

总结

综上所述，Partitioner主要的作用是在shuffle过程中对数据的partition进行重新分区，其主要实现的函数为：

1. 获得重新分区的分区个数；
2. 针对某个k-v对根据其中的key，将它按特定的方法重新分区。

Partitioner的具体实现类HashPartitioner和RangePartitioner对以上两个方法进行了重写。HashPartitioner主要是通过对原partition中数据的key值进行hash后，根据key的hash值将放入bucket中，再将不同partition中的bucket合并实现重新分区；RangePartitioner是先根据所有partition中数据的分布情况，尽可能均匀地构造出重分区的分隔符，再将数据的key值根据分隔符进行重新分区。HashPartitioner是大部分transformation的默认实现，sortBy、sortByKey使用RangePartitioner实现，也可以自定义Partitioner。以下是自定义Partitioner的一个例子，该Partitioner以key的长度进行分区：

class CustomPartitioner(partitions: Int) extends Partitioner {

  def numPartitions: Int = partitions

  def getPartition(key: Any): Int = {
    val k = key.asInstanceOf[String]
    return k.length() & (partitions -1)
  }
}

具体的调用如下：

    val data = sc.parallelize(List("a c", "a b", "bb c", "bb d", "c d"), 2)
    data.flatMap(_.split(" ")).map((_, 1))
      .partitionBy(new CustomPartitioner(2)).reduceByKey(_ + _)
        .collect()

我们可以看到，Partitioner并不作用于reduceByKey而是作用于map，是因为它是对shuffle之前的partition进行重分区策略的定义，在reduceByKey过程中，会调用当前RDD的dep(父RDD)中定义的Partitioner进行重分区。