Support partitioning functionality

[Dandandan]

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

[tustvold]

Makes sense to me, of course assuming it empirically shows a performance improvement over using take. It occurs to me this would trade write locality for read locality, which may not be beneficial.

[Dandandan]

Yeah agree, we'll have to benchmark it.

I expect though that write locality won't be as bad as the current read locality...

[richox]

it is better to take multiple input arrays than a single array? it makes many small fragment arrays (which has less performance and wastes a lot of memory) if the number of partitions is big.