This is an automated email from the ASF dual-hosted git repository. Rachelint pushed a commit to branch reuse-hash-through-hold-aggr in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 098dcbb5ee0ccaf9a22b889990b4d435e9ccc855 Author: kamille <[email protected]> AuthorDate: Fri Jul 10 19:31:36 2026 +0800 feat: reuse partial aggregate hashes across repartition --- .../benches/dictionary_group_values.rs | 6 ++-- datafusion/physical-plan/benches/multi_group_by.rs | 2 +- .../aggregate_hash_table/common_ordered.rs | 8 +++-- .../aggregates/aggregate_hash_table/final_table.rs | 21 ++++++++++--- .../aggregate_hash_table/partial_reduce_table.rs | 22 +++++++++++--- .../aggregate_hash_table/partial_table.rs | 30 +++++++++++++++---- .../src/aggregates/group_values/mod.rs | 7 ++++- .../aggregates/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 ++++-- .../group_values/single_group_by/bytes_view.rs | 3 +- .../group_values/single_group_by/primitive.rs | 17 +++++++---- .../src/aggregates/grouped_hash_stream.rs | 17 +++++++---- datafusion/physical-plan/src/aggregates/mod.rs | 35 ++++++++++++++++++++++ datafusion/physical-plan/src/recursive_query.rs | 7 +++-- datafusion/physical-plan/src/repartition/mod.rs | 28 ++++++++++++----- 17 files changed, 209 insertions(+), 56 deletions(-) diff --git a/datafusion/physical-plan/benches/dictionary_group_values.rs b/datafusion/physical-plan/benches/dictionary_group_values.rs index ded52aebd1..daf813a09d 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 92d0448775..76b94ed6ce 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_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index c83303c51d..340922eb7f 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<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( 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 568b866b10..ec1fd2bb4e 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 @@ -17,12 +17,13 @@ use std::sync::Arc; +use arrow::array::UInt64Array; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; use datafusion_expr::EmitTo; -use crate::aggregates::AggregateExec; +use crate::aggregates::{AggregateExec, group_hash_column_index}; use super::common::{ AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, FinalMarker, @@ -126,11 +127,23 @@ impl AggregateHashTable<FinalMarker> { let evaluated_batch = self.evaluate_batch(batch)?; let state = self.state.building_mut(); + let input_hashes = group_hash_column_index(batch.schema().as_ref()).map(|idx| { + batch + .column(idx) + .as_any() + .downcast_ref::<UInt64Array>() + .expect("group hash column must be UInt64Array") + }); + 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(); 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 4d94c55943..3d2e1ac589 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,7 +22,9 @@ 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, group_hash_column_index, +}; use super::common::{ AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, @@ -89,6 +91,7 @@ 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)?; + output.push(create_group_hash_array(&output)?); for acc in state.accumulators.iter_mut() { output.extend(acc.state(emit_to_all)?); @@ -120,12 +123,23 @@ impl AggregateHashTable<PartialReduceMarker> { ) -> Result<()> { let evaluated_batch = self.evaluate_batch(batch)?; let state = self.state.building_mut(); + let input_hashes = group_hash_column_index(batch.schema().as_ref()).map(|idx| { + batch + .column(idx) + .as_any() + .downcast_ref::<arrow::array::UInt64Array>() + .expect("group hash column must be UInt64Array") + }); 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(); 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 f11eef8c14..2fc843d5dd 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 @@ -19,7 +19,7 @@ use std::collections::HashMap; use std::marker::PhantomData; use std::sync::Arc; -use arrow::array::{ArrayRef, BooleanArray, new_null_array}; +use arrow::array::{ArrayRef, BooleanArray, UInt64Array, new_null_array}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; @@ -27,7 +27,10 @@ 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_hash_column_index, group_id_array, + max_duplicate_ordinal, +}; use super::common::{ AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, @@ -98,6 +101,7 @@ 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)?; + output.push(create_group_hash_array(&output)?); for acc in state.accumulators.iter_mut() { output.extend(acc.state(emit_to)?); @@ -168,11 +172,23 @@ impl AggregateHashTable<PartialMarker> { let evaluated_batch = self.evaluate_batch(batch)?; let state = self.state.building_mut(); + let input_hashes = group_hash_column_index(batch.schema().as_ref()).map(|idx| { + batch + .column(idx) + .as_any() + .downcast_ref::<UInt64Array>() + .expect("group hash column must be UInt64Array") + }); + 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 +258,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 +296,8 @@ impl AggregateHashTable<PartialSkipMarker> { .into_iter() .next() .unwrap_or_default(); + let group_hashes = create_group_hash_array(&output)?; + output.push(group_hashes); let state = self.state.building_mut(); for (acc, values) in state diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index ee253e5d7a..1e6ffad6e0 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<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; 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 f275d777c3..9ca2c3bb46 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<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(); } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 4976a098ec..dc8f0692c2 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<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)| { 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 e993c0c53d..2faee68287 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<usize>) -> Result<()> { + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec<usize>, + _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 b881a51b25..5152a31311 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<O: OffsetSizeTrait> GroupValuesBytes<O> { } impl<O: OffsetSizeTrait> GroupValues for GroupValuesBytes<O> { - 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<()> { assert_eq!(cols.len(), 1); // look up / add entries in the table @@ -108,7 +113,7 @@ impl<O: OffsetSizeTrait> GroupValues for GroupValuesBytes<O> { 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 7a56f7c52c..7a0ba85692 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<usize>, + _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 e254aebcfd..0ad03da204 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<T: ArrowPrimitiveType> GroupValues for GroupValuesPrimitive<T> where T::Native: HashValue, { - 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<()> { assert_eq!(cols.len(), 1); groups.clear(); - for v in cols[0].as_primitive::<T>() { + for (row, v) in cols[0].as_primitive::<T>().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 0d00e5c4d0..94a84309b1 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, }; 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,9 @@ impl GroupedHashAggregateStream { let timer = self.group_by_metrics.emitting_time.timer(); let mut output = self.group_values.emit(emit_to)?; + if self.mode.output_mode() == AggregateOutputMode::Partial || spilling { + output.push(create_group_hash_array(&output)?); + } if let EmitTo::First(n) = emit_to { self.group_ordering.remove_groups(n); } @@ -1104,7 +1110,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 +1389,7 @@ impl GroupedHashAggregateStream { "group_values expected to have single element" ); let mut output = group_values.swap_remove(0); + output.push(create_group_hash_array(&output)?); let iter = self .accumulators diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index d7c72253ec..44e6bf1a76 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -18,6 +18,7 @@ //! Aggregates functionalities use std::borrow::Cow; +use std::collections::HashMap as StdHashMap; use std::sync::Arc; use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; @@ -104,6 +105,39 @@ 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) const GROUP_HASH_GROUP_COUNT_META_KEY: &str = + "datafusion.group_hash.group_expr_count"; + +pub(crate) fn create_group_hash_array(group_values: &[ArrayRef]) -> Result<ArrayRef> { + 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(num_group_exprs: usize) -> Field { + let mut metadata = StdHashMap::new(); + metadata.insert( + GROUP_HASH_GROUP_COUNT_META_KEY.to_string(), + num_group_exprs.to_string(), + ); + Field::new(GROUP_HASH_COLUMN_NAME, DataType::UInt64, false).with_metadata(metadata) +} + +pub(crate) fn group_hash_column_index(schema: &Schema) -> Option<usize> { + schema + .fields() + .iter() + .position(|field| field.name() == GROUP_HASH_COLUMN_NAME) +} + /// Whether an aggregate stage consumes raw input data or intermediate /// accumulator state from a previous aggregation stage. /// @@ -2028,6 +2062,7 @@ fn create_schema( } AggregateOutputMode::Partial => { // in partial mode, the fields of the accumulator's state + fields.push(group_hash_field(group_by.num_group_exprs()).into()); for expr in aggr_expr { fields.extend(expr.state_fields()?.iter().cloned()); } diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 7289ac43e5..7328c3aad7 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 1617af3a68..7a4a5106d0 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; @@ -807,17 +808,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::<PrimitiveArray<arrow::datatypes::UInt64Type>>( + ) + .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()); --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
