RDD-自定义分区器案例

package org.example

import org.apache.spark.{Partitioner, SparkConf, SparkContext}


case class Order(id: Int, price: Double, info: String) {
  override def toString: String = s"$id, $price, $info"
}

class orderPartitioner extends Partitioner{

  override def numPartitions: Int = 3

  override def getPartition(key: Any): Int = {
    //0-1000 => 1
    //1001-2000 => 2
    //3
    if (key.asInstanceOf[Int] <= 1000) {
      0
    } else if (key.toString.toInt <= 2000) {
      1
    } else {
      2
    }
  }
}

object PartitionOrder {
  def main(args: Array[String]): Unit = {
    val conf = new SparkConf().setAppName("Partition").setMaster("local[*]")
    val sc = new SparkContext(conf)

    //读入data/order.csv 创建RDD
    val orderRDD = sc.textFile("data/order.csv")

    val rdd1 = orderRDD.map(line => {
      val fields = line.split(",")
      val order = Order(fields(0).toInt, fields(1).toDouble, fields(2))
      (order.id, order)
    })

    val rdd2 = rdd1.partitionBy(new orderPartitioner)

    rdd2.map(x => x._2).saveAsTextFile("data/output1")

    rdd2.mapPartitions(iter => {
      var count = 0

      var sum = 0.0
      iter.foreach(x => {
        sum += x._2.price
        count += 1
      })
      Iterator(s"${count}件, ${sum}元")
    })saveAsTextFile("data/output2")
  }