Reuse hash values during Aggregate(partial) -> Repartition -> Aggregate(final)

datafusion/core/tests/execution/mod.rs

@@ -17,5 +17,6 @@
 
 mod coop;
 mod datasource_split;
+mod hash_reuse;
 mod logical_plan;
 mod register_arrow;

datafusion/core/tests/physical_optimizer/partition_statistics.rs

@@ -911,7 +911,6 @@ mod test {
                 ColumnStatistics::new_unknown(),
             ],
         };
-
         assert_eq!(*p0_statistics, expected_p0_statistics);
 
         let expected_p1_statistics = Statistics {
@@ -930,7 +929,6 @@ mod test {
                 ColumnStatistics::new_unknown(),
             ],
         };
-
         let p1_statistics = StatisticsContext::new().compute(
             aggregate_exec_partial.as_ref(),
             &StatisticsArgs::new().with_partition(Some(1)),
@@ -1009,7 +1007,6 @@ mod test {
                 ColumnStatistics::new_unknown(),
             ],
         };
-
         assert_eq!(
             empty_stat,
             *StatisticsContext::new().compute(

datafusion/physical-expr-common/src/binary_map.rs

@@ -299,6 +299,34 @@ where
         MP: FnMut(Option<&[u8]>) -> V,
         OP: FnMut(V),
     {
+        let mut hashes = std::mem::take(&mut self.hashes_buffer);
+        hashes.clear();
+        hashes.resize(values.len(), 0);
+        create_hashes([values], &self.random_state, &mut hashes)
+            // hash is supported for all types and create_hashes only
+            // returns errors for unsupported types
+            .unwrap();
+        self.insert_if_new_with_hashes(
+            values,
+            &hashes,
+            make_payload_fn,
+            observe_payload_fn,
+        );
+        self.hashes_buffer = hashes;
+    }
+
+    /// Like [`Self::insert_if_new`], but uses a hash supplied for each input value.
+    pub fn insert_if_new_with_hashes<MP, OP>(
+        &mut self,
+        values: &ArrayRef,
+        hashes: &[u64],
+        make_payload_fn: MP,
+        observe_payload_fn: OP,
+    ) where
+        MP: FnMut(Option<&[u8]>) -> V,
+        OP: FnMut(V),
+    {
+        assert_eq!(values.len(), hashes.len());
         // Sanity array type
         match self.output_type {
             OutputType::Binary => {
@@ -308,6 +336,7 @@ where
                 ));
                 self.insert_if_new_inner::<MP, OP, GenericBinaryType<O>>(
                     values,
+                    hashes,
                     make_payload_fn,
                     observe_payload_fn,
                 )
@@ -319,6 +348,7 @@ where
                 ));
                 self.insert_if_new_inner::<MP, OP, GenericStringType<O>>(
                     values,
+                    hashes,
                     make_payload_fn,
                     observe_payload_fn,
                 )
@@ -338,29 +368,20 @@ where
     fn insert_if_new_inner<MP, OP, B>(
         &mut self,
         values: &ArrayRef,
+        hashes: &[u64],
         mut make_payload_fn: MP,
         mut observe_payload_fn: OP,
     ) where
         MP: FnMut(Option<&[u8]>) -> V,
         OP: FnMut(V),
         B: ByteArrayType,
     {
-        // step 1: compute hashes
-        let batch_hashes = &mut self.hashes_buffer;
-        batch_hashes.clear();
-        batch_hashes.resize(values.len(), 0);
-        create_hashes([values], &self.random_state, batch_hashes)
-            // hash is supported for all types and create_hashes only
-            // returns errors for unsupported types
-            .unwrap();
-
-        // step 2: insert each value into the set, if not already present
         let values = values.as_bytes::<B>();
 
         // Ensure lengths are equivalent
-        assert_eq!(values.len(), batch_hashes.len());
+        assert_eq!(values.len(), hashes.len());
 
-        for (value, &hash) in values.iter().zip(batch_hashes.iter()) {
+        for (value, &hash) in values.iter().zip(hashes) {
             // handle null value
             let Some(value) = value else {
                 let payload = if let Some(&(payload, _offset)) = self.null.as_ref() {

datafusion/physical-expr-common/src/binary_view_map.rs

@@ -212,12 +212,41 @@ where
         MP: FnMut(Option<&[u8]>) -> V,
         OP: FnMut(V),
     {
+        let mut hashes = std::mem::take(&mut self.hashes_buffer);
+        hashes.clear();
+        hashes.resize(values.len(), 0);
+        create_hashes([values], &self.random_state, &mut hashes)
+            // hash is supported for all types and create_hashes only
+            // returns errors for unsupported types
+            .unwrap();
+        self.insert_if_new_with_hashes(
+            values,
+            &hashes,
+            make_payload_fn,
+            observe_payload_fn,
+        );
+        self.hashes_buffer = hashes;
+    }
+
+    /// Like [`Self::insert_if_new`], but uses a hash supplied for each input value.
+    pub fn insert_if_new_with_hashes<MP, OP>(
+        &mut self,
+        values: &ArrayRef,
+        hashes: &[u64],
+        make_payload_fn: MP,
+        observe_payload_fn: OP,
+    ) where
+        MP: FnMut(Option<&[u8]>) -> V,
+        OP: FnMut(V),
+    {
+        assert_eq!(values.len(), hashes.len());
         // Sanity check array type
         match self.output_type {
             OutputType::BinaryView => {
                 assert!(matches!(values.data_type(), DataType::BinaryView));
                 self.insert_if_new_inner::<MP, OP, BinaryViewType>(
                     values,
+                    hashes,
                     make_payload_fn,
                     observe_payload_fn,
                 )
@@ -226,6 +255,7 @@ where
                 assert!(matches!(values.data_type(), DataType::Utf8View));
                 self.insert_if_new_inner::<MP, OP, StringViewType>(
                     values,
+                    hashes,
                     make_payload_fn,
                     observe_payload_fn,
                 )
@@ -245,34 +275,25 @@ where
     fn insert_if_new_inner<MP, OP, B>(
         &mut self,
         values: &ArrayRef,
+        hashes: &[u64],
         mut make_payload_fn: MP,
         mut observe_payload_fn: OP,
     ) where
         MP: FnMut(Option<&[u8]>) -> V,
         OP: FnMut(V),
         B: ByteViewType,
     {
-        // step 1: compute hashes
-        let batch_hashes = &mut self.hashes_buffer;
-        batch_hashes.clear();
-        batch_hashes.resize(values.len(), 0);
-        create_hashes([values], &self.random_state, batch_hashes)
-            // hash is supported for all types and create_hashes only
-            // returns errors for unsupported types
-            .unwrap();
-
-        // step 2: insert each value into the set, if not already present
         let values = values.as_byte_view::<B>();
 
         // Get raw views buffer for direct comparison
         let input_views = values.views();
 
         // Ensure lengths are equivalent
-        assert_eq!(values.len(), self.hashes_buffer.len());
+        assert_eq!(values.len(), hashes.len());
 
         for i in 0..values.len() {
             let view_u128 = input_views[i];
-            let hash = self.hashes_buffer[i];
+            let hash = hashes[i];
 
             // handle null value via validity bitmap check
             if values.is_null(i) {

datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs

@@ -52,7 +52,7 @@ use datafusion_physical_expr::{
 };
 use datafusion_physical_plan::ExecutionPlanProperties;
 use datafusion_physical_plan::aggregates::{
-    AggregateExec, AggregateMode, PhysicalGroupBy,
+    AggregateExec, AggregateMode, AggregateOutputMode, PhysicalGroupBy,
 };
 use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion_physical_plan::execution_plan::EmissionType;
@@ -728,6 +728,63 @@ fn partial_aggregate_output_satisfies_final_partitioning(
         .is_satisfied()
 }
 
+/// Returns whether avoiding repeated hashing is expected to offset the cost of
+/// propagating an additional `UInt64` column.
+fn aggregate_hash_output_is_worthwhile(aggregate: &AggregateExec) -> Result<bool> {
+    let input_schema = aggregate.input().schema();
+    for expr in aggregate.group_expr().input_exprs() {
+        if !expr.data_type(&input_schema)?.is_primitive() {
+            return Ok(true);
+        }
+    }
+    Ok(false)
+}
+
+/// Configures the partial aggregate below distribution-preserving operators to
+/// emit reusable group hashes.
+///
+/// This runs while enforcing the final aggregate's distribution, when the
+/// complete partial -> exchange -> final shape is known. Keeping the decision
+/// here avoids runtime operators independently inferring whether a private hash
+/// column should be part of their output.
+fn configure_partial_aggregate_hash_output(
+    mut ctx: DistributionContext,
+) -> Result<DistributionContext> {
+    if let Some(aggregate) = ctx.plan.downcast_ref::<AggregateExec>() {
+        if aggregate.mode() == &AggregateMode::Partial {
+            let emit_hashes = aggregate_hash_output_is_worthwhile(aggregate)?;
+            ctx.plan = Arc::new(aggregate.clone().with_hash_reuse(emit_hashes)?);
+        }
+        return Ok(ctx);
+    }
+
+    if ctx.children.len() != 1 {
+        return Ok(ctx);
+    }
+
+    if let Some(repartition) = ctx.plan.downcast_ref::<RepartitionExec>() {
+        let child = configure_partial_aggregate_hash_output(ctx.children.swap_remove(0))?;
+        ctx.plan = Arc::new(
+            repartition
+                .clone()
+                .with_new_input(Arc::clone(&child.plan))?
+                .with_hash_reuse(true),
+        );
+        ctx.children.push(child);
+        return Ok(ctx);
+    }
+
+    if ctx.plan.is::<CoalescePartitionsExec>() || ctx.plan.is::<SortPreservingMergeExec>()
+    {
+        let child = configure_partial_aggregate_hash_output(ctx.children.swap_remove(0))?;
+        ctx.plan =
+            Arc::clone(&ctx.plan).with_new_children(vec![Arc::clone(&child.plan)])?;
+        ctx.children.push(child);
+    }
+
+    Ok(ctx)
+}
+
 /// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator
 /// on top of the given plan node to satisfy a single partition requirement
 /// while preserving ordering constraints.
@@ -1407,7 +1464,7 @@ pub fn ensure_distribution(
         target_partitions,
     )?;
 
-    let children = children
+    let mut children = children
         .into_iter()
         .map(
             |DistributionChildState {
@@ -1483,6 +1540,16 @@ pub fn ensure_distribution(
         )
         .collect::<Result<Vec<_>>>()?;
 
+    if plan
+        .downcast_ref::<AggregateExec>()
+        .is_some_and(|aggregate| {
+            matches!(aggregate.mode().output_mode(), AggregateOutputMode::Final)
+        })
+    {
+        let child = children.swap_remove(0);
+        children.push(configure_partial_aggregate_hash_output(child)?);
+    }
+
     let children_plans = children
         .iter()
         .map(|c| Arc::clone(&c.plan))

datafusion/physical-plan/benches/dictionary_group_values.rs

@@ -27,6 +27,7 @@ use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef};
 use criterion::{
     BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main,
 };
+use datafusion_common::hash_utils::{RandomState, create_hashes};
 use datafusion_expr::EmitTo;
 use datafusion_physical_plan::aggregates::group_values::new_group_values;
 use datafusion_physical_plan::aggregates::order::GroupOrdering;
@@ -41,6 +42,18 @@ const CARDS_RELATIVE: [usize; 4] = [20, 75, 300, 1000];
 const N_BATCHES: usize = 4;
 // Fixed for reproducibility.
 const SEED: u64 = 0xD1C7;
+const AGGREGATION_HASH_SEED: u64 = 15395726432021054657;
+
+fn create_aggregation_hashes(array: &ArrayRef, hashes: &mut Vec<u64>) {
+    hashes.clear();
+    hashes.resize(array.len(), 0);
+    create_hashes(
+        std::slice::from_ref(array),
+        &RandomState::with_seed(AGGREGATION_HASH_SEED),
+        hashes,
+    )
+    .unwrap();
+}
 
 fn dict_schema() -> SchemaRef {
     let dict_ty =
@@ -109,10 +122,13 @@ fn bench_intern_emit(c: &mut Criterion) {
                                 new_group_values(schema.clone(), &GroupOrdering::None)
                                     .unwrap(),
                                 Vec::<usize>::with_capacity(size),
+                                Vec::<u64>::with_capacity(size),
                             )
                         },
-                        |(gv, groups)| {
-                            gv.intern(std::slice::from_ref(&array), groups).unwrap();
+                        |(gv, groups, hashes)| {
+                            create_aggregation_hashes(&array, hashes);
+                            gv.intern(std::slice::from_ref(&array), groups, hashes)
+                                .unwrap();
                             black_box(&*groups);
                             black_box(gv.emit(EmitTo::All).unwrap());
                         },
@@ -154,11 +170,14 @@ fn bench_repeated_intern_emit(c: &mut Criterion) {
                                 new_group_values(schema.clone(), &GroupOrdering::None)
                                     .unwrap(),
                                 Vec::<usize>::with_capacity(size),
+                                Vec::<u64>::with_capacity(size),
                             )
                         },
-                        |(gv, groups)| {
+                        |(gv, groups, hashes)| {
                             for arr in &batches {
-                                gv.intern(std::slice::from_ref(arr), groups).unwrap();
+                                create_aggregation_hashes(arr, hashes);
+                                gv.intern(std::slice::from_ref(arr), groups, hashes)
+                                    .unwrap();
                                 black_box(&*groups);
                             }
                             black_box(gv.emit(EmitTo::All).unwrap());