From 020092516259b5c6d3654fa18157496a9886750d Mon Sep 17 00:00:00 2001 From: kamille Date: Fri, 10 Jul 2026 19:31:36 +0800 Subject: [PATCH] feat: reuse partial aggregate hashes across repartition --- .../benches/dictionary_group_values.rs | 6 +- .../physical-plan/benches/multi_group_by.rs | 2 +- .../aggregates/aggregate_hash_table/common.rs | 71 ++++++++++++++++++- .../aggregate_hash_table/common_ordered.rs | 8 ++- .../aggregate_hash_table/final_table.rs | 16 +++-- .../partial_reduce_table.rs | 22 ++++-- .../aggregate_hash_table/partial_table.rs | 45 ++++++++---- .../src/aggregates/group_values/mod.rs | 7 +- .../group_values/multi_group_by/mod.rs | 31 +++++--- .../src/aggregates/group_values/row.rs | 15 +++- .../group_values/single_group_by/boolean.rs | 7 +- .../group_values/single_group_by/bytes.rs | 9 ++- .../single_group_by/bytes_view.rs | 3 +- .../group_values/single_group_by/primitive.rs | 17 +++-- .../src/aggregates/grouped_hash_stream.rs | 33 +++++++-- .../physical-plan/src/aggregates/mod.rs | 48 ++++++++++++- .../physical-plan/src/recursive_query.rs | 7 +- .../physical-plan/src/repartition/mod.rs | 31 +++++--- 18 files changed, 308 insertions(+), 70 deletions(-) diff --git a/datafusion/physical-plan/benches/dictionary_group_values.rs b/datafusion/physical-plan/benches/dictionary_group_values.rs index ded52aebd1100..daf813a09d274 100644 --- a/datafusion/physical-plan/benches/dictionary_group_values.rs +++ b/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()); diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 92d0448775599..76b94ed6ce31b 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/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); } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index f29f3e7ff8af1..7a7fbad108094 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -18,9 +18,10 @@ use std::marker::PhantomData; use std::sync::Arc; -use arrow::array::{ArrayRef, AsArray, new_null_array}; +use arrow::array::{ArrayRef, AsArray, UInt64Array, new_null_array}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; +use datafusion_common::hash_utils::create_hashes; use datafusion_common::{Result, internal_err}; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::{EmitTo, GroupsAccumulator}; @@ -31,7 +32,8 @@ use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_val use crate::aggregates::grouped_hash_stream::create_group_accumulator; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ - AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, + AGGREGATION_HASH_SEED, AggregateExec, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, schema_with_group_hash, }; /// Marker for raw rows -> partial state aggregation. @@ -136,6 +138,8 @@ impl AggregateHashTable { group_by: Arc::clone(&agg.group_by), group_values, batch_group_indices: Default::default(), + group_hashes: Default::default(), + batch_hashes: Default::default(), accumulators, }), _mode: PhantomData, @@ -183,6 +187,8 @@ impl AggregateHashTable { acc + state.group_values.size() + state.batch_group_indices.allocated_size() + + state.group_hashes.allocated_size() + + state.batch_hashes.allocated_size() } AggregateHashTableState::OutputtingMaterialized(output) => { output.memory_size() @@ -212,6 +218,7 @@ impl AggregateHashTable { }; state.batch_group_indices = Vec::new(); + state.batch_hashes = Vec::new(); self.state = AggregateHashTableState::Outputting(state); } } @@ -287,6 +294,12 @@ pub(super) struct AggregateHashTableBuffer { /// accumulator to update that group's aggregate state. pub(super) batch_group_indices: Vec, + /// Hash value for each stored group, indexed by group index. + pub(super) group_hashes: Vec, + + /// Scratch hash vector for the current input batch. + pub(super) batch_hashes: Vec, + /// One item per aggregate expression. /// /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input @@ -343,6 +356,60 @@ impl MaterializedAggregateOutput { } } +pub(super) fn try_new_internal_batch( + output_schema: SchemaRef, + mut columns: Vec, +) -> Result { + 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 AggregateHashTableBuffer { + pub(super) fn ensure_batch_hashes( + &mut self, + group_values: &[ArrayRef], + input_hashes: Option<&UInt64Array>, + ) -> Result<&[u64]> { + self.batch_hashes.clear(); + if let Some(input_hashes) = input_hashes { + self.batch_hashes.extend_from_slice(input_hashes.values()); + } else { + let num_rows = group_values.first().map(|array| array.len()).unwrap_or(0); + self.batch_hashes.resize(num_rows, 0); + create_hashes(group_values, &AGGREGATION_HASH_SEED, &mut self.batch_hashes)?; + } + Ok(&self.batch_hashes) + } + + pub(super) fn record_new_group_hashes(&mut self, starting_num_groups: usize) { + for (row, group_index) in self.batch_group_indices.iter().enumerate() { + if *group_index >= starting_num_groups { + if self.group_hashes.len() <= *group_index { + self.group_hashes.resize(*group_index + 1, 0); + } + self.group_hashes[*group_index] = self.batch_hashes[row]; + } + } + } + + pub(super) fn emit_hashes(&mut self, emit_to: EmitTo) -> ArrayRef { + let hashes = match emit_to { + EmitTo::All => std::mem::take(&mut self.group_hashes), + EmitTo::First(n) => self.group_hashes.drain(..n).collect(), + }; + Arc::new(UInt64Array::from(hashes)) + } +} + impl HashAggregateAccumulator { pub(super) fn new( aggregate_expr: Arc, diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index c83303c51d6e8..340922eb7ff55 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -266,9 +266,11 @@ impl OrderedAggregateTable { ) -> 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( diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs index 568b866b10517..3e47d5d1f05ed 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs +++ b/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,20 @@ impl AggregateHashTable { &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)?; + let starting_num_groups = state.group_values.len(); + state.ensure_batch_hashes(group_values, input_hashes)?; + state.group_values.intern( + group_values, + &mut state.batch_group_indices, + Some(&state.batch_hashes), + )?; + state.record_new_group_hashes(starting_num_groups); let group_indices = &state.batch_group_indices; let total_num_groups = state.group_values.len(); diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs index 4d94c559436fb..df2176036f166 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs @@ -22,11 +22,11 @@ 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, - MaterializedAggregateOutput, PartialReduceMarker, + MaterializedAggregateOutput, PartialReduceMarker, try_new_internal_batch, }; /// Methods specific to the aggregate hash table used in the partial-reduce stage. @@ -89,13 +89,15 @@ impl AggregateHashTable { 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 = state.emit_hashes(emit_to_all); 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 +120,20 @@ impl AggregateHashTable { &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)?; + let starting_num_groups = state.group_values.len(); + state.ensure_batch_hashes(group_values, input_hashes)?; + state.group_values.intern( + group_values, + &mut state.batch_group_indices, + Some(&state.batch_hashes), + )?; + state.record_new_group_hashes(starting_num_groups); let group_indices = &state.batch_group_indices; let total_num_groups = state.group_values.len(); diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index f11eef8c14277..72f0fc863b460 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/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 { 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 = state.emit_hashes(emit_to); 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)) } @@ -155,6 +160,8 @@ impl AggregateHashTable { group_by: Arc::clone(&state.group_by), group_values, batch_group_indices: Default::default(), + group_hashes: Default::default(), + batch_hashes: Default::default(), accumulators, }), _mode: PhantomData, @@ -165,14 +172,20 @@ impl AggregateHashTable { &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)?; + let starting_num_groups = state.group_values.len(); + state.ensure_batch_hashes(group_values, input_hashes)?; + state.group_values.intern( + group_values, + &mut state.batch_group_indices, + Some(&state.batch_hashes), + )?; + state.record_new_group_hashes(starting_num_groups); let group_indices = &state.batch_group_indices; let total_num_groups = state.group_values.len(); @@ -215,12 +228,13 @@ impl AggregateHashTable { } let max_ordinal = max_duplicate_ordinal(state.group_by.groups()); + let groups = state.group_by.groups().to_vec(); let mut ordinals: HashMap<&[bool], usize> = HashMap::new(); let group_schema = state.group_by.group_schema(&self.input_schema)?; let n_expr = state.group_by.expr().len(); let mut any_interned = false; - for group in state.group_by.groups() { + for group in &groups { let ordinal = { let entry = ordinals.entry(group.as_slice()).or_insert(0); let ordinal = *entry; @@ -240,9 +254,14 @@ impl AggregateHashTable { .collect(); cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); - state - .group_values - .intern(&cols, &mut state.batch_group_indices)?; + let starting_num_groups = state.group_values.len(); + state.ensure_batch_hashes(&cols, None)?; + state.group_values.intern( + &cols, + &mut state.batch_group_indices, + Some(&state.batch_hashes), + )?; + state.record_new_group_hashes(starting_num_groups); any_interned = true; } @@ -280,6 +299,7 @@ impl AggregateHashTable { .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 +309,10 @@ impl AggregateHashTable { { 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, )?) } diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index ee253e5d7afdd..1e6ffad6e024d 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/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) -> Result<()>; + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec, + hashes: Option<&[u64]>, + ) -> Result<()>; /// Returns the number of bytes of memory used by this [`GroupValues`] fn size(&self) -> usize; diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index f275d777c3279..9ca2c3bb46fb6 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -348,6 +348,7 @@ impl GroupValuesColumn { &mut self, cols: &[ArrayRef], groups: &mut Vec, + hashes: Option<&[u64]>, ) -> Result<()> { let n_rows = cols[0].len(); @@ -357,8 +358,12 @@ impl GroupValuesColumn { // 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 GroupValuesColumn { &mut self, cols: &[ArrayRef], groups: &mut Vec, + hashes: Option<&[u64]>, ) -> Result<()> { let n_rows = cols[0].len(); @@ -458,8 +464,12 @@ impl GroupValuesColumn { 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> { } impl GroupValues for GroupValuesColumn { - fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec, + 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(); } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 4976a098ecee5..dc8f0692c2e95 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/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) -> Result<()> { + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec, + 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)| { diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs index e993c0c53d199..2faee682874af 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/boolean.rs +++ b/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) -> Result<()> { + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec, + _hashes: Option<&[u64]>, + ) -> Result<()> { let array = cols[0].as_boolean(); groups.clear(); diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs index b881a51b25474..5152a3131196a 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs @@ -45,7 +45,12 @@ impl GroupValuesBytes { } impl GroupValues for GroupValuesBytes { - fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec, + _hashes: Option<&[u64]>, + ) -> Result<()> { assert_eq!(cols.len(), 1); // look up / add entries in the table @@ -108,7 +113,7 @@ impl GroupValues for GroupValuesBytes { self.num_groups = 0; let mut group_indexes = vec![]; - self.intern(&[remaining_group_values], &mut group_indexes)?; + self.intern(&[remaining_group_values], &mut group_indexes, None)?; // Verify that the group indexes were assigned in the correct order assert_eq!(0, group_indexes[0]); diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs index 7a56f7c52c11a..7a0ba8569215e 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs @@ -47,6 +47,7 @@ impl GroupValues for GroupValuesBytesView { &mut self, cols: &[ArrayRef], groups: &mut Vec, + _hashes: Option<&[u64]>, ) -> datafusion_common::Result<()> { assert_eq!(cols.len(), 1); @@ -110,7 +111,7 @@ impl GroupValues for GroupValuesBytesView { self.num_groups = 0; let mut group_indexes = vec![]; - self.intern(&[remaining_group_values], &mut group_indexes)?; + self.intern(&[remaining_group_values], &mut group_indexes, None)?; // Verify that the group indexes were assigned in the correct order assert_eq!(0, group_indexes[0]); diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index e254aebcfd7ce..0ad03da204344 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -135,11 +135,16 @@ impl GroupValues for GroupValuesPrimitive where T::Native: HashValue, { - fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec, + hashes: Option<&[u64]>, + ) -> Result<()> { assert_eq!(cols.len(), 1); groups.clear(); - for v in cols[0].as_primitive::() { + for (row, v) in cols[0].as_primitive::().iter().enumerate() { let group_id = match v { None => *self.null_group.get_or_insert_with(|| { let group_id = self.values.len(); @@ -151,8 +156,10 @@ where // so the bit-equal `is_eq` matches and the stored value is // the canonical representative. let key = key.canonicalize(); - let state = &self.random_state; - let hash = key.hash(state); + let hash = hashes.map(|hashes| hashes[row]).unwrap_or_else(|| { + let state = &self.random_state; + key.hash(state) + }); let insert = self.map.entry( hash, |&(g, h)| unsafe { @@ -273,7 +280,7 @@ mod tests { // Intern 20 distinct values; `new()` pre-allocates capacity 128 for `values`. let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..20i32)); let mut groups = vec![]; - gv.intern(&[arr], &mut groups)?; + gv.intern(&[arr], &mut groups, None)?; let capacity_before = gv.values.capacity(); // 128 // n=4, n*2=8 <= len=20 -> drain branch diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 0d00e5c4d0d86..26c7e70583a34 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -28,8 +28,8 @@ use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_val use crate::aggregates::order::GroupOrderingFull; use crate::aggregates::{ AggregateInputMode, AggregateMode, AggregateOutputMode, PhysicalGroupBy, - create_schema, evaluate_group_by, evaluate_many, evaluate_optional, group_id_array, - max_duplicate_ordinal, + create_group_hash_array, create_schema, evaluate_group_by, evaluate_many, + evaluate_optional, group_id_array, max_duplicate_ordinal, schema_with_group_hash, }; use crate::metrics::{BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput}; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; @@ -883,8 +883,11 @@ impl GroupedHashAggregateStream { // calculate the group indices for each input row let starting_num_groups = self.group_values.len(); - self.group_values - .intern(group_values, &mut self.current_group_indices)?; + self.group_values.intern( + group_values, + &mut self.current_group_indices, + None, + )?; let group_indices = &self.current_group_indices; // Update ordering information if necessary @@ -1034,6 +1037,12 @@ impl GroupedHashAggregateStream { let timer = self.group_by_metrics.emitting_time.timer(); let mut output = self.group_values.emit(emit_to)?; + let group_hashes = + if self.mode.output_mode() == AggregateOutputMode::Partial || spilling { + Some(create_group_hash_array(&output)?) + } else { + None + }; if let EmitTo::First(n) = emit_to { self.group_ordering.remove_groups(n); } @@ -1048,11 +1057,20 @@ impl GroupedHashAggregateStream { output.extend(acc.state(emit_to)?) } } + if let Some(group_hashes) = group_hashes { + output.push(group_hashes); + } drop(timer); // emit reduces the memory usage. Ignore Err from update_memory_reservation. Even if it is // over the target memory size after emission, we can emit again rather than returning Err. let _ = self.update_memory_reservation(); + let schema = + if self.mode.output_mode() == AggregateOutputMode::Partial || spilling { + schema_with_group_hash(&schema) + } else { + schema + }; let batch = RecordBatch::try_new(schema, output)?; debug_assert!(batch.num_rows() > 0); @@ -1104,7 +1122,7 @@ impl GroupedHashAggregateStream { let starting_groups = self.group_values.len(); self.group_values - .intern(&cols, &mut self.current_group_indices)?; + .intern(&cols, &mut self.current_group_indices, None)?; let total_groups = self.group_values.len(); if total_groups > starting_groups { self.group_ordering.new_groups( @@ -1383,6 +1401,7 @@ impl GroupedHashAggregateStream { "group_values expected to have single element" ); let mut output = group_values.swap_remove(0); + let group_hashes = create_group_hash_array(&output)?; let iter = self .accumulators @@ -1394,8 +1413,10 @@ impl GroupedHashAggregateStream { let opt_filter = opt_filter.as_ref().map(|filter| filter.as_boolean()); output.extend(acc.convert_to_state(values, opt_filter)?); } + output.push(group_hashes); - let states_batch = RecordBatch::try_new(self.schema(), output)?; + let states_batch = + RecordBatch::try_new(schema_with_group_hash(&self.schema()), output)?; Ok(states_batch) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index d7c72253ecc0c..e1a329f0477a3 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -53,7 +53,7 @@ use arrow_schema::FieldRef; use datafusion_common::stats::Precision; use datafusion_common::{ Constraint, Constraints, Result, ScalarValue, assert_eq_or_internal_err, - internal_err, not_impl_err, + internal_err, not_impl_err, project_schema, }; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryLimit; @@ -104,6 +104,52 @@ const AGGREGATION_HASH_SEED: datafusion_common::hash_utils::RandomState = // This seed is chosen to be a large 64-bit number datafusion_common::hash_utils::RandomState::with_seed(15395726432021054657); +/// Internal partial-aggregation column that carries precomputed group hashes +/// across repartition and into downstream aggregate stages. +pub(crate) const GROUP_HASH_COLUMN_NAME: &str = "__datafusion_group_hash"; +pub(crate) fn create_group_hash_array(group_values: &[ArrayRef]) -> Result { + let num_rows = group_values.first().map(|array| array.len()).unwrap_or(0); + let mut hashes = vec![0; num_rows]; + datafusion_common::hash_utils::create_hashes( + group_values, + &AGGREGATION_HASH_SEED, + &mut hashes, + )?; + Ok(Arc::new(UInt64Array::from(hashes))) +} + +pub(crate) fn group_hash_field() -> Field { + Field::new(GROUP_HASH_COLUMN_NAME, DataType::UInt64, false) +} + +pub(crate) fn schema_with_group_hash(schema: &SchemaRef) -> SchemaRef { + let mut fields = schema.fields().to_vec(); + fields.push(group_hash_field().into()); + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) +} + +pub(crate) fn group_hash_column_index(schema: &Schema) -> Option { + let last_index = schema.fields().len().checked_sub(1)?; + (schema.field(last_index).name() == GROUP_HASH_COLUMN_NAME).then_some(last_index) +} + +pub(crate) fn strip_group_hash_column( + batch: &RecordBatch, +) -> Result<(RecordBatch, Option<&UInt64Array>)> { + let Some(hash_index) = group_hash_column_index(batch.schema().as_ref()) else { + return Ok((batch.clone(), None)); + }; + let hashes = batch + .column(hash_index) + .as_any() + .downcast_ref::() + .expect("group hash column must be UInt64Array"); + let projection: Vec = (0..hash_index).collect(); + let schema = project_schema(&batch.schema(), Some(&projection))?; + let columns = batch.columns()[..hash_index].to_vec(); + Ok((RecordBatch::try_new(schema, columns)?, Some(hashes))) +} + /// Whether an aggregate stage consumes raw input data or intermediate /// accumulator state from a previous aggregation stage. /// diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 7289ac43e510c..7328c3aad7736 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -466,8 +466,11 @@ impl DistinctDeduplicator { "failed to reserve {additional} recursive query group ids: {e}" ) })?; - self.group_values - .intern(batch.columns(), &mut self.intern_output_buffer)?; + self.group_values.intern( + batch.columns(), + &mut self.intern_output_buffer, + None, + )?; let mask = new_groups_mask(&self.intern_output_buffer, size_before); self.intern_output_buffer.clear(); // We update the reservation to reflect the new size of the hash table. diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 1617af3a68baa..b220820c8fc9c 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -31,6 +31,7 @@ use super::metrics::{self, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; use super::{ DisplayAs, ExecutionPlanProperties, RecordBatchStream, SendableRecordBatchStream, }; +use crate::aggregates::group_hash_column_index; use crate::coalesce::LimitedBatchCoalescer; use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; use crate::hash_utils::create_hashes; @@ -245,6 +246,9 @@ impl SharedCoalescer { fn push_and_drain(&self, batch: RecordBatch) -> Result> { let mut acc = Vec::new(); let mut c = self.inner.lock(); + if c.schema() != batch.schema() { + return Ok(vec![batch]); + } c.push_batch(batch)?; while let Some(b) = c.next_completed_batch() { acc.push(b); @@ -807,17 +811,28 @@ impl BatchPartitioner { // Tracking time required for distributing indexes across output partitions let timer = self.timer.timer(); - let arrays = - evaluate_expressions_to_arrays(exprs.as_slice(), &batch)?; - hash_buffer.clear(); hash_buffer.resize(batch.num_rows(), 0); - create_hashes( - &arrays, - REPARTITION_RANDOM_STATE.random_state(), - hash_buffer, - )?; + if let Some(hash_index) = + group_hash_column_index(batch.schema().as_ref()) + { + let hashes = batch + .column(hash_index) + .as_any() + .downcast_ref::>( + ) + .expect("group hash column must be UInt64Array"); + hash_buffer.copy_from_slice(hashes.values()); + } else { + let arrays = + evaluate_expressions_to_arrays(exprs.as_slice(), &batch)?; + create_hashes( + &arrays, + REPARTITION_RANDOM_STATE.random_state(), + hash_buffer, + )?; + } indices.iter_mut().for_each(|v| v.clear());