feat: reuse partial aggregate hashes across repartition

datafusion/physical-plan/benches/dictionary_group_values.rs

@@ -112,7 +112,8 @@ fn bench_intern_emit(c: &mut Criterion) {
                             )
                         },
                         |(gv, groups)| {
-                            gv.intern(std::slice::from_ref(&array), groups).unwrap();
+                            gv.intern(std::slice::from_ref(&array), groups, None)
+                                .unwrap();
                             black_box(&*groups);
                             black_box(gv.emit(EmitTo::All).unwrap());
                         },
@@ -158,7 +159,8 @@ fn bench_repeated_intern_emit(c: &mut Criterion) {
                         },
                         |(gv, groups)| {
                             for arr in &batches {
-                                gv.intern(std::slice::from_ref(arr), groups).unwrap();
+                                gv.intern(std::slice::from_ref(arr), groups, None)
+                                    .unwrap();
                                 black_box(&*groups);
                             }
                             black_box(gv.emit(EmitTo::All).unwrap());

datafusion/physical-plan/benches/multi_group_by.rs

@@ -97,7 +97,7 @@ fn bench_intern(
 ) {
     for batch in batches {
         groups.clear();
-        gv.intern(batch, groups).unwrap();
+        gv.intern(batch, groups, None).unwrap();
     }
     black_box(&*groups);
 }

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs

@@ -32,6 +32,7 @@ use crate::aggregates::grouped_hash_stream::create_group_accumulator;
 use crate::aggregates::order::GroupOrdering;
 use crate::aggregates::{
     AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by,
+    schema_with_group_hash,
 };
 
 /// Marker for raw rows -> partial state aggregation.
@@ -343,6 +344,23 @@ impl MaterializedAggregateOutput {
     }
 }
 
+pub(super) fn try_new_internal_batch(
+    output_schema: SchemaRef,
+    mut columns: Vec<ArrayRef>,
+) -> Result<RecordBatch> {
+    if columns.len() == output_schema.fields().len() + 1 {
+        Ok(RecordBatch::try_new(
+            schema_with_group_hash(&output_schema),
+            columns,
+        )?)
+    } else {
+        Ok(RecordBatch::try_new(
+            output_schema,
+            std::mem::take(&mut columns),
+        )?)
+    }
+}
+
 impl HashAggregateAccumulator {
     pub(super) fn new(
         aggregate_expr: Arc<AggregateFunctionExpr>,

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs

@@ -266,9 +266,11 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
     ) -> Result<()> {
         for group_values in &evaluated_batch.grouping_set_args {
             let starting_num_groups = self.buffer.group_values.len();
-            self.buffer
-                .group_values
-                .intern(group_values, &mut self.buffer.group_indices)?;
+            self.buffer.group_values.intern(
+                group_values,
+                &mut self.buffer.group_indices,
+                None,
+            )?;
             let total_num_groups = self.buffer.group_values.len();
             if total_num_groups > starting_num_groups {
                 self.buffer.group_ordering.new_groups(

datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs

@@ -22,7 +22,7 @@ use arrow::record_batch::RecordBatch;
 use datafusion_common::{Result, internal_err};
 use datafusion_expr::EmitTo;
 
-use crate::aggregates::AggregateExec;
+use crate::aggregates::{AggregateExec, strip_group_hash_column};
 
 use super::common::{
     AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, FinalMarker,
@@ -123,14 +123,19 @@ impl AggregateHashTable<FinalMarker> {
         &mut self,
         batch: &RecordBatch,
     ) -> Result<()> {
-        let evaluated_batch = self.evaluate_batch(batch)?;
+        let (batch, input_hashes) = strip_group_hash_column(batch)?;
+        let evaluated_batch = self.evaluate_batch(&batch)?;
         let state = self.state.building_mut();
 
         let timer = self.group_by_metrics.aggregation_time.timer();
         for group_values in &evaluated_batch.grouping_set_args {
-            state
-                .group_values
-                .intern(group_values, &mut state.batch_group_indices)?;
+            state.group_values.intern(
+                group_values,
+                &mut state.batch_group_indices,
+                input_hashes
+                    .map(|hashes| hashes.values())
+                    .map(|values| &**values),
+            )?;
             let group_indices = &state.batch_group_indices;
             let total_num_groups = state.group_values.len();
 

datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs

@@ -22,11 +22,13 @@ use arrow::record_batch::RecordBatch;
 use datafusion_common::{Result, internal_err};
 use datafusion_expr::EmitTo;
 
-use crate::aggregates::AggregateExec;
+use crate::aggregates::{
+    AggregateExec, create_group_hash_array, strip_group_hash_column,
+};
 
 use super::common::{
     AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState,
-    MaterializedAggregateOutput, PartialReduceMarker,
+    MaterializedAggregateOutput, PartialReduceMarker, try_new_internal_batch,
 };
 
 /// Methods specific to the aggregate hash table used in the partial-reduce stage.
@@ -89,13 +91,15 @@ impl AggregateHashTable<PartialReduceMarker> {
         let emit_to_all = EmitTo::All;
         let timer = self.group_by_metrics.emitting_time.timer();
         let mut output = state.group_values.emit(emit_to_all)?;
+        let group_hashes = create_group_hash_array(&output)?;
 
         for acc in state.accumulators.iter_mut() {
             output.extend(acc.state(emit_to_all)?);
         }
+        output.push(group_hashes);
         drop(timer);
 
-        let batch = RecordBatch::try_new(output_schema, output)?;
+        let batch = try_new_internal_batch(output_schema, output)?;
         debug_assert!(batch.num_rows() > 0);
         Ok(MaterializedAggregateOutput::new(batch))
     }
@@ -118,14 +122,19 @@ impl AggregateHashTable<PartialReduceMarker> {
         &mut self,
         batch: &RecordBatch,
     ) -> Result<()> {
-        let evaluated_batch = self.evaluate_batch(batch)?;
+        let (batch, input_hashes) = strip_group_hash_column(batch)?;
+        let evaluated_batch = self.evaluate_batch(&batch)?;
         let state = self.state.building_mut();
 
         let timer = self.group_by_metrics.aggregation_time.timer();
         for group_values in &evaluated_batch.grouping_set_args {
-            state
-                .group_values
-                .intern(group_values, &mut state.batch_group_indices)?;
+            state.group_values.intern(
+                group_values,
+                &mut state.batch_group_indices,
+                input_hashes
+                    .map(|hashes| hashes.values())
+                    .map(|values| &**values),
+            )?;
             let group_indices = &state.batch_group_indices;
             let total_num_groups = state.group_values.len();
 

datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs

@@ -27,12 +27,15 @@ use datafusion_expr::EmitTo;
 
 use crate::aggregates::group_values::new_group_values;
 use crate::aggregates::order::GroupOrdering;
-use crate::aggregates::{AggregateExec, group_id_array, max_duplicate_ordinal};
+use crate::aggregates::{
+    AggregateExec, create_group_hash_array, group_id_array, max_duplicate_ordinal,
+    schema_with_group_hash, strip_group_hash_column,
+};
 
 use super::common::{
     AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState,
     EvaluatedAccumulatorArgs, HashAggregateAccumulator, MaterializedAggregateOutput,
-    PartialMarker, PartialSkipMarker,
+    PartialMarker, PartialSkipMarker, try_new_internal_batch,
 };
 
 /// Implementation specific to partial aggregation, where the table stores
@@ -98,13 +101,15 @@ impl AggregateHashTable<PartialMarker> {
         let emit_to = EmitTo::All;
         let timer = self.group_by_metrics.emitting_time.timer();
         let mut output = state.group_values.emit(emit_to)?;
+        let group_hashes = create_group_hash_array(&output)?;
 
         for acc in state.accumulators.iter_mut() {
             output.extend(acc.state(emit_to)?);
         }
+        output.push(group_hashes);
         drop(timer);
 
-        let batch = RecordBatch::try_new(output_schema, output)?;
+        let batch = try_new_internal_batch(output_schema, output)?;
         debug_assert!(batch.num_rows() > 0);
         Ok(MaterializedAggregateOutput::new(batch))
     }
@@ -165,14 +170,19 @@ impl AggregateHashTable<PartialMarker> {
         &mut self,
         batch: &RecordBatch,
     ) -> Result<()> {
-        let evaluated_batch = self.evaluate_batch(batch)?;
+        let (batch, input_hashes) = strip_group_hash_column(batch)?;
+        let evaluated_batch = self.evaluate_batch(&batch)?;
         let state = self.state.building_mut();
 
         let _timer = self.group_by_metrics.aggregation_time.timer();
         for group_values in &evaluated_batch.grouping_set_args {
-            state
-                .group_values
-                .intern(group_values, &mut state.batch_group_indices)?;
+            state.group_values.intern(
+                group_values,
+                &mut state.batch_group_indices,
+                input_hashes
+                    .map(|hashes| hashes.values())
+                    .map(|values| &**values),
+            )?;
             let group_indices = &state.batch_group_indices;
             let total_num_groups = state.group_values.len();
 
@@ -242,7 +252,7 @@ impl AggregateHashTable<PartialMarker> {
 
             state
                 .group_values
-                .intern(&cols, &mut state.batch_group_indices)?;
+                .intern(&cols, &mut state.batch_group_indices, None)?;
             any_interned = true;
         }
 
@@ -280,6 +290,7 @@ impl AggregateHashTable<PartialSkipMarker> {
             .into_iter()
             .next()
             .unwrap_or_default();
+        let group_hashes = create_group_hash_array(&output)?;
 
         let state = self.state.building_mut();
         for (acc, values) in state
@@ -289,9 +300,10 @@ impl AggregateHashTable<PartialSkipMarker> {
         {
             output.extend(acc.convert_to_state(values)?);
         }
+        output.push(group_hashes);
 
         Ok(RecordBatch::try_new(
-            Arc::clone(&self.output_schema),
+            schema_with_group_hash(&self.output_schema),
             output,
         )?)
     }

datafusion/physical-plan/src/aggregates/group_values/mod.rs

@@ -97,7 +97,12 @@ pub trait GroupValues: Send {
     /// If a row has the same value as a previous row, the same group id is
     /// assigned. If a row has a new value, the next available group id is
     /// assigned.
-    fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()>;
+    fn intern(
+        &mut self,
+        cols: &[ArrayRef],
+        groups: &mut Vec<usize>,
+        hashes: Option<&[u64]>,
+    ) -> Result<()>;
 
     /// Returns the number of bytes of memory used by this [`GroupValues`]
     fn size(&self) -> usize;

datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs

@@ -348,6 +348,7 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> {
         &mut self,
         cols: &[ArrayRef],
         groups: &mut Vec<usize>,
+        hashes: Option<&[u64]>,
     ) -> Result<()> {
         let n_rows = cols[0].len();
 
@@ -357,8 +358,12 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> {
         // 1.1 Calculate the group keys for the group values
         let batch_hashes = &mut self.hashes_buffer;
         batch_hashes.clear();
-        batch_hashes.resize(n_rows, 0);
-        create_hashes(cols, &self.random_state, batch_hashes)?;
+        if let Some(hashes) = hashes {
+            batch_hashes.extend_from_slice(hashes);
+        } else {
+            batch_hashes.resize(n_rows, 0);
+            create_hashes(cols, &self.random_state, batch_hashes)?;
+        }
 
         for (row, &target_hash) in batch_hashes.iter().enumerate() {
             let entry = self
@@ -449,6 +454,7 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> {
         &mut self,
         cols: &[ArrayRef],
         groups: &mut Vec<usize>,
+        hashes: Option<&[u64]>,
     ) -> Result<()> {
         let n_rows = cols[0].len();
 
@@ -458,8 +464,12 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> {
 
         let mut batch_hashes = mem::take(&mut self.hashes_buffer);
         batch_hashes.clear();
-        batch_hashes.resize(n_rows, 0);
-        create_hashes(cols, &self.random_state, &mut batch_hashes)?;
+        if let Some(hashes) = hashes {
+            batch_hashes.extend_from_slice(hashes);
+        } else {
+            batch_hashes.resize(n_rows, 0);
+            create_hashes(cols, &self.random_state, &mut batch_hashes)?;
+        }
 
         // General steps for one round `vectorized equal_to & append`:
         //   1. Collect vectorized context by checking hash values of `cols` in `map`,
@@ -1078,14 +1088,19 @@ fn make_group_column(field: &Field) -> Result<Box<dyn GroupColumn>> {
 }
 
 impl<const STREAMING: bool> GroupValues for GroupValuesColumn<STREAMING> {
-    fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+    fn intern(
+        &mut self,
+        cols: &[ArrayRef],
+        groups: &mut Vec<usize>,
+        hashes: Option<&[u64]>,
+    ) -> Result<()> {
         // `try_new` and the reset points in `emit` / `clear_shrink` keep
         // `self.group_values` populated with one builder per schema field,
         // so no lazy initialization is needed here.
         if !STREAMING {
-            self.vectorized_intern(cols, groups)
+            self.vectorized_intern(cols, groups, hashes)
         } else {
-            self.scalarized_intern(cols, groups)
+            self.scalarized_intern(cols, groups, hashes)
         }
     }
 
@@ -1878,7 +1893,7 @@ mod tests {
 
         fn load_to_group_values(&self, group_values: &mut impl GroupValues) {
             for batch in self.test_batches.iter() {
-                group_values.intern(batch, &mut vec![]).unwrap();
+                group_values.intern(batch, &mut vec![], None).unwrap();
             }
         }
 

datafusion/physical-plan/src/aggregates/group_values/row.rs

@@ -116,7 +116,12 @@ impl GroupValuesRows {
 }
 
 impl GroupValues for GroupValuesRows {
-    fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+    fn intern(
+        &mut self,
+        cols: &[ArrayRef],
+        groups: &mut Vec<usize>,
+        hashes: Option<&[u64]>,
+    ) -> Result<()> {
         // Normalize -0.0 → +0.0 so RowConverter (IEEE 754 totalOrder) and
         // primitive hashing both group ±0 together. No-op for non-float
         // columns.
@@ -141,8 +146,12 @@ impl GroupValues for GroupValuesRows {
         // 1.1 Calculate the group keys for the group values
         let batch_hashes = &mut self.hashes_buffer;
         batch_hashes.clear();
-        batch_hashes.resize(n_rows, 0);
-        create_hashes(cols, &self.random_state, batch_hashes)?;
+        if let Some(hashes) = hashes {
+            batch_hashes.extend_from_slice(hashes);
+        } else {
+            batch_hashes.resize(n_rows, 0);
+            create_hashes(cols, &self.random_state, batch_hashes)?;
+        }
 
         for (row, &target_hash) in batch_hashes.iter().enumerate() {
             let entry = self.map.find_mut(target_hash, |(exist_hash, group_idx)| {

datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs

@@ -42,7 +42,12 @@ impl GroupValuesBoolean {
 }
 
 impl GroupValues for GroupValuesBoolean {
-    fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+    fn intern(
+        &mut self,
+        cols: &[ArrayRef],
+        groups: &mut Vec<usize>,
+        _hashes: Option<&[u64]>,
+    ) -> Result<()> {
         let array = cols[0].as_boolean();
         groups.clear();