rdd-partition-split.scala

// Split an rdd according to its partition number
def splitByPartition[T:ClassTag](rdd: RDD[T], hotPartitions:Int): (RDD[T], RDD[T]) = {
  val splits = rdd.mapPartitions { iter =>
    val partId = TaskContext.get.partitionId
    val left  = if (partId <  hotPartitions) iter else empty
    val right = if (partId >= hotPartitions) iter else empty
    Seq(left, right).iterator
  }

  val left = splits.mapPartitions { iter => iter.next().toIterator}
  val right = splits.mapPartitions { iter =>
    iter.next()
    iter.next().toIterator
  }
  (left, right)
}