Spark DataFrame - How to partition the data based on condition

Question

Have some employee data set. in that i need to partition based employee salary based on some condition. Created DataFrame and converted to Custom DataFrame Object. Created Custom Partition for salary.

class SalaryPartition(override val numPartitions: Int) extends Partitioner {

  override def getPartition(key: Any): Int =
    {
      import com.csc.emp.spark.tutorial.PartitonObj._
      key.asInstanceOf[Emp].EMPLOYEE_ID match {
        case salary if salary < 10000 => 1
        case salary if salary >= 10001 && salary < 20000 => 2
        case _ => 3
      }

    }

}

Question how can i invoke\\call my custome partition. Couldn't find partitionBy in dataframe. Have any alternative way?

Answer 1

Just code for my comment:

val empDS = List(Emp(5, 1000), Emp(4, 15000), Emp(3, 30000), Emp(2, 2000)).toDS()
println(s"Original partitions number: ${empDS.rdd.partitions.size}")
println("-- Original partition: data --")
empDS.rdd.mapPartitionsWithIndex((index, it) => {
  it.foreach(r => println(s"Partition $index: $r")); it
}).count()

val getSalaryGrade = (salary: Int) => salary match {
  case salary if salary < 10000 => 1
  case salary if salary >= 10001 && salary < 20000 => 2
  case _ => 3
}
val getSalaryGradeUDF = udf(getSalaryGrade)
val salaryGraded = empDS.withColumn("salaryGrade", getSalaryGradeUDF($"salary"))

val repartitioned = salaryGraded.repartition($"salaryGrade")
println
println(s"Partitions number after: ${repartitioned.rdd.partitions.size}")
println("-- Reparitioned partition: data --")

repartitioned.as[Emp].rdd.mapPartitionsWithIndex((index, it) => {
  it.foreach(r => println(s"Partition $index: $r")); it
}).count()

Output is:

Original partitions number: 2
-- Original partition: data --
Partition 1: Emp(3,30000)
Partition 0: Emp(5,1000)
Partition 1: Emp(2,2000)
Partition 0: Emp(4,15000)

Partitions number after: 5
-- Reparitioned partition: data --
Partition 1: Emp(3,30000)
Partition 3: Emp(5,1000)
Partition 3: Emp(2,2000)
Partition 4: Emp(4,15000)

Note: guess, several partitions possible with the same "salaryGrade".

Advice: "groupBy" or similar looks like more reliable solution.

For stay with Dataset entities, "groupByKey" can be used:

empDS.groupByKey(x => getSalaryGrade(x.salary)).mapGroups((index, it) => {
  it.foreach(r => println(s"Group $index: $r")); index
}).count()

Output:

Group 1: Emp(5,1000)
Group 3: Emp(3,30000)
Group 1: Emp(2,2000)
Group 2: Emp(4,15000)