Spark自定义分区器简单示例

继上一篇：Spark分区器探索（HashPartitioner、RangePartitioner），现来看看自定义分区器。 继承 org.apache.spark.Partitioner 类并实现下面三个方法。 （1）numPartitions: Int:设置分区数。 （2）getPartition(key: Any): Int:返回给定key计算出的分区编号(0到numPartitions-1)。 （3）equals():判断分区器是否相等的方法。用于判断当前使用分区器对象与其他分区器实例是否相等相同，以此判断两个 RDD 的分区方式是否相同。

示例代码如下：

import org.apache.spark.rdd.RDD import org.apache.spark.{ SparkConf, SparkContext} class MyPartitioner(numParts:Int) extends org.apache.spark.Partitioner{ //number of partitions override def numPartitions: Int = numParts //get partition id base on key override def getPartition(key: Any): Int = { /* val ckey: String = key.toString ckey.substring(ckey.length-1).toInt%numParts*/ key.toString.toInt%numPartitions } override def equals(other: Any): Boolean = other match { case h: MyPartitioner => h.numPartitions == numPartitions case _ => false } } object partitionTest { def main(args: Array[String]): Unit = { val conf = new SparkConf().setAppName("TestPartitioner").setMaster("local[*]") val sc = new SparkContext(conf) //map to kv val r1: RDD[(Int, Int)] = sc.parallelize(1 to 10,5).map((_,1)) //partition by MyPatitioner, and map to resource val rp: RDD[Int] = r1.partitionBy(new MyPartitioner(2)).map(_._1) rp.mapPartitionsWithIndex((index,items)=>items.map((index,_))).collect.foreach(println) sc.stop() } }

运行结果：

(0,2) (0,4) (0,6) (0,8) (0,10) (1,1) (1,3) (1,5) (1,7) (1,9)

在以上代码中创建了一个分区数为5的RDD，随后运用partitionBy(new MyPartitioner(2)重新分为两个分区并输出对应分区数据。