experiment: avoid final partition reordering

Rachel Int · Rachel Int · commit ded74525110f · 2026-07-20T20:06:32.000+08:00
diff --git a/datafusion/physical-plan/src/aggregates/hash_aggregate.rs b/datafusion/physical-plan/src/aggregates/hash_aggregate.rs
@@ -42,9 +42,7 @@ use super::aggregate_hash_table::{
     AggregateHashTable, FinalMarker, PartialMarker, PartialSkipMarker,
 };
 use super::skip_partial::SkipAggregationProbe;
-use super::{
-    AggregateExec, SubpartitionReorderBuffer, partition_runs, reorder_by_subpartition,
-};
+use super::{AggregateExec, partition_runs, partition_runs_from_column};
 use crate::metrics::{
     BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput, SpillMetrics,
 };
@@ -56,7 +54,6 @@ struct FinalPartitionRunState {
     run_sizes: BTreeMap<usize, usize>,
     total_runs_size: usize,
     active_partition_size: usize,
-    subpartition_reorder_buffer: SubpartitionReorderBuffer,
 }
 
 impl FinalPartitionRunState {
@@ -66,7 +63,6 @@ impl FinalPartitionRunState {
             run_sizes: BTreeMap::new(),
             total_runs_size: 0,
             active_partition_size: 0,
-            subpartition_reorder_buffer: SubpartitionReorderBuffer::new(),
         }
     }
 
@@ -866,18 +862,17 @@ impl FinalHashAggregateStream {
     }
 
     fn stage_partition_runs(&mut self, batch: &RecordBatch) -> Result<Option<usize>> {
-        if let Some(state) = self.partition_run_state.as_mut()
-            && let Some(batch) =
-                reorder_by_subpartition(batch, &mut state.subpartition_reorder_buffer)?
-        {
+        if let Some((batch, runs)) = partition_runs_from_column(batch)? {
             let mut offset = 0;
             let mut staged_memory = 0;
-            let num_runs = state.subpartition_reorder_buffer.runs().len();
-            for run_idx in 0..num_runs {
-                let run = state.subpartition_reorder_buffer.runs()[run_idx];
+            for run in runs {
                 let run_batch = batch.slice(offset, run.len);
                 offset += run.len;
-                staged_memory += state.stage_batch(run_batch, run.relative_partition);
+                staged_memory += self
+                    .partition_run_state
+                    .as_mut()
+                    .map(|state| state.stage_batch(run_batch, run.relative_partition))
+                    .unwrap_or(0);
             }
             return Ok(Some(staged_memory));
         }
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -72,8 +72,8 @@ use itertools::Itertools;
 #[cfg(test)]
 pub(crate) use partition_runs::set_partition_runs_metadata;
 pub(crate) use partition_runs::{
-    PartitionRun, SubpartitionReorderBuffer, append_subpartition_column, partition_runs,
-    reorder_by_subpartition, subpartition_schema,
+    PartitionRun, append_subpartition_column, partition_runs, partition_runs_from_column,
+    subpartition_schema,
 };
 use topk::hash_table::is_supported_hash_key_type;
 use topk::heap::is_supported_heap_type;
diff --git a/datafusion/physical-plan/src/aggregates/partition_runs.rs b/datafusion/physical-plan/src/aggregates/partition_runs.rs
@@ -21,15 +21,12 @@
 //! internal contracts used to describe relative aggregate subpartitions passed
 //! from repartition to final hash aggregation.
 //!
-//! Entry points: [`append_subpartition_column`], [`reorder_by_subpartition`],
+//! Entry points: [`append_subpartition_column`], [`partition_runs_from_column`],
 //! [`set_partition_runs_metadata`] and [`partition_runs`].
 
 use std::sync::Arc;
 
-use arrow::array::{
-    Array, ArrayRef, MutableArrayData, RecordBatch, RecordBatchOptions, UInt32Array,
-    make_array,
-};
+use arrow::array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, UInt32Array};
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion_common::cast::as_uint32_array;
 use datafusion_common::{Result, internal_err};
@@ -109,41 +106,11 @@ pub(crate) fn append_subpartition_column(
     RecordBatch::try_new_with_options(schema, columns, &options).map_err(Into::into)
 }
 
-/// Reusable scratch space for [`reorder_by_subpartition`].
-#[derive(Default)]
-pub(crate) struct SubpartitionReorderBuffer {
-    counts: Vec<usize>,
-    offsets: Vec<usize>,
-    source_runs: Vec<SourcePartitionRun>,
-    runs: Vec<PartitionRun>,
-}
-
-impl SubpartitionReorderBuffer {
-    /// Create an empty reorder scratch buffer.
-    pub(crate) fn new() -> Self {
-        Self::default()
-    }
-
-    /// Return partition runs produced by the latest reorder.
-    pub(crate) fn runs(&self) -> &[PartitionRun] {
-        &self.runs
-    }
-
-    fn clear(&mut self) {
-        self.counts.clear();
-        self.offsets.clear();
-        self.source_runs.clear();
-        self.runs.clear();
-    }
-}
-
-/// Reorder `batch` by its internal subpartition column and strip that column.
-pub(crate) fn reorder_by_subpartition(
+/// Read contiguous partition runs from the internal subpartition column and
+/// return a batch that omits the internal column without reordering any rows.
+pub(crate) fn partition_runs_from_column(
     batch: &RecordBatch,
-    buffer: &mut SubpartitionReorderBuffer,
-) -> Result<Option<RecordBatch>> {
-    buffer.clear();
-
+) -> Result<Option<(RecordBatch, Vec<PartitionRun>)>> {
     let Some(subpartition_idx) = subpartition_column_index(batch.schema_ref()) else {
         return Ok(None);
     };
@@ -155,129 +122,29 @@ pub(crate) fn reorder_by_subpartition(
         );
     }
 
-    let num_rows = batch.num_rows();
-    if num_rows > u32::MAX as usize {
-        return internal_err!(
-            "Hash aggregate batch row count {num_rows} does not fit UInt32"
-        );
-    }
-
-    for relative_partition in subpartitions.values().iter().copied() {
-        let relative_partition = relative_partition as usize;
-        if relative_partition >= buffer.counts.len() {
-            buffer.counts.resize(relative_partition + 1, 0);
-        }
-        buffer.counts[relative_partition] += 1;
-    }
-
-    buffer.offsets.resize(buffer.counts.len(), 0);
-    let mut output_offset = 0;
-    for (relative_partition, len) in buffer.counts.iter().copied().enumerate() {
-        buffer.offsets[relative_partition] = output_offset;
-        if len != 0 {
-            buffer
-                .runs
-                .push(PartitionRun::new(relative_partition, len)?);
-        }
-        output_offset += len;
-    }
-    debug_assert_eq!(output_offset, num_rows);
-
-    let mut is_identity = true;
-    let mut values = subpartitions.values().iter().copied().enumerate();
-    if let Some((run_start, first_partition)) = values.next() {
+    let mut runs = Vec::new();
+    let mut values = subpartitions.values().iter().copied();
+    if let Some(first_partition) = values.next() {
         let mut current_partition = first_partition as usize;
-        let mut current_start = run_start;
         let mut current_len = 1;
-        for (row_idx, relative_partition) in values {
+        for relative_partition in values {
             let relative_partition = relative_partition as usize;
             if relative_partition == current_partition {
                 current_len += 1;
             } else {
-                push_source_run(
-                    buffer,
-                    current_partition,
-                    current_start,
-                    current_len,
-                    &mut is_identity,
-                );
+                runs.push(PartitionRun::new(current_partition, current_len)?);
                 current_partition = relative_partition;
-                current_start = row_idx;
                 current_len = 1;
             }
         }
-        push_source_run(
-            buffer,
-            current_partition,
-            current_start,
-            current_len,
-            &mut is_identity,
-        );
+        runs.push(PartitionRun::new(current_partition, current_len)?);
     }
 
     let schema = strip_subpartition_schema(batch.schema_ref(), subpartition_idx);
     let columns = strip_subpartition_columns(batch, subpartition_idx);
-    let output = if is_identity {
-        RecordBatch::try_new(schema, columns)?
-    } else {
-        buffer
-            .source_runs
-            .sort_unstable_by_key(|run| run.output_start);
-        let columns =
-            reorder_columns_by_source_runs(&columns, &buffer.source_runs, num_rows);
-        let options = RecordBatchOptions::new().with_row_count(Some(num_rows));
-        RecordBatch::try_new_with_options(schema, columns, &options)?
-    };
-
-    Ok(Some(output))
-}
-
-#[derive(Debug, Clone, Copy)]
-struct SourcePartitionRun {
-    start: usize,
-    len: usize,
-    output_start: usize,
-}
-
-fn push_source_run(
-    buffer: &mut SubpartitionReorderBuffer,
-    relative_partition: usize,
-    start: usize,
-    len: usize,
-    is_identity: &mut bool,
-) {
-    let output_start = buffer.offsets[relative_partition];
-    buffer.offsets[relative_partition] += len;
-    *is_identity &= output_start == start;
-    buffer.source_runs.push(SourcePartitionRun {
-        start,
-        len,
-        output_start,
-    });
-}
-
-fn reorder_columns_by_source_runs(
-    columns: &[ArrayRef],
-    source_runs: &[SourcePartitionRun],
-    num_rows: usize,
-) -> Vec<ArrayRef> {
-    columns
-        .iter()
-        .map(|column| reorder_array_by_source_runs(column, source_runs, num_rows))
-        .collect()
-}
-
-fn reorder_array_by_source_runs(
-    array: &ArrayRef,
-    source_runs: &[SourcePartitionRun],
-    num_rows: usize,
-) -> ArrayRef {
-    let array_data = array.to_data();
-    let mut mutable = MutableArrayData::new(vec![&array_data], false, num_rows);
-    for run in source_runs {
-        mutable.extend(0, run.start, run.start + run.len);
-    }
-    make_array(mutable.freeze())
+    let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
+    let batch = RecordBatch::try_new_with_options(schema, columns, &options)?;
+    Ok(Some((batch, runs)))
 }
 
 /// Return partition runs encoded on `schema`, if present.
@@ -436,7 +303,7 @@ mod tests {
     }
 
     #[test]
-    fn test_reorder_by_subpartition_groups_rows_and_strips_column() -> Result<()> {
+    fn test_partition_runs_from_column_preserves_row_order() -> Result<()> {
         let schema =
             Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, false)]));
         let batch = RecordBatch::try_new(
@@ -452,14 +319,17 @@ mod tests {
             ],
         )?;
 
-        let mut buffer = SubpartitionReorderBuffer::new();
-        let batch = reorder_by_subpartition(&batch, &mut buffer)?.unwrap();
+        let (batch, runs) = partition_runs_from_column(&batch)?.unwrap();
         let values = as_uint32_array(batch.column(0).as_ref())?;
-        assert_eq!(values.values(), &[11, 12, 10, 13]);
+        assert_eq!(values.values(), &[10, 11, 12, 13]);
         assert_eq!(batch.num_columns(), 1);
         assert_eq!(
-            buffer.runs(),
-            &[PartitionRun::new(0, 2)?, PartitionRun::new(1, 2)?]
+            runs,
+            [
+                PartitionRun::new(1, 1)?,
+                PartitionRun::new(0, 2)?,
+                PartitionRun::new(1, 1)?,
+            ]
         );
         Ok(())
     }
