Is your feature request related to a problem or challenge? Please describe what you are trying to do.

In DataFusion, a common operation is to repartition a RecordBatch by hashing one or more columns and dividing them into partition record batches using the "formula" hash % num_partitions.

The current approach is to create the indices that match and use them to take the individual arrays (see BatchPartitioner in datafusion).

This is relatively expensive however, as we visit the arrays num_partitions times in different places, leading to cache inefficient.

Describe the solution you'd like

I would like to have a method in arrow-rs that takes an Array with partition numbers (0..) and spreads the rows of the Array(s) into a Vec<Array> and/or Vec<RecordBatch>.

Describe alternatives you've considered

Additional context