[Data] boost hash_partition w/ sort_indices + zero-copy slices

[owenowenisme]

Description

hash_partition previously did three expensive things in sequence:
N = num_partition | R = num_rows

1. Built per-partition index arrays via N × np.where(part_ids == p) — O(N · R) scans
2. Defragmented the input via try_combine_chunked_columns(table) — a full-table copy
3. Ran N independent table.take(indices[p]) calls

This change replaces all three with:

1. pyarrow.compute.sort_indices(partition_ids) — radix sort on integers, one O(R) pass
2. One take_table(table, sort_indices) on the original (possibly chunked) input
3. N zero-copy Table.slice() calls

The N takes together form a permutation of the table, so consolidating them into one sort + N zero-copy slices is equivalent and strictly cheaper (fixed take overhead paid once instead of N times). The defrag copy can also be removed: the original Arrow problem (apache/arrow#35126) is that every take on a chunked table internally concatenates all chunks first, so try_combine_chunked_columns exists to pay that concat once externally and let the subsequent N takes use the fast path. By calling take only once, the internal concat happens just once anyway — the external defrag becomes redundant. And because the take output already arranges each partition's rows contiguously, we can carve out the N partitions with zero-copy slices instead of materializing a second copy — which would be another 1GB for a 1GB input.

Benchmark: 1GB block → 1000 partitions

Single thread, PyArrow 23.0.1. K = number of chunks in the input table; K=256 mirrors realistic multi-chunk input.

Block shape	K	Time before	Time after	Speedup	Peak Arrow before	Peak Arrow after
16M rows × 8 int64	1	6542 ms	939 ms	6.96×	1024 MB (no copy needed)	1152 MB
16M rows × 8 int64	256	6887 ms	1057 ms	6.51×	2048 MB	1280 MB
8M rows × 16 int64	1	4818 ms	825 ms	5.84×	1024 MB (no copy needed)	1088 MB
8M rows × 16 int64	256	5086 ms	982 ms	5.18×	2048 MB	1152 MB
2M rows × 64 int64	1	2468 ms	322 ms	7.66×	1026 MB (no copy needed)	1040 MB
2M rows × 64 int64	256	2216 ms	369 ms	6.01×	2050 MB	1056 MB

• Throughput: 5–8× faster across all shapes.
• Peak Arrow allocation on chunked inputs (K=256): ~2.0 GB → ~1.1 GB (~40% reduction) — the input no longer has to coexist with a defragmented copy.

	Before	After	Speedup
aggregate_groups (84 groups, mean)	61 s	40 s	1.53×

Related issues

Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234".

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[gemini-code-assist]

Code Review

This pull request optimizes the hash_partition function in python/ray/data/_internal/arrow_ops/transform_pyarrow.py by replacing multiple independent take operations with a single take followed by zero-copy slices. The implementation uses pyarrow.compute.sort_indices to sort the table by partition ID and numpy.bincount to determine partition boundaries. I have no feedback to provide as there were no review comments to evaluate.