Native hash shuffle gathers rows one at a time even when a whole batch goes to one partition #7474

What is the problem the feature request solves?

For hash partitioning, MultiPartitionShuffleRepartitioner records one (batch, row) pair per row (PartitionIndices::Rows in native/shuffle/src/partitioners/multi_partition.rs) and builds each partition's output with interleave_record_batch (partitioned_batch_iterator.rs). Arrow's interleave copies primitive and string columns one row at a time (interleave_primitive, interleave_bytes) and merges contiguous runs only in its generic interleave_fallback, which Boolean columns use. So when the rows of an input batch already sit together, because they all go to the same partition, the writer still gathers them row by row. Positional round robin avoids this by recording runs.

That case is common for Iceberg writes. Spark plans RebalancePartitions on the distribution columns when a V2 write asks for a distribution it does not strictly require (DistributionAndOrderingUtils), which is what a hash-distributed Iceberg table gets for its partition columns. A job that writes one partition value per run, such as an hourly or daily ETL, then rebalances on a column that is constant for the whole run. Every row of every map task hashes to one reduce partition, which AQE splits back up.

In a production Spark 3.4 job with that shape (16 columns, mostly strings and longs, about 10 million rows per map task), partition interleaving time was 54% of native shuffle writer time in the map stage, about 4 s of a 46 s task, and 8% of the stage's task time. Each AQE-split reduce task read exactly one block, so all rows of each map task went to one partition.

Describe the potential solution

Record runs for rows that are already contiguous, without reordering anything. For example, when partition_starts shows that a whole input batch maps to one partition, store one run (batch, 0, num_rows) for that partition instead of num_rows row pairs, and at flush slice it, or concatenate runs up to batch_size, instead of interleaving. More generally, a batch could be split into runs wherever consecutive rows share a partition id, keeping row pairs for short runs.

This differs from the reorder proposed in #6917, which was declined because sorting each batch by partition writes small batches unless every partition keeps a coalescer. Nothing is reordered or buffered per partition here. Only batches whose rows are already grouped take the run path, and a run takes less index memory than the row pairs it replaces.

The open question is the index representation. PartitionIndices is chosen once per repartitioner today (rows for hash, runs for positional round robin), so this needs a per-partition mix of runs and rows, or a per-batch choice. A shuffle_writer bench case with a constant hash key would show the gain. The 4 s per task above is an upper bound.

Additional context

Related: #6917 (row gather, closed), #5905 (shuffle performance EPIC, which has no item for this), #7234 (removed the per-partition tail copy) and #6430 (Boolean gather specialization, which still gathers row by row for long runs).

Spark's own writer also pays per-row cost here, since UnsafeShuffleWriter serializes every record and sorts by partition id, so this is an optimization rather than a gap against Spark.