How can I repartition RDD by key and then pack it to shards?
I have many files containing millions of rows in format:
id, created_date, some_value_a, some_value_b, some_value_c

This way of repartitioning was super slow and created for me over million of small ~500b files:
    rdd_df = rdd.toDF([&quot;id&quot;, &quot;created_time&quot;, &quot;a&quot;, &quot;b&quot;, &quot;c&quot;])
    rdd_df.write.partitionBy(&quot;id&quot;).csv(&quot;output&quot;) 

I would like to achieve output files, where each file contains like 10000 unique IDs and all their rows.
How could I achieve something like this?

You can repartition by adding a Random Salt key.
val totRows = rdd_df.count

val maxRowsForAnId = rdd_df.groupBy(&quot;id&quot;).count().agg(max(&quot;count&quot;))
val numParts1 = totRows/maxRowsForAnId


val totalUniqueIds = rdd_df.select(&quot;id&quot;).distinct.count
val numParts2 = totRows/(10000*totalUniqueIds)

val numPart = numParts1.min(numParts2)


rdd_df
 .repartition(numPart,col(&quot;id&quot;),rand)
 .csv(&quot;output&quot;) 

The main concept is each partition will be written as 1 file. SO you would have  bring your required rows in to 1 partition by repartition(numPart,col(&quot;id&quot;),rand).
The first 4-5 operations is just to calculate how many partitions we need to achieve almost 10000 ids per file.

Calculate assuming 10000 ids per partition
Corner case : if a single id has too many rows and doesn't fit in the above calculated partition size.
Hence we calculate no of paritition according to the largest count of ID present
Take min of the 2 noOfPartitons

rand is necessary so, that we can bring multiple IDs in a single partition
NOTE : Although this will give you larger files and each file will contain a set of unique ids for sure. But this involves shuffling , due to which your operation actually might be slower than the code you have mentioned in question.
