Filtering empty partitions in RDD - scala

Question

Is there a way to filter empty partitions in RDD? I have some empty partitions after partitioning and I can't use them in action method.
I use Apache Spark in Scala

Answer 1

This is my sample data
val sc = spark.sparkContext
val myDataFrame = spark.range(20).toDF("mycol").repartition($"mycol")
myDataFrame.show(false)
Output :
+-----+
|mycol|
+-----+
|19 |
|0 |
|7 |
|6 |
|9 |
|17 |
|5 |
|1 |
|10 |
|3 |
|12 |
|8 |
|11 |
|2 |
|4 |
|13 |
|18 |
|14 |
|15 |
|16 |
+-----+
In the above code when you do repartition on column then 200 paritions will be created since spark.sql.shuffle.partitions = 200 in that many are not used or empty partitions since data is just 10 numbers (we are trying to fit 20 numbers in to 200 partitions means.... most of the partitions are empty.... :-))
1) Prepare a long accumulator variable to quickly count non empty partitions.
2) Add all non empty partitions in to accumulator variable like below example.
val nonEmptyPartitions = sc.longAccumulator("nonEmptyPartitions")
myDataFrame.foreachPartition(partition =>
if (partition.length > 0) nonEmptyPartitions.add(1))
drop non empty partitions (means coalesce them... less shuffle/ minimum shuffle ).
print them.
val finalDf = myDataFrame.coalesce(nonEmptyPartitions.value.toInt)
println(s"nonEmptyPart : ${nonEmptyPartitions.value.toInt}")
println(s"df.rdd.partitions.length : ${myDataFrame.rdd.getNumPartitions}")
println(s"finalDf.rdd.partitions.length : ${finalDf.rdd.getNumPartitions}")
print them ...
Result :
nonEmptyPart : 20
df.rdd.partitions.length : 200
finalDf.rdd.partitions.length : 20
Proof that all non empty partitions are dropped...
myDataFrame.withColumn("partitionId", org.apache.spark.sql.functions.spark_partition_id)
.groupBy("partitionId")
.count
.show
Result printed partition wise record count :
+-----------+-----+
|partitionId|count|
+-----------+-----+
|128 |1 |
|190 |1 |
|140 |1 |
|164 |1 |
|5 |1 |
|154 |1 |
|112 |1 |
|107 |1 |
|4 |1 |
|49 |1 |
|69 |1 |
|77 |1 |
|45 |1 |
|121 |1 |
|143 |1 |
|58 |1 |
|11 |1 |
|150 |1 |
|68 |1 |
|116 |1 |
+-----------+-----+
Note :
Usage spark_partition_id is for demo/debug purpose only not for production purpose.
I reduced 200 partitions (due to repartition on column ) to 20 non empty partitions.
Conclusion :
Finally you got rid of extra empty partitions which doesnt have any data and avoided un necessary schedule to dummy tasks on empty partitions.

Answer 2

From the little info you provide, I can think about two options. Use mapPartitions and just catching empty iterators and returning them, while working on the non-empty ones.
rdd.mapPartitions { case iter => if(iter.isEmpty) { iter } else { ??? } }
Or you can use repartition, to get rid of the empty partitions.
rdd.repartition(10) // or any proper number

Answer 3

If you dont know the distinct values within the column, and wish to avoid having empty partitions, you can use countApproxDistinct() as:
df.repartition(df.rdd.countApproxDistinct().toInt)
If you wish to filter the existing empty partitions and repartition, you can use as solution suggeste by Sasa
OR:
df.repartition(df.mapPartitions(part => List(part.length).iterator).collect().count(_ != 0)).df.getNumPartitions)
However, in later case the partitions may or may not contain records by value.