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 82bdb2b046567ce0084bab889096f0e632c41b87 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 +- .../src/aggregates/aggregate_hash_table/common.rs | 18 ++++++++ .../aggregate_hash_table/common_ordered.rs | 8 ++-- .../aggregates/aggregate_hash_table/final_table.rs | 15 ++++--- .../aggregate_hash_table/partial_reduce_table.rs | 23 +++++++---- .../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 | 33 ++++++++++++--- datafusion/physical-plan/src/aggregates/mod.rs | 48 +++++++++++++++++++++- datafusion/physical-plan/src/recursive_query.rs | 7 +++- datafusion/physical-plan/src/repartition/mod.rs | 31 ++++++++++---- 18 files changed, 245 insertions(+), 65 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.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index f29f3e7ff8..256462e4c3 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/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>, 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..9f6f95f4d9 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,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(); 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..86de9d2dc7 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,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(); 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..b268051b35 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<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, )?) } 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..26c7e70583 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 d7c72253ec..e1a329f047 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<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() -> 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<usize> { + 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::<UInt64Array>() + .expect("group hash column must be UInt64Array"); + let projection: Vec<usize> = (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 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..b220820c8f 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<Vec<RecordBatch>> { 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::<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]
