Spark自定义RDD重分区

    在某些计算场景中，我们可能需要将两个有关联的数据输入的部分数据，也就是说RDD中的部分数据，需要聚合在同一个partition进行匹配计算，这个时候，我们就需要根据实际的业务需求，自定义RDD重分区。

下面结合代码，看看具体怎么实现重分区，spark内部提供了一个分区抽象类Partitioner:

package org.apache.spark
/**
 * An object that defines how the elements in a key-value pair RDD are partitioned by key.
 * Maps each key to a partition ID, from 0 to `numPartitions - 1`.
 */
abstract class Partitioner extends Serializable {
  def numPartitions: Int
  def getPartition(key: Any): Int
  def equals(other: Any): Boolean
}

    def numPartitions: Int：这个方法需要返回你想要创建分区的个数；
    def getPartition(key: Any): Int：这个函数需要对输入的key做计算，然后返回该key的分区ID
    equals()：这个是Java标准的判断相等的函数，之所以要求用户实现这个函数是因为Spark内部会比较两个RDD的分区是否一样。

具体实现一个，如下：

package com.ais.common.tools.partitioner

import com.ais.common.tools.{HashTool, IPTool}
import org.apache.spark.Partitioner


class IndexPartitioner(partitions: Int) extends Partitioner {
  def numPartitions = partitions

  def getPartition(key: Any): Int = {
    key match {
      case null => 0
      case iKey: Int => iKey % numPartitions
      case textKey: String => (HashTool.hash(textKey) % numPartitions).toInt
      case _ => 0
    }
  }

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

    这个重写的分区策略就是：

    如果key是个整形数值，则和分区数取余分配；如果key是个字符型的值，则计算他的哈希值再和分区数取余分配。这样我们只要在将两个RDD的key值保持一直，然后进行重分区，就能确保key一样的数据shuffe到同一个分区。当然也可以根据自己实际的业务来定义更复杂的分区策略。