perf(index): bucket shuffle rows by partition id instead of sorting keys

[LuciferYang]

Closes #9381

sort_to_interleave_indices built a (part_id, batch_idx, row_idx) tuple for every row in a shuffle flush group and sorted them by partition id. The key is bounded by num_partitions (the loop right after the sort rejects anything else), so the grouping can be had by counting rows per partition, turning the counts into prefix sums, and scattering each row into its partition's run. Two linear passes, no per-row tuple.

Measured in release, both implementations alternating in one process, uniformly distributed partition ids: 1M rows over 1024 batches into 256 partitions goes 15.2 ms to 2.00 ms, the same rows into 4096 partitions 15.3 ms to 1.92 ms, and 8.4M rows over 8192 batches into 4096 partitions 125.5 ms to 20.3 ms. The function runs once per flush group inside spawn_cpu, and a group holds up to shuffle_partition_batches batches (10240 by default), so the large shape is the realistic one.

Both rejections the old code performed are kept, now in the counting pass: a null partition id still errors instead of being read through values() as partition 0, and an id outside [0, num_partitions) still errors. One thing does become stricter in a useful way: sort_unstable_by_key left the order of rows within a partition unspecified, while bucketing emits them batch by batch and in row order within a batch, so the same input now produces the same partition layout every time.

How was this patch tested?

test_two_file_shuffler_groups_two_batches_by_partition covers what the existing tests did not: two non-empty batches in a single flush group, with partition ids interleaved across both, asserting the exact rows and order in each partition. It is the only test in the module that fails if the scatter walks the batches in the wrong order, since every partition size stays correct in that case. The existing null and out-of-range tests still cover the rejections, and a wrong prefix sum fails six of the existing shuffle tests.

[lance-gatekeeper]

✅ Gate recommendation: approve.

The prefix-sum scatter preserves the shuffle contract—null/range rejection, per-partition counts, and row-to-batch mapping—while removing the comparison sort and per-row key allocation. Its single preallocated output is preferable to per-partition buffers because it retains the existing bounded interleave representation.