From 0ecb6ec49064bfab157ce0855b996e1a57b8f59e 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 --- .../physical-expr-common/src/binary_map.rs | 43 +++++--- .../src/binary_view_map.rs | 43 +++++--- .../benches/dictionary_group_values.rs | 23 ++++- .../physical-plan/benches/multi_group_by.rs | 8 +- .../aggregates/aggregate_hash_table/common.rs | 75 +++++++++++++- .../aggregate_hash_table/common_ordered.rs | 25 ++++- .../aggregate_hash_table/final_table.rs | 16 +-- .../partial_reduce_table.rs | 28 ++++-- .../aggregate_hash_table/partial_table.rs | 46 ++++++--- .../src/aggregates/group_values/mod.rs | 7 +- .../group_values/multi_group_by/mod.rs | 54 ++++------ .../src/aggregates/group_values/row.rs | 32 ++---- .../group_values/single_group_by/boolean.rs | 7 +- .../group_values/single_group_by/bytes.rs | 17 +++- .../single_group_by/bytes_view.rs | 11 ++- .../group_values/single_group_by/primitive.rs | 19 ++-- .../src/aggregates/grouped_hash_stream.rs | 47 +++++++-- .../physical-plan/src/aggregates/mod.rs | 98 ++++++++++++++++++- .../physical-plan/src/recursive_query.rs | 9 +- .../physical-plan/src/repartition/mod.rs | 31 ++++-- 20 files changed, 483 insertions(+), 156 deletions(-) diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index ad184d6500d56..85b9b7aa7da3d 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -298,6 +298,32 @@ where ) where MP: FnMut(Option<&[u8]>) -> V, OP: FnMut(V), + { + let mut batch_hashes = std::mem::take(&mut self.hashes_buffer); + batch_hashes.clear(); + batch_hashes.resize(values.len(), 0); + create_hashes([values], &self.random_state, &mut batch_hashes) + // hash is supported for all types and create_hashes only + // returns errors for unsupported types + .unwrap(); + self.insert_if_new_with_hashes( + values, + &batch_hashes, + make_payload_fn, + observe_payload_fn, + ); + self.hashes_buffer = batch_hashes; + } + + pub fn insert_if_new_with_hashes( + &mut self, + values: &ArrayRef, + hashes: &[u64], + make_payload_fn: MP, + observe_payload_fn: OP, + ) where + MP: FnMut(Option<&[u8]>) -> V, + OP: FnMut(V), { // Sanity array type match self.output_type { @@ -308,6 +334,7 @@ where )); self.insert_if_new_inner::>( values, + hashes, make_payload_fn, observe_payload_fn, ) @@ -319,6 +346,7 @@ where )); self.insert_if_new_inner::>( values, + hashes, make_payload_fn, observe_payload_fn, ) @@ -338,6 +366,7 @@ where fn insert_if_new_inner( &mut self, values: &ArrayRef, + hashes: &[u64], mut make_payload_fn: MP, mut observe_payload_fn: OP, ) where @@ -345,22 +374,12 @@ where 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::(); // 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.iter()) { // handle null value let Some(value) = value else { let payload = if let Some(&(payload, _offset)) = self.null.as_ref() { diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 9d4b556393a24..2a3ffe824ce7a 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -211,6 +211,32 @@ where ) where MP: FnMut(Option<&[u8]>) -> V, OP: FnMut(V), + { + let mut batch_hashes = std::mem::take(&mut self.hashes_buffer); + batch_hashes.clear(); + batch_hashes.resize(values.len(), 0); + create_hashes([values], &self.random_state, &mut batch_hashes) + // hash is supported for all types and create_hashes only + // returns errors for unsupported types + .unwrap(); + self.insert_if_new_with_hashes( + values, + &batch_hashes, + make_payload_fn, + observe_payload_fn, + ); + self.hashes_buffer = batch_hashes; + } + + pub fn insert_if_new_with_hashes( + &mut self, + values: &ArrayRef, + hashes: &[u64], + make_payload_fn: MP, + observe_payload_fn: OP, + ) where + MP: FnMut(Option<&[u8]>) -> V, + OP: FnMut(V), { // Sanity check array type match self.output_type { @@ -218,6 +244,7 @@ where assert!(matches!(values.data_type(), DataType::BinaryView)); self.insert_if_new_inner::( values, + hashes, make_payload_fn, observe_payload_fn, ) @@ -226,6 +253,7 @@ where assert!(matches!(values.data_type(), DataType::Utf8View)); self.insert_if_new_inner::( values, + hashes, make_payload_fn, observe_payload_fn, ) @@ -245,6 +273,7 @@ where fn insert_if_new_inner( &mut self, values: &ArrayRef, + hashes: &[u64], mut make_payload_fn: MP, mut observe_payload_fn: OP, ) where @@ -252,27 +281,17 @@ where 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::(); // 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) { diff --git a/datafusion/physical-plan/benches/dictionary_group_values.rs b/datafusion/physical-plan/benches/dictionary_group_values.rs index ded52aebd1100..bfc5862b06c1e 100644 --- a/datafusion/physical-plan/benches/dictionary_group_values.rs +++ b/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; @@ -109,10 +110,17 @@ fn bench_intern_emit(c: &mut Criterion) { new_group_values(schema.clone(), &GroupOrdering::None) .unwrap(), Vec::::with_capacity(size), + Vec::::new(), ) }, - |(gv, groups)| { - gv.intern(std::slice::from_ref(&array), groups).unwrap(); + |(gv, groups, hashes)| { + let seed = RandomState::with_seed(15395726432021054657); + hashes.clear(); + hashes.resize(array.len(), 0); + create_hashes(std::slice::from_ref(&array), &seed, hashes) + .unwrap(); + gv.intern(std::slice::from_ref(&array), groups, hashes) + .unwrap(); black_box(&*groups); black_box(gv.emit(EmitTo::All).unwrap()); }, @@ -154,11 +162,18 @@ fn bench_repeated_intern_emit(c: &mut Criterion) { new_group_values(schema.clone(), &GroupOrdering::None) .unwrap(), Vec::::with_capacity(size), + Vec::::new(), ) }, - |(gv, groups)| { + |(gv, groups, hashes)| { for arr in &batches { - gv.intern(std::slice::from_ref(arr), groups).unwrap(); + let seed = RandomState::with_seed(15395726432021054657); + hashes.clear(); + hashes.resize(arr.len(), 0); + create_hashes(std::slice::from_ref(arr), &seed, hashes) + .unwrap(); + gv.intern(std::slice::from_ref(arr), groups, hashes) + .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..6ceff31fd9f15 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -28,6 +28,7 @@ use arrow::array::{ArrayRef, Int32Array}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::hash_utils::{RandomState, create_hashes}; use datafusion_physical_plan::aggregates::group_values::GroupValues; use datafusion_physical_plan::aggregates::group_values::GroupValuesRows; use datafusion_physical_plan::aggregates::group_values::multi_group_by::GroupValuesColumn; @@ -95,9 +96,14 @@ fn bench_intern( batches: &[Vec], groups: &mut Vec, ) { + let seed = RandomState::with_seed(15395726432021054657); + let mut hashes = Vec::new(); for batch in batches { groups.clear(); - gv.intern(batch, groups).unwrap(); + hashes.clear(); + hashes.resize(batch[0].len(), 0); + create_hashes(batch, &seed, &mut hashes).unwrap(); + gv.intern(batch, groups, &hashes).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..ae1068c00e6b6 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,64 @@ impl MaterializedAggregateOutput { } } +pub(super) fn try_new_internal_batch( + output_schema: SchemaRef, + mut columns: Vec, + has_group_hash: bool, +) -> Result { + if has_group_hash { + 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 create_batch_hashes( + &mut self, + group_values: &[ArrayRef], + ) -> Result<&[u64]> { + self.batch_hashes.clear(); + 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 reuse_batch_hashes(&mut self, input_hashes: &UInt64Array) -> &[u64] { + self.batch_hashes.clear(); + self.batch_hashes.extend_from_slice(input_hashes.values()); + &self.batch_hashes + } + + pub(super) fn record_new_group_hashes( + &mut self, + starting_num_groups: usize, + total_num_groups: usize, + ) { + self.group_hashes.resize(total_num_groups, 0); + for (row, group_index) in self.batch_group_indices.iter().enumerate() { + if *group_index >= starting_num_groups { + 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..ed6500d6a4e44 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 @@ -35,7 +35,7 @@ use crate::aggregates::grouped_hash_stream::create_group_accumulator; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, - evaluate_group_by, + create_group_hash_array, evaluate_group_by, schema_with_group_hash, }; use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; @@ -266,9 +266,13 @@ 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)?; + let mut hashes = Vec::new(); + crate::aggregates::create_group_hashes(group_values, &mut hashes)?; + self.buffer.group_values.intern( + group_values, + &mut self.buffer.group_indices, + &hashes, + )?; let total_num_groups = self.buffer.group_values.len(); if total_num_groups > starting_num_groups { self.buffer.group_ordering.new_groups( @@ -330,6 +334,11 @@ impl OrderedAggregateTable { let timer = self.group_by_metrics.emitting_time.timer(); let mut output = self.buffer.group_values.emit(emit_to)?; + let group_hashes = if is_final || self.buffer.accumulators.is_empty() { + None + } else { + Some(create_group_hash_array(&output)?) + }; if should_remove_groups { match emit_to { EmitTo::First(n) => self.buffer.group_ordering.remove_groups(n), @@ -347,9 +356,15 @@ impl OrderedAggregateTable { output.extend(acc.state(emit_to)?); } } + let output_schema = if let Some(group_hashes) = group_hashes { + output.push(group_hashes); + schema_with_group_hash(&self.output_schema) + } else { + Arc::clone(&self.output_schema) + }; drop(timer); - let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; + let batch = RecordBatch::try_new(output_schema, output)?; debug_assert!(batch.num_rows() > 0); Ok(Some(batch)) 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..24076121366f3 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,16 +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 group_indices = &state.batch_group_indices; + state.reuse_batch_hashes(input_hashes); + state.group_values.intern( + group_values, + &mut state.batch_group_indices, + &state.batch_hashes, + )?; let total_num_groups = state.group_values.len(); + let group_indices = &state.batch_group_indices; for (acc, values) in state .accumulators 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..85fb9c303f44f 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, try_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, true)?; debug_assert!(batch.num_rows() > 0); Ok(MaterializedAggregateOutput::new(batch)) } @@ -118,16 +120,26 @@ impl AggregateHashTable { &mut self, batch: &RecordBatch, ) -> Result<()> { - let evaluated_batch = self.evaluate_batch(batch)?; + let (batch, input_hashes) = try_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 group_indices = &state.batch_group_indices; + let starting_num_groups = state.group_values.len(); + if let Some(input_hashes) = input_hashes { + state.reuse_batch_hashes(input_hashes); + } else { + state.create_batch_hashes(group_values)?; + } + state.group_values.intern( + group_values, + &mut state.batch_group_indices, + &state.batch_hashes, + )?; let total_num_groups = state.group_values.len(); + state.record_new_group_hashes(starting_num_groups, total_num_groups); + let group_indices = &state.batch_group_indices; for (acc, values) in state .accumulators 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..5ca6ea933e6be 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, +}; 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, true)?; 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, @@ -170,11 +177,16 @@ impl AggregateHashTable { 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 group_indices = &state.batch_group_indices; + let starting_num_groups = state.group_values.len(); + state.create_batch_hashes(group_values)?; + state.group_values.intern( + group_values, + &mut state.batch_group_indices, + &state.batch_hashes, + )?; let total_num_groups = state.group_values.len(); + state.record_new_group_hashes(starting_num_groups, total_num_groups); + let group_indices = &state.batch_group_indices; for (acc, values) in state .accumulators @@ -215,12 +227,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 +253,15 @@ 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.create_batch_hashes(&cols)?; + state.group_values.intern( + &cols, + &mut state.batch_group_indices, + &state.batch_hashes, + )?; + let total_num_groups = state.group_values.len(); + state.record_new_group_hashes(starting_num_groups, total_num_groups); any_interned = true; } @@ -280,8 +299,8 @@ impl AggregateHashTable { .into_iter() .next() .unwrap_or_default(); - let state = self.state.building_mut(); + let group_hashes = create_group_hash_array(&output)?; for (acc, values) in state .accumulators .iter_mut() @@ -289,9 +308,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..c7bbeb655fe31 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: &[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..1d36229aac5c4 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 @@ -39,10 +39,8 @@ use arrow::datatypes::{ TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; -use datafusion_common::hash_utils::RandomState; -use datafusion_common::hash_utils::create_hashes; use datafusion_common::{Result, internal_datafusion_err, not_impl_err}; -use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_execution::memory_pool::proxy::HashTableAllocExt; use datafusion_expr::EmitTo; use datafusion_physical_expr::binary_map::OutputType; @@ -214,12 +212,6 @@ pub struct GroupValuesColumn { /// /// [`GroupValuesRows`]: crate::aggregates::group_values::GroupValuesRows group_values: Vec>, - - /// reused buffer to store hashes - hashes_buffer: Vec, - - /// Random state for creating hashes - random_state: RandomState, } /// Buffers to store intermediate results in `vectorized_append` @@ -281,8 +273,6 @@ impl GroupValuesColumn { vectorized_operation_buffers: VectorizedOperationBuffers::default(), map_size: 0, group_values, - hashes_buffer: Default::default(), - random_state: crate::aggregates::AGGREGATION_HASH_SEED, }) } @@ -348,19 +338,12 @@ impl GroupValuesColumn { &mut self, cols: &[ArrayRef], groups: &mut Vec, + hashes: &[u64], ) -> Result<()> { - let n_rows = cols[0].len(); - // tracks to which group each of the input rows belongs groups.clear(); - // 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)?; - - for (row, &target_hash) in batch_hashes.iter().enumerate() { + for (row, &target_hash) in hashes.iter().enumerate() { let entry = self .map .find_mut(target_hash, |(exist_hash, group_idx_view)| { @@ -449,6 +432,7 @@ impl GroupValuesColumn { &mut self, cols: &[ArrayRef], groups: &mut Vec, + hashes: &[u64], ) -> Result<()> { let n_rows = cols[0].len(); @@ -456,11 +440,6 @@ impl GroupValuesColumn { groups.clear(); groups.resize(n_rows, usize::MAX); - 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)?; - // General steps for one round `vectorized equal_to & append`: // 1. Collect vectorized context by checking hash values of `cols` in `map`, // mainly fill `vectorized_append_row_indices`, `vectorized_equal_to_row_indices` @@ -482,7 +461,7 @@ impl GroupValuesColumn { // // 1. Collect vectorized context by checking hash values of `cols` in `map` - self.collect_vectorized_process_context(&batch_hashes, groups); + self.collect_vectorized_process_context(hashes, groups); // 2. Perform `vectorized_append` self.vectorized_append(cols)?; @@ -492,9 +471,7 @@ impl GroupValuesColumn { // 4. Perform scalarized inter for remaining rows // (about remaining rows, can see comments for `remaining_row_indices`) - self.scalarized_intern_remaining(cols, &batch_hashes, groups)?; - - self.hashes_buffer = batch_hashes; + self.scalarized_intern_remaining(cols, hashes, groups)?; Ok(()) } @@ -1078,20 +1055,25 @@ 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: &[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) } } fn size(&self) -> usize { let group_values_size: usize = self.group_values.iter().map(|v| v.size()).sum(); - group_values_size + self.map_size + self.hashes_buffer.allocated_size() + group_values_size + self.map_size } fn is_empty(&self) -> bool { @@ -1224,8 +1206,6 @@ impl GroupValues for GroupValuesColumn { self.map.clear(); self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared self.map_size = self.map.capacity() * size_of::<(u64, usize)>(); - self.hashes_buffer.clear(); - self.hashes_buffer.shrink_to(num_rows); // Such structures are only used in `non-streaming` case if !STREAMING { @@ -1878,7 +1858,9 @@ 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(); + let mut hashes = Vec::new(); + crate::aggregates::create_group_hashes(batch, &mut hashes).unwrap(); + group_values.intern(batch, &mut vec![], &hashes).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..2bcb1ed6cf102 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -24,10 +24,8 @@ use arrow::compute::cast; use arrow::datatypes::{DataType, SchemaRef}; use arrow::row::{RowConverter, Rows, SortField}; use datafusion_common::Result; -use datafusion_common::hash_utils::RandomState; -use datafusion_common::hash_utils::create_hashes; use datafusion_common::utils::normalize_float_zero; -use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_execution::memory_pool::proxy::HashTableAllocExt; use datafusion_expr::EmitTo; use hashbrown::hash_table::HashTable; use log::debug; @@ -72,14 +70,8 @@ pub struct GroupValuesRows { /// [`Row`]: arrow::row::Row group_values: Option, - /// reused buffer to store hashes - hashes_buffer: Vec, - /// reused buffer to store rows rows_buffer: Rows, - - /// Random state for creating hashes - random_state: RandomState, } impl GroupValuesRows { @@ -108,15 +100,18 @@ impl GroupValuesRows { map, map_size: 0, group_values: None, - hashes_buffer: Default::default(), rows_buffer, - random_state: crate::aggregates::AGGREGATION_HASH_SEED, }) } } impl GroupValues for GroupValuesRows { - fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec, + hashes: &[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. @@ -128,8 +123,6 @@ impl GroupValues for GroupValuesRows { let group_rows = &mut self.rows_buffer; group_rows.clear(); self.row_converter.append(group_rows, cols)?; - let n_rows = group_rows.num_rows(); - let mut group_values = match self.group_values.take() { Some(group_values) => group_values, None => self.row_converter.empty_rows(0, 0), @@ -138,13 +131,7 @@ impl GroupValues for GroupValuesRows { // tracks to which group each of the input rows belongs groups.clear(); - // 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)?; - - for (row, &target_hash) in batch_hashes.iter().enumerate() { + for (row, &target_hash) in hashes.iter().enumerate() { let entry = self.map.find_mut(target_hash, |(exist_hash, group_idx)| { // Somewhat surprisingly, this closure can be called even if the // hash doesn't match, so check the hash first with an integer @@ -189,7 +176,6 @@ impl GroupValues for GroupValuesRows { + group_values_size + self.map_size + self.rows_buffer.size() - + self.hashes_buffer.allocated_size() } fn is_empty(&self) -> bool { @@ -262,8 +248,6 @@ impl GroupValues for GroupValuesRows { self.map.clear(); self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared self.map_size = self.map.capacity() * size_of::<(u64, usize)>(); - self.hashes_buffer.clear(); - self.hashes_buffer.shrink_to(num_rows); } } 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..1a698ef94bdee 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: &[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..fb7e2c40bc417 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,15 +45,21 @@ 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: &[u64], + ) -> Result<()> { assert_eq!(cols.len(), 1); // look up / add entries in the table let arr = &cols[0]; groups.clear(); - self.map.insert_if_new( + self.map.insert_if_new_with_hashes( arr, + hashes, // called for each new group |_value| { // assign new group index on each insert @@ -108,7 +114,12 @@ impl GroupValues for GroupValuesBytes { self.num_groups = 0; let mut group_indexes = vec![]; - self.intern(&[remaining_group_values], &mut group_indexes)?; + let mut hashes = Vec::new(); + crate::aggregates::create_group_hashes( + std::slice::from_ref(&remaining_group_values), + &mut hashes, + )?; + self.intern(&[remaining_group_values], &mut group_indexes, &hashes)?; // 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..d51dab51c02c8 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: &[u64], ) -> datafusion_common::Result<()> { assert_eq!(cols.len(), 1); @@ -54,8 +55,9 @@ impl GroupValues for GroupValuesBytesView { let arr = &cols[0]; groups.clear(); - self.map.insert_if_new( + self.map.insert_if_new_with_hashes( arr, + hashes, // called for each new group |_value| { // assign new group index on each insert @@ -110,7 +112,12 @@ impl GroupValues for GroupValuesBytesView { self.num_groups = 0; let mut group_indexes = vec![]; - self.intern(&[remaining_group_values], &mut group_indexes)?; + let mut hashes = Vec::new(); + crate::aggregates::create_group_hashes( + std::slice::from_ref(&remaining_group_values), + &mut hashes, + )?; + self.intern(&[remaining_group_values], &mut group_indexes, &hashes)?; // 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..5611d4c17c597 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 @@ -114,8 +114,6 @@ pub struct GroupValuesPrimitive { null_group: Option, /// The values for each group index values: Vec, - /// The random state used to generate hashes - random_state: RandomState, } impl GroupValuesPrimitive { @@ -126,7 +124,6 @@ impl GroupValuesPrimitive { map: HashTable::with_capacity(128), values: Vec::with_capacity(128), null_group: None, - random_state: crate::aggregates::AGGREGATION_HASH_SEED, } } } @@ -135,11 +132,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: &[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 +153,7 @@ 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[row]; let insert = self.map.entry( hash, |&(g, h)| unsafe { @@ -273,7 +274,9 @@ 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)?; + let mut hashes = Vec::new(); + crate::aggregates::create_group_hashes(std::slice::from_ref(&arr), &mut hashes)?; + gv.intern(&[arr], &mut groups, &hashes)?; 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..1f8ff3c09e37c 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,13 @@ 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)?; + let mut hashes = Vec::new(); + aggregates::create_group_hashes(group_values, &mut hashes)?; + self.group_values.intern( + group_values, + &mut self.current_group_indices, + &hashes, + )?; let group_indices = &self.current_group_indices; // Update ordering information if necessary @@ -1034,6 +1039,13 @@ 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.accumulators.is_empty() + && (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 +1060,21 @@ 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.accumulators.is_empty() + && (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); @@ -1103,8 +1125,10 @@ impl GroupedHashAggregateStream { cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); let starting_groups = self.group_values.len(); + let mut hashes = Vec::new(); + aggregates::create_group_hashes(&cols, &mut hashes)?; self.group_values - .intern(&cols, &mut self.current_group_indices)?; + .intern(&cols, &mut self.current_group_indices, &hashes)?; let total_groups = self.group_values.len(); if total_groups > starting_groups { self.group_ordering.new_groups( @@ -1383,6 +1407,9 @@ impl GroupedHashAggregateStream { "group_values expected to have single element" ); let mut output = group_values.swap_remove(0); + let group_hashes = (!self.accumulators.is_empty()) + .then(|| create_group_hash_array(&output)) + .transpose()?; let iter = self .accumulators @@ -1394,8 +1421,16 @@ impl GroupedHashAggregateStream { let opt_filter = opt_filter.as_ref().map(|filter| filter.as_boolean()); output.extend(acc.convert_to_state(values, opt_filter)?); } + if let Some(group_hashes) = group_hashes { + output.push(group_hashes); + } - let states_batch = RecordBatch::try_new(self.schema(), output)?; + let states_schema = if self.accumulators.is_empty() { + self.schema() + } else { + schema_with_group_hash(&self.schema()) + }; + let states_batch = RecordBatch::try_new(states_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..9d3cbbde05dc7 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,80 @@ 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_hashes( + group_values: &[ArrayRef], + hashes: &mut Vec, +) -> Result<()> { + let num_rows = group_values.first().map(|array| array.len()).unwrap_or(0); + hashes.clear(); + hashes.resize(num_rows, 0); + datafusion_common::hash_utils::create_hashes( + group_values, + &AGGREGATION_HASH_SEED, + hashes, + )?; + Ok(()) +} + +pub(crate) fn create_group_hash_array(group_values: &[ArrayRef]) -> Result { + let mut hashes = Vec::new(); + create_group_hashes(group_values, &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 try_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)); + }; + + strip_group_hash_column_at(batch, hash_index) + .map(|(batch, hashes)| (batch, Some(hashes))) +} + +pub(crate) fn strip_group_hash_column( + batch: &RecordBatch, +) -> Result<(RecordBatch, &UInt64Array)> { + let hash_index = group_hash_column_index(batch.schema().as_ref()) + .expect("input must include internal group hash column"); + strip_group_hash_column_at(batch, hash_index) +} + +fn strip_group_hash_column_at( + batch: &RecordBatch, + hash_index: usize, +) -> Result<(RecordBatch, &UInt64Array)> { + 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)?, hashes)) +} + /// Whether an aggregate stage consumes raw input data or intermediate /// accumulator state from a previous aggregation stage. /// @@ -2518,7 +2592,10 @@ mod tests { }; use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::Int32Type; - use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; + use datafusion_common::test_util::{ + batches_to_sort_string as format_batches_to_sort_string, + batches_to_string as format_batches_to_string, + }; use datafusion_common::{DataFusionError, internal_err}; use datafusion_execution::config::SessionConfig; use datafusion_execution::memory_pool::FairSpillPool; @@ -2558,6 +2635,23 @@ mod tests { Ok(schema) } + fn batches_without_group_hash(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .map(|batch| try_strip_group_hash_column(batch).map(|(batch, _hashes)| batch)) + .collect::>>() + .unwrap() + } + + fn batches_to_sort_string(batches: &[RecordBatch]) -> String { + let batches = batches_without_group_hash(batches); + format_batches_to_sort_string(&batches) + } + + fn batches_to_string(batches: &[RecordBatch]) -> String { + let batches = batches_without_group_hash(batches); + format_batches_to_string(&batches) + } /// some mock data to aggregates fn some_data() -> (Arc, Vec) { // define a schema. diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 7289ac43e510c..c8e2430f1185b 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -466,8 +466,13 @@ impl DistinctDeduplicator { "failed to reserve {additional} recursive query group ids: {e}" ) })?; - self.group_values - .intern(batch.columns(), &mut self.intern_output_buffer)?; + let mut hashes = Vec::new(); + crate::aggregates::create_group_hashes(batch.columns(), &mut hashes)?; + self.group_values.intern( + batch.columns(), + &mut self.intern_output_buffer, + &hashes, + )?; 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());