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.partitionerimport com.ais.common.tools.{HashTool, IPTool}import org.apache.spark.Partitionerclass 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到同一个分区。当然也可以根据自己实际的业务来定义更复杂的分区策略。