This is an automated email from the ASF dual-hosted git repository. gabotechs pushed a commit to branch gabotechs/reuse-hashes in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 680e0ddeab95322c42da625eff80c7f5685cdba1 Author: Gabriel <[email protected]> AuthorDate: Sat Jul 18 18:07:31 2026 +0200 Reuse hashes in aggregations --- datafusion/core/tests/dataframe/mod.rs | 4 + .../physical_optimizer/partition_statistics.rs | 20 +- datafusion/physical-expr-common/src/binary_map.rs | 45 ++- .../physical-expr-common/src/binary_view_map.rs | 45 ++- .../benches/dictionary_group_values.rs | 27 +- datafusion/physical-plan/benches/multi_group_by.rs | 27 +- .../src/aggregates/aggregate_hash_table/common.rs | 59 +++- .../aggregate_hash_table/common_ordered.rs | 58 +++- .../aggregate_hash_table/partial_table.rs | 40 ++- .../src/aggregates/group_values/mod.rs | 8 +- .../aggregates/group_values/multi_group_by/mod.rs | 54 ++-- .../src/aggregates/group_values/row.rs | 31 +-- .../group_values/single_group_by/boolean.rs | 8 +- .../group_values/single_group_by/bytes.rs | 15 +- .../group_values/single_group_by/bytes_view.rs | 9 +- .../group_values/single_group_by/primitive.rs | 20 +- .../src/aggregates/grouped_hash_stream.rs | 93 ++++++- .../src/aggregates/grouped_topk_stream.rs | 17 +- datafusion/physical-plan/src/aggregates/mod.rs | 303 ++++++++++++++++++++- datafusion/physical-plan/src/recursive_query.rs | 22 +- datafusion/physical-plan/src/repartition/hash.rs | 177 ++++++++++++ datafusion/physical-plan/src/repartition/mod.rs | 60 ++-- .../test_files/push_down_filter_parquet.slt | 6 +- 23 files changed, 971 insertions(+), 177 deletions(-) diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index db26413ac9..c392b163aa 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -7180,6 +7180,10 @@ async fn test_duplicate_state_fields_for_dfschema_construct() -> Result<()> { let binding = partial_agg.schema(); let actual_field_names: Vec<_> = binding.fields().iter().map(|f| f.name()).collect(); + let (hash_field_name, actual_field_names) = actual_field_names + .split_last() + .expect("Partial aggregate schema should not be empty"); + assert!(hash_field_name.starts_with("__datafusion_internal_hash_")); assert_eq!(actual_field_names, expected_field_names); // Ensure that DFSchema::try_from does not fail diff --git a/datafusion/core/tests/physical_optimizer/partition_statistics.rs b/datafusion/core/tests/physical_optimizer/partition_statistics.rs index 6cabcdb710..2aaa303268 100644 --- a/datafusion/core/tests/physical_optimizer/partition_statistics.rs +++ b/datafusion/core/tests/physical_optimizer/partition_statistics.rs @@ -911,8 +911,12 @@ mod test { ColumnStatistics::new_unknown(), ], }; + let mut expected_partial_p0_statistics = expected_p0_statistics.clone(); + expected_partial_p0_statistics + .column_statistics + .push(ColumnStatistics::new_unknown()); - assert_eq!(*p0_statistics, expected_p0_statistics); + assert_eq!(*p0_statistics, expected_partial_p0_statistics); let expected_p1_statistics = Statistics { num_rows: Precision::Inexact(2), @@ -930,12 +934,16 @@ mod test { ColumnStatistics::new_unknown(), ], }; + let mut expected_partial_p1_statistics = expected_p1_statistics.clone(); + expected_partial_p1_statistics + .column_statistics + .push(ColumnStatistics::new_unknown()); let p1_statistics = StatisticsContext::new().compute( aggregate_exec_partial.as_ref(), &StatisticsArgs::new().with_partition(Some(1)), )?; - assert_eq!(*p1_statistics, expected_p1_statistics); + assert_eq!(*p1_statistics, expected_partial_p1_statistics); validate_statistics_with_data( aggregate_exec_partial.clone(), @@ -1009,16 +1017,20 @@ mod test { ColumnStatistics::new_unknown(), ], }; + let mut empty_partial_stat = empty_stat.clone(); + empty_partial_stat + .column_statistics + .push(ColumnStatistics::new_unknown()); assert_eq!( - empty_stat, + empty_partial_stat, *StatisticsContext::new().compute( agg_partial.as_ref(), &StatisticsArgs::new().with_partition(Some(0)) )? ); assert_eq!( - empty_stat, + empty_partial_stat, *StatisticsContext::new().compute( agg_partial.as_ref(), &StatisticsArgs::new().with_partition(Some(1)) diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index ad184d6500..77bce64eed 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -299,6 +299,34 @@ where MP: FnMut(Option<&[u8]>) -> V, OP: FnMut(V), { + let mut hashes = std::mem::take(&mut self.hashes_buffer); + hashes.clear(); + hashes.resize(values.len(), 0); + create_hashes([values], &self.random_state, &mut hashes) + // hash is supported for all types and create_hashes only + // returns errors for unsupported types + .unwrap(); + self.insert_if_new_with_hashes( + values, + &hashes, + make_payload_fn, + observe_payload_fn, + ); + self.hashes_buffer = hashes; + } + + /// Like [`Self::insert_if_new`], but uses a hash supplied for each input value. + pub fn insert_if_new_with_hashes<MP, OP>( + &mut self, + values: &ArrayRef, + hashes: &[u64], + make_payload_fn: MP, + observe_payload_fn: OP, + ) where + MP: FnMut(Option<&[u8]>) -> V, + OP: FnMut(V), + { + assert_eq!(values.len(), hashes.len()); // Sanity array type match self.output_type { OutputType::Binary => { @@ -308,6 +336,7 @@ where )); self.insert_if_new_inner::<MP, OP, GenericBinaryType<O>>( values, + hashes, make_payload_fn, observe_payload_fn, ) @@ -319,6 +348,7 @@ where )); self.insert_if_new_inner::<MP, OP, GenericStringType<O>>( values, + hashes, make_payload_fn, observe_payload_fn, ) @@ -338,6 +368,7 @@ where fn insert_if_new_inner<MP, OP, B>( &mut self, values: &ArrayRef, + hashes: &[u64], mut make_payload_fn: MP, mut observe_payload_fn: OP, ) where @@ -345,22 +376,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::<B>(); // Ensure lengths are equivalent - assert_eq!(values.len(), batch_hashes.len()); + assert_eq!(values.len(), hashes.len()); - for (value, &hash) in values.iter().zip(batch_hashes.iter()) { + for (value, &hash) in values.iter().zip(hashes) { // handle null value let Some(value) = value else { let payload = if let Some(&(payload, _offset)) = self.null.as_ref() { diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 9d4b556393..34b4df8b25 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -212,12 +212,41 @@ where MP: FnMut(Option<&[u8]>) -> V, OP: FnMut(V), { + let mut hashes = std::mem::take(&mut self.hashes_buffer); + hashes.clear(); + hashes.resize(values.len(), 0); + create_hashes([values], &self.random_state, &mut hashes) + // hash is supported for all types and create_hashes only + // returns errors for unsupported types + .unwrap(); + self.insert_if_new_with_hashes( + values, + &hashes, + make_payload_fn, + observe_payload_fn, + ); + self.hashes_buffer = hashes; + } + + /// Like [`Self::insert_if_new`], but uses a hash supplied for each input value. + pub fn insert_if_new_with_hashes<MP, OP>( + &mut self, + values: &ArrayRef, + hashes: &[u64], + make_payload_fn: MP, + observe_payload_fn: OP, + ) where + MP: FnMut(Option<&[u8]>) -> V, + OP: FnMut(V), + { + assert_eq!(values.len(), hashes.len()); // Sanity check array type match self.output_type { OutputType::BinaryView => { assert!(matches!(values.data_type(), DataType::BinaryView)); self.insert_if_new_inner::<MP, OP, BinaryViewType>( values, + hashes, make_payload_fn, observe_payload_fn, ) @@ -226,6 +255,7 @@ where assert!(matches!(values.data_type(), DataType::Utf8View)); self.insert_if_new_inner::<MP, OP, StringViewType>( values, + hashes, make_payload_fn, observe_payload_fn, ) @@ -245,6 +275,7 @@ where fn insert_if_new_inner<MP, OP, B>( &mut self, values: &ArrayRef, + hashes: &[u64], mut make_payload_fn: MP, mut observe_payload_fn: OP, ) where @@ -252,27 +283,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::<B>(); // Get raw views buffer for direct comparison let input_views = values.views(); // Ensure lengths are equivalent - assert_eq!(values.len(), self.hashes_buffer.len()); + assert_eq!(values.len(), hashes.len()); for i in 0..values.len() { let view_u128 = input_views[i]; - let hash = self.hashes_buffer[i]; + let hash = hashes[i]; // handle null value via validity bitmap check if values.is_null(i) { diff --git a/datafusion/physical-plan/benches/dictionary_group_values.rs b/datafusion/physical-plan/benches/dictionary_group_values.rs index ded52aebd1..756860da9f 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; @@ -41,6 +42,18 @@ const CARDS_RELATIVE: [usize; 4] = [20, 75, 300, 1000]; const N_BATCHES: usize = 4; // Fixed for reproducibility. const SEED: u64 = 0xD1C7; +const AGGREGATION_HASH_SEED: u64 = 15395726432021054657; + +fn create_aggregation_hashes(array: &ArrayRef, hashes: &mut Vec<u64>) { + hashes.clear(); + hashes.resize(array.len(), 0); + create_hashes( + std::slice::from_ref(array), + &RandomState::with_seed(AGGREGATION_HASH_SEED), + hashes, + ) + .unwrap(); +} fn dict_schema() -> SchemaRef { let dict_ty = @@ -109,10 +122,13 @@ fn bench_intern_emit(c: &mut Criterion) { new_group_values(schema.clone(), &GroupOrdering::None) .unwrap(), Vec::<usize>::with_capacity(size), + Vec::<u64>::with_capacity(size), ) }, - |(gv, groups)| { - gv.intern(std::slice::from_ref(&array), groups).unwrap(); + |(gv, groups, hashes)| { + create_aggregation_hashes(&array, hashes); + gv.intern(std::slice::from_ref(&array), groups, hashes) + .unwrap(); black_box(&*groups); black_box(gv.emit(EmitTo::All).unwrap()); }, @@ -154,11 +170,14 @@ fn bench_repeated_intern_emit(c: &mut Criterion) { new_group_values(schema.clone(), &GroupOrdering::None) .unwrap(), Vec::<usize>::with_capacity(size), + Vec::<u64>::with_capacity(size), ) }, - |(gv, groups)| { + |(gv, groups, hashes)| { for arr in &batches { - gv.intern(std::slice::from_ref(arr), groups).unwrap(); + create_aggregation_hashes(arr, hashes); + gv.intern(std::slice::from_ref(arr), groups, hashes) + .unwrap(); black_box(&*groups); } black_box(gv.emit(EmitTo::All).unwrap()); diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 92d0448775..d491dc9684 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; @@ -35,6 +36,7 @@ use std::hint::black_box; use std::sync::Arc; const DEFAULT_BATCH_SIZE: usize = 8192; +const AGGREGATION_HASH_SEED: u64 = 15395726432021054657; fn make_schema(num_cols: usize) -> SchemaRef { let fields: Vec<Field> = (0..num_cols) @@ -94,10 +96,15 @@ fn bench_intern( gv: &mut Box<dyn GroupValues>, batches: &[Vec<ArrayRef>], groups: &mut Vec<usize>, + hashes: &mut Vec<u64>, ) { + let random_state = RandomState::with_seed(AGGREGATION_HASH_SEED); for batch in batches { groups.clear(); - gv.intern(batch, groups).unwrap(); + hashes.clear(); + hashes.resize(batch[0].len(), 0); + create_hashes(batch, &random_state, hashes).unwrap(); + gv.intern(batch, groups, hashes).unwrap(); } black_box(&*groups); } @@ -131,9 +138,10 @@ fn bench_issue_17850_regression(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + Vec::<u64>::with_capacity(DEFAULT_BATCH_SIZE), ) }, - |(gv, groups)| bench_intern(gv, batches, groups), + |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), criterion::BatchSize::LargeInput, ); }, @@ -174,9 +182,10 @@ fn bench_low_cardinality(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + Vec::<u64>::with_capacity(DEFAULT_BATCH_SIZE), ) }, - |(gv, groups)| bench_intern(gv, batches, groups), + |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), criterion::BatchSize::LargeInput, ); }, @@ -213,9 +222,10 @@ fn bench_batch_size_sensitivity(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::<usize>::with_capacity(batch_size), + Vec::<u64>::with_capacity(batch_size), ) }, - |(gv, groups)| bench_intern(gv, batches, groups), + |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), criterion::BatchSize::LargeInput, ); }, @@ -253,9 +263,10 @@ fn bench_column_scaling(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + Vec::<u64>::with_capacity(DEFAULT_BATCH_SIZE), ) }, - |(gv, groups)| bench_intern(gv, batches, groups), + |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), criterion::BatchSize::LargeInput, ); }, @@ -291,9 +302,10 @@ fn bench_high_cardinality_scaling(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + Vec::<u64>::with_capacity(DEFAULT_BATCH_SIZE), ) }, - |(gv, groups)| bench_intern(gv, batches, groups), + |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), criterion::BatchSize::LargeInput, ); }, @@ -332,9 +344,10 @@ fn bench_group_count_sweep(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + Vec::<u64>::with_capacity(DEFAULT_BATCH_SIZE), ) }, - |(gv, groups)| bench_intern(gv, batches, groups), + |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), criterion::BatchSize::LargeInput, ); }, 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 eaf39929ce..71dfb1571d 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -31,8 +31,10 @@ 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, + AggregateExec, GroupHashTracker, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, outputs_group_hashes, }; +use crate::repartition::{ExpressionHasher, HashMetrics}; /// Marker for raw rows -> partial state aggregation. pub(in crate::aggregates) struct PartialMarker; @@ -138,6 +140,12 @@ impl<AggrMode> AggregateHashTable<AggrMode> { group_by: Arc::clone(&agg.group_by), group_values, batch_group_indices: Default::default(), + group_hash_tracker: outputs_group_hashes(agg.mode, &agg.group_by) + .then(GroupHashTracker::default), + hasher: ExpressionHasher::new_with_metrics( + agg.group_by.input_exprs(), + HashMetrics::new(&agg.metrics, partition), + ), accumulators, }), _mode: PhantomData, @@ -145,10 +153,10 @@ impl<AggrMode> AggregateHashTable<AggrMode> { } /// See comments in [`EvaluatedAggregateBatch`] - pub(super) fn evaluate_batch( + pub(super) fn evaluate_batch<'a>( &self, - batch: &RecordBatch, - ) -> Result<EvaluatedAggregateBatch> { + batch: &'a RecordBatch, + ) -> Result<EvaluatedAggregateBatch<'a>> { let state = self.state.building(); let timer = self.group_by_metrics.time_calculating_group_ids.timer(); // outer vec: one per each grouping set @@ -170,6 +178,9 @@ impl<AggrMode> AggregateHashTable<AggrMode> { Ok(EvaluatedAggregateBatch { grouping_set_args, accumulator_args, + precomputed_group_hashes: state + .hasher + .precomputed_group_by(&state.group_by, batch), }) } @@ -189,9 +200,24 @@ impl<AggrMode> AggregateHashTable<AggrMode> { 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 previous_group_count = state.group_values.len(); + let hashes = match evaluated_batch.precomputed_group_hashes { + Some(hashes) => hashes, + None => state.hasher.compute_hashes(group_values)?, + }; + state.group_values.intern( + group_values, + &mut state.batch_group_indices, + hashes, + )?; + if let Some(group_hash_tracker) = &mut state.group_hash_tracker { + group_hash_tracker.record_new_groups( + previous_group_count, + state.group_values.len(), + &state.batch_group_indices, + hashes, + ); + } let group_indices = &state.batch_group_indices; let total_num_groups = state.group_values.len(); @@ -238,6 +264,9 @@ impl<AggrMode> AggregateHashTable<AggrMode> { for acc in state.accumulators.iter_mut() { columns.extend(materialize_accumulator_fn(acc, emit_to)?); } + if let Some(group_hash_tracker) = &mut state.group_hash_tracker { + columns.push(group_hash_tracker.emit(emit_to)); + } drop(timer); let batch = RecordBatch::try_new(output_schema, columns)?; @@ -274,6 +303,11 @@ impl<AggrMode> AggregateHashTable<AggrMode> { acc + state.group_values.size() + state.batch_group_indices.allocated_size() + + state + .group_hash_tracker + .as_ref() + .map_or(0, |v| v.allocated_size()) + + state.hasher.allocated_size() } AggregateHashTableState::OutputtingMaterialized(output) => { output.memory_size() @@ -374,13 +408,16 @@ pub(super) struct EvaluatedAccumulatorArgs { /// /// e.g., `select k+1, sum(v*v) from t group by (k+1)`, this function evaluates /// `k+1`, `v*v` -pub(super) struct EvaluatedAggregateBatch { +pub(super) struct EvaluatedAggregateBatch<'a> { /// One entry per grouping set; each entry contains all evaluated group key /// arrays for the current input batch. pub(super) grouping_set_args: Vec<Vec<ArrayRef>>, /// Evaluated arguments and filters, one entry per aggregate expression. pub(super) accumulator_args: Vec<EvaluatedAccumulatorArgs>, + + /// Hashes carried by the input batch for a simple GROUP BY, when available. + pub(super) precomputed_group_hashes: Option<&'a [u64]>, } /// Buffer for the aggregate hash table's group keys and accumulator states. @@ -403,6 +440,12 @@ pub(super) struct AggregateHashTableBuffer { /// accumulator to update that group's aggregate state. pub(super) batch_group_indices: Vec<usize>, + /// Hashes retained for the internal hash column when output requires it. + pub(super) group_hash_tracker: Option<GroupHashTracker>, + + /// Computes group hashes when the input has no reusable hash column. + pub(super) hasher: ExpressionHasher, + /// One item per aggregate expression. /// /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input 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..f25c0f9bb0 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 @@ -34,9 +34,10 @@ 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, AggregateMode, PhysicalGroupBy, aggregate_expressions, - evaluate_group_by, + AggregateExec, AggregateMode, GroupHashTracker, PhysicalGroupBy, + aggregate_expressions, evaluate_group_by, outputs_group_hashes, }; +use crate::repartition::{ExpressionHasher, HashMetrics}; use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; @@ -114,6 +115,12 @@ pub(super) struct OrderedAggregateTableBuffer { /// Scratch group id vector for the current input batch. pub(super) group_indices: Vec<usize>, + /// Hashes retained for the internal hash column when output requires it. + pub(super) group_hash_tracker: Option<GroupHashTracker>, + + /// Computes group hashes when the input has no reusable hash column. + pub(super) hasher: ExpressionHasher, + /// One item per aggregate expression. /// /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input @@ -175,6 +182,12 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { group_ordering, group_values, group_indices: vec![], + group_hash_tracker: outputs_group_hashes(*aggregate_mode, &agg.group_by) + .then(GroupHashTracker::default), + hasher: ExpressionHasher::new_with_metrics( + agg.group_by.input_exprs(), + HashMetrics::new(&agg.metrics, partition), + ), accumulators, }, _mode: PhantomData, @@ -185,10 +198,10 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { /// /// e.g., `select k+1, sum(v*v) from t group by (k+1)`, this function /// evaluates `k+1`, `v*v`. - pub(super) fn evaluate_batch( + pub(super) fn evaluate_batch<'a>( &self, - batch: &RecordBatch, - ) -> Result<EvaluatedAggregateBatch> { + batch: &'a RecordBatch, + ) -> Result<EvaluatedAggregateBatch<'a>> { let timer = self.group_by_metrics.time_calculating_group_ids.timer(); let grouping_set_args = evaluate_group_by(&self.buffer.group_by, batch)?; drop(timer); @@ -205,6 +218,10 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { Ok(EvaluatedAggregateBatch { grouping_set_args, accumulator_args, + precomputed_group_hashes: self + .buffer + .hasher + .precomputed_group_by(&self.buffer.group_by, batch), }) } @@ -232,6 +249,12 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { + self.buffer.group_values.size() + self.buffer.group_ordering.size() + self.buffer.group_indices.allocated_size() + + self + .buffer + .group_hash_tracker + .as_ref() + .map_or(0, |v| v.allocated_size()) + + self.buffer.hasher.allocated_size() } /// Returns the [`EmitTo`], clamped to the specified batch size @@ -261,15 +284,29 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { /// - `false`: update aggregate states from raw input for partial aggregation. pub(super) fn aggregate_evaluated_batch( &mut self, - evaluated_batch: &EvaluatedAggregateBatch, + evaluated_batch: &EvaluatedAggregateBatch<'_>, is_final: bool, ) -> 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 hashes = match evaluated_batch.precomputed_group_hashes { + Some(hashes) => hashes, + None => self.buffer.hasher.compute_hashes(group_values)?, + }; + self.buffer.group_values.intern( + group_values, + &mut self.buffer.group_indices, + hashes, + )?; let total_num_groups = self.buffer.group_values.len(); + if let Some(group_hash_tracker) = &mut self.buffer.group_hash_tracker { + group_hash_tracker.record_new_groups( + starting_num_groups, + total_num_groups, + &self.buffer.group_indices, + hashes, + ); + } if total_num_groups > starting_num_groups { self.buffer.group_ordering.new_groups( group_values, @@ -347,6 +384,9 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { output.extend(acc.state(emit_to)?); } } + if let Some(group_hash_tracker) = &mut self.buffer.group_hash_tracker { + output.push(group_hash_tracker.emit(emit_to)) + } drop(timer); let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; 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 ffac42feaa..41949dc67a 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,14 +19,16 @@ 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}; 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, GroupHashTracker, group_id_array, max_duplicate_ordinal, +}; use super::common::{ AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, @@ -100,6 +102,11 @@ impl AggregateHashTable<PartialMarker> { group_by: Arc::clone(&state.group_by), group_values, batch_group_indices: Default::default(), + group_hash_tracker: state + .group_hash_tracker + .is_some() + .then(GroupHashTracker::default), + hasher: state.hasher.new_for_exprs(state.group_by.input_exprs()), accumulators, }), _mode: PhantomData, @@ -167,9 +174,19 @@ impl AggregateHashTable<PartialMarker> { .collect(); cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); + let previous_group_count = state.group_values.len(); + let hashes = state.hasher.compute_hashes(&cols)?; state .group_values - .intern(&cols, &mut state.batch_group_indices)?; + .intern(&cols, &mut state.batch_group_indices, hashes)?; + if let Some(group_hash_tracker) = &mut state.group_hash_tracker { + group_hash_tracker.record_new_groups( + previous_group_count, + state.group_values.len(), + &state.batch_group_indices, + hashes, + ); + } any_interned = true; } @@ -202,13 +219,25 @@ impl AggregateHashTable<PartialSkipMarker> { 1, "group_values expected to have single element" ); + let state = self.state.building_mut(); + let hash_array = if state.group_hash_tracker.is_some() { + let hashes = match evaluated_batch.precomputed_group_hashes { + Some(hashes) => hashes, + None => state + .hasher + .compute_hashes(&evaluated_batch.grouping_set_args[0])?, + }; + Some(Arc::new(UInt64Array::from(hashes.to_vec())) as ArrayRef) + } else { + None + }; + let mut output = evaluated_batch .grouping_set_args .into_iter() .next() .unwrap_or_default(); - let state = self.state.building_mut(); for (acc, values) in state .accumulators .iter_mut() @@ -216,6 +245,9 @@ impl AggregateHashTable<PartialSkipMarker> { { output.extend(acc.convert_to_state(values)?); } + if let Some(hash_array) = hash_array { + output.push(hash_array); + } Ok(RecordBatch::try_new( Arc::clone(&self.output_schema), diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index ee253e5d7a..d6024eea4a 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -93,11 +93,17 @@ pub trait GroupValues: Send { /// /// When the function returns, `groups` must contain the group id for each /// row in `cols`. + /// `hashes` contains the precomputed hash for each input row. /// /// 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: &[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..7e99ee184c 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<const STREAMING: bool> { /// /// [`GroupValuesRows`]: crate::aggregates::group_values::GroupValuesRows group_values: Vec<Box<dyn GroupColumn>>, - - /// reused buffer to store hashes - hashes_buffer: Vec<u64>, - - /// Random state for creating hashes - random_state: RandomState, } /// Buffers to store intermediate results in `vectorized_append` @@ -281,8 +273,6 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> { vectorized_operation_buffers: VectorizedOperationBuffers::default(), map_size: 0, group_values, - hashes_buffer: Default::default(), - random_state: crate::aggregates::AGGREGATION_HASH_SEED, }) } @@ -348,19 +338,15 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> { &mut self, cols: &[ArrayRef], groups: &mut Vec<usize>, + hashes: &[u64], ) -> Result<()> { let n_rows = cols[0].len(); + assert_eq!(n_rows, hashes.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,18 +435,15 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> { &mut self, cols: &[ArrayRef], groups: &mut Vec<usize>, + hashes: &[u64], ) -> Result<()> { let n_rows = cols[0].len(); + assert_eq!(n_rows, hashes.len()); // tracks to which group each of the input rows belongs 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 +465,7 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> { // // 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 +475,7 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> { // 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 +1059,25 @@ 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: &[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 +1210,6 @@ impl<const STREAMING: bool> GroupValues for GroupValuesColumn<STREAMING> { 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 { @@ -1877,8 +1861,10 @@ mod tests { } fn load_to_group_values(&self, group_values: &mut impl GroupValues) { + let mut hasher = crate::repartition::ExpressionHasher::new(vec![]); for batch in self.test_batches.iter() { - group_values.intern(batch, &mut vec![]).unwrap(); + let hashes = hasher.compute_hashes(batch).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 4976a098ec..0be654bdbf 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<Rows>, - /// reused buffer to store hashes - hashes_buffer: Vec<u64>, - /// 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<usize>) -> Result<()> { + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec<usize>, + 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. @@ -129,6 +124,7 @@ impl GroupValues for GroupValuesRows { group_rows.clear(); self.row_converter.append(group_rows, cols)?; let n_rows = group_rows.num_rows(); + assert_eq!(n_rows, hashes.len()); let mut group_values = match self.group_values.take() { Some(group_values) => group_values, @@ -138,13 +134,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 +179,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 +251,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 e993c0c53d..ac94c8d4ac 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,8 +42,14 @@ 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: &[u64], + ) -> Result<()> { let array = cols[0].as_boolean(); + assert_eq!(array.len(), hashes.len()); groups.clear(); for value in array.iter() { 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..71bcb98509 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<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: &[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,10 @@ impl<O: OffsetSizeTrait> GroupValues for GroupValuesBytes<O> { self.num_groups = 0; let mut group_indexes = vec![]; - self.intern(&[remaining_group_values], &mut group_indexes)?; + let cols = [remaining_group_values]; + let mut hasher = crate::repartition::ExpressionHasher::new(vec![]); + let hashes = hasher.compute_hashes(&cols)?; + self.intern(&cols, &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 7a56f7c52c..e4a2397b06 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: &[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,10 @@ impl GroupValues for GroupValuesBytesView { self.num_groups = 0; let mut group_indexes = vec![]; - self.intern(&[remaining_group_values], &mut group_indexes)?; + let cols = [remaining_group_values]; + let mut hasher = crate::repartition::ExpressionHasher::new(vec![]); + let hashes = hasher.compute_hashes(&cols)?; + self.intern(&cols, &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 e254aebcfd..b8f1356d95 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<T: ArrowPrimitiveType> { null_group: Option<usize>, /// The values for each group index values: Vec<T::Native>, - /// The random state used to generate hashes - random_state: RandomState, } impl<T: ArrowPrimitiveType> GroupValuesPrimitive<T> { @@ -126,7 +124,6 @@ impl<T: ArrowPrimitiveType> GroupValuesPrimitive<T> { map: HashTable::with_capacity(128), values: Vec::with_capacity(128), null_group: None, - random_state: crate::aggregates::AGGREGATION_HASH_SEED, } } } @@ -135,11 +132,17 @@ 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: &[u64], + ) -> Result<()> { assert_eq!(cols.len(), 1); + assert_eq!(cols[0].len(), hashes.len()); groups.clear(); - for v in cols[0].as_primitive::<T>() { + for (v, &hash) in cols[0].as_primitive::<T>().into_iter().zip(hashes) { let group_id = match v { None => *self.null_group.get_or_insert_with(|| { let group_id = self.values.len(); @@ -151,8 +154,6 @@ 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 insert = self.map.entry( hash, |&(g, h)| unsafe { @@ -273,7 +274,10 @@ 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 cols = [arr]; + let mut hasher = crate::repartition::ExpressionHasher::new(vec![]); + let hashes = hasher.compute_hashes(&cols)?; + gv.intern(&cols, &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 0d00e5c4d0..ea365bdb14 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -27,11 +27,12 @@ use super::{AggregateExec, format_human_display}; use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; 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, + AggregateInputMode, AggregateMode, AggregateOutputMode, GroupHashTracker, + PhysicalGroupBy, create_schema, evaluate_group_by, evaluate_many, evaluate_optional, + group_id_array, max_duplicate_ordinal, }; use crate::metrics::{BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput}; +use crate::repartition::{ExpressionHasher, HashMetrics}; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::spill_manager::{GetSlicedSize, SpillManager}; use crate::stream::EmptyRecordBatchStream; @@ -334,6 +335,12 @@ pub(crate) struct GroupedHashAggregateStream { /// processed. Reused across batches here to avoid reallocations current_group_indices: Vec<usize>, + /// Hashes retained for the internal hash column when output requires it. + group_hash_tracker: Option<GroupHashTracker>, + + /// Computes group hashes when the input has no reusable hash column. + hasher: ExpressionHasher, + /// Accumulators, one for each `AggregateFunctionExpr` in the query /// /// For example, if the query has aggregates, `SUM(x)`, @@ -587,6 +594,13 @@ impl GroupedHashAggregateStream { } else { None }; + let track_group_hashes = !agg_group_by.is_true_no_grouping() + && (agg.mode.output_mode() == AggregateOutputMode::Partial + || oom_mode == OutOfMemoryMode::Spill); + let hasher = ExpressionHasher::new_with_metrics( + agg_group_by.input_exprs(), + HashMetrics::new(&agg.metrics, partition), + ); Ok(GroupedHashAggregateStream { schema: agg_schema, @@ -601,6 +615,8 @@ impl GroupedHashAggregateStream { oom_mode, group_values, current_group_indices: Default::default(), + group_hash_tracker: track_group_hashes.then(GroupHashTracker::default), + hasher, exec_state, baseline_metrics, group_by_metrics, @@ -878,17 +894,39 @@ impl GroupedHashAggregateStream { evaluate_optional(&self.filter_expressions, batch)? }; + let group_by = match self.spill_state.is_stream_merging { + true => &self.spill_state.merging_group_by, + false => &self.group_by, + }; + + let precomputed_hashes = self.hasher.precomputed_group_by(group_by, batch); + for group_values in &group_by_values { let groups_start_time = Instant::now(); // 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 hashes = match precomputed_hashes { + Some(hashes) => hashes, + None => self.hasher.compute_hashes(group_values)?, + }; + self.group_values.intern( + group_values, + &mut self.current_group_indices, + hashes, + )?; let group_indices = &self.current_group_indices; // Update ordering information if necessary let total_num_groups = self.group_values.len(); + if let Some(group_hash_tracker) = &mut self.group_hash_tracker { + group_hash_tracker.record_new_groups( + starting_num_groups, + total_num_groups, + group_indices, + hashes, + ); + } if total_num_groups > starting_num_groups { self.group_ordering.new_groups( group_values, @@ -991,7 +1029,12 @@ impl GroupedHashAggregateStream { let groups_and_acc_size = acc + self.group_values.size() + self.group_ordering.size() - + self.current_group_indices.allocated_size(); + + self.current_group_indices.allocated_size() + + self + .group_hash_tracker + .as_ref() + .map_or(0, |v| v.allocated_size()) + + self.hasher.allocated_size(); // Reserve extra headroom for sorting during potential spill. // When OOM triggers, group_aggregate_batch has already processed the @@ -1048,6 +1091,12 @@ impl GroupedHashAggregateStream { output.extend(acc.state(emit_to)?) } } + if let Some(group_hash_tracker) = &mut self.group_hash_tracker { + let group_hashes = group_hash_tracker.emit(emit_to); + if spilling || self.mode.output_mode() == AggregateOutputMode::Partial { + output.push(group_hashes); + } + } drop(timer); // emit reduces the memory usage. Ignore Err from update_memory_reservation. Even if it is @@ -1103,9 +1152,18 @@ impl GroupedHashAggregateStream { cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); let starting_groups = self.group_values.len(); + let hashes = self.hasher.compute_hashes(&cols)?; 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 let Some(group_hash_tracker) = &mut self.group_hash_tracker { + group_hash_tracker.record_new_groups( + starting_groups, + total_groups, + &self.current_group_indices, + hashes, + ); + } if total_groups > starting_groups { self.group_ordering.new_groups( &cols, @@ -1238,6 +1296,10 @@ impl GroupedHashAggregateStream { self.group_values.clear_shrink(num_rows); self.current_group_indices.clear(); self.current_group_indices.shrink_to(num_rows); + if let Some(group_hash_tracker) = &mut self.group_hash_tracker { + group_hash_tracker.clear_shrink(num_rows); + } + self.hasher.clear_shrink(num_rows); } /// Clear memory and shrink capacities to zero. @@ -1289,6 +1351,9 @@ impl GroupedHashAggregateStream { // Mark that we're switching to stream merging mode. self.spill_state.is_stream_merging = true; + self.hasher = self + .hasher + .new_for_exprs(self.spill_state.merging_group_by.input_exprs()); self.input = StreamingMergeBuilder::new() .with_schema(Arc::clone(&self.spill_state.spill_schema)) @@ -1372,7 +1437,7 @@ impl GroupedHashAggregateStream { } /// Transforms input batch to intermediate aggregate state, without grouping it - fn transform_to_states(&self, batch: &RecordBatch) -> Result<RecordBatch> { + fn transform_to_states(&mut self, batch: &RecordBatch) -> Result<RecordBatch> { let mut group_values = evaluate_group_by(&self.group_by, batch)?; let input_values = evaluate_many(&self.aggregate_arguments, batch)?; let filter_values = evaluate_optional(&self.filter_expressions, batch)?; @@ -1382,6 +1447,15 @@ impl GroupedHashAggregateStream { 1, "group_values expected to have single element" ); + let group_hashes = if self.group_hash_tracker.is_some() { + let hashes = match self.hasher.precomputed_group_by(&self.group_by, batch) { + Some(hashes) => hashes, + None => self.hasher.compute_hashes(&group_values[0])?, + }; + Some(Arc::new(UInt64Array::from(hashes.to_vec())) as ArrayRef) + } else { + None + }; let mut output = group_values.swap_remove(0); let iter = self @@ -1394,6 +1468,9 @@ 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)?; diff --git a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs index 97f4662c11..1d512fdf53 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs @@ -23,9 +23,10 @@ use crate::aggregates::topk::priority_map::PriorityMap; use crate::aggregates::topk_types_supported; use crate::aggregates::{ AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, - evaluate_many, + evaluate_many, outputs_group_hashes, }; use crate::metrics::BaselineMetrics; +use crate::repartition::{ExpressionHasher, HashMetrics}; use crate::stream::EmptyRecordBatchStream; use crate::{RecordBatchStream, SendableRecordBatchStream}; use arrow::array::{Array, ArrayRef, RecordBatch, new_null_array}; @@ -55,6 +56,8 @@ pub struct GroupedTopKAggregateStream { aggregate_arguments: Vec<Vec<Arc<dyn PhysicalExpr>>>, group_by: Arc<PhysicalGroupBy>, priority_map: PriorityMap, + output_hasher: ExpressionHasher, + output_group_hashes: bool, /// Whether a NULL group key has been seen for a group-by-only aggregation. null_group_seen: bool, } @@ -109,6 +112,11 @@ impl GroupedTopKAggregateStream { // Note: Null values in aggregate columns are filtered by the aggregation layer // before reaching the heap, so the heap implementations don't need explicit null handling. let priority_map = PriorityMap::new(kt, vt, limit, desc)?; + let output_group_hashes = outputs_group_hashes(aggr.mode, &group_by); + let output_hasher = ExpressionHasher::new_with_metrics( + group_by.output_exprs(), + HashMetrics::new(&aggr.metrics, partition), + ); Ok(GroupedTopKAggregateStream { partition, @@ -122,6 +130,8 @@ impl GroupedTopKAggregateStream { aggregate_arguments, group_by, priority_map, + output_hasher, + output_group_hashes, null_group_seen: false, }) } @@ -175,6 +185,11 @@ impl GroupedTopKAggregateStream { } } + if self.output_group_hashes { + let hashes = self.output_hasher.compute_hashes(&cols[..1])?.to_vec(); + cols.push(Arc::new(arrow::array::UInt64Array::from(hashes))); + } + Ok(cols) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 7ae481e96c..c5bb889940 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -37,6 +37,7 @@ use crate::filter_pushdown::{ FilterPushdownPropagation, PushedDownPredicate, }; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use crate::repartition::ExpressionHasher; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, @@ -52,13 +53,14 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use arrow_schema::FieldRef; use datafusion_common::stats::Precision; +use datafusion_common::utils::split_vec_min_alloc; use datafusion_common::{ Constraint, Constraints, Result, ScalarValue, assert_eq_or_internal_err, internal_err, not_impl_err, }; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryLimit; -use datafusion_expr::{Accumulator, Aggregate}; +use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use datafusion_physical_expr::equivalence::ProjectionMapping; use datafusion_physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, lit}; @@ -71,6 +73,7 @@ use datafusion_physical_expr_common::sort_expr::{ }; use datafusion_expr::utils::AggregateOrderSensitivity; +use datafusion_expr::{Accumulator, Aggregate, EmitTo}; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use itertools::Itertools; use topk::hash_table::is_supported_hash_key_type; @@ -101,10 +104,62 @@ pub fn topk_types_supported(key_type: &DataType, value_type: &DataType) -> bool is_supported_hash_key_type(key_type) && is_supported_heap_type(value_type) } -/// Hard-coded seed for aggregations to ensure hash values differ from `RepartitionExec`, avoiding collisions. -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); +fn outputs_group_hashes(mode: AggregateMode, group_by: &PhysicalGroupBy) -> bool { + mode.output_mode() == AggregateOutputMode::Partial && !group_by.is_true_no_grouping() +} + +/// Tracks hashes for interned groups solely for inclusion in aggregate output. +/// +/// This type does not compute hashes or decide whether an input hash can be +/// reused; that is the responsibility of `ExpressionHasher`. When +/// `GroupValues` interns a new group, this tracker copies the hash used for that +/// group into the corresponding group-id position. It therefore contains one +/// hash per distinct group and remains aligned with `GroupValues` as groups are +/// emitted. +/// +/// Aggregations create this tracker only when partial or spill output must carry +/// an internal hash column for a downstream operator to reuse. +#[derive(Default)] +struct GroupHashTracker { + values: Vec<u64>, +} + +impl GroupHashTracker { + fn record_new_groups( + &mut self, + previous_group_count: usize, + total_group_count: usize, + group_indices: &[usize], + hashes: &[u64], + ) { + assert_eq!(group_indices.len(), hashes.len()); + assert_eq!(self.values.len(), previous_group_count); + self.values.resize(total_group_count, 0); + + for (&group_index, &hash) in group_indices.iter().zip(hashes) { + if group_index >= previous_group_count { + self.values[group_index] = hash; + } + } + } + + fn emit(&mut self, emit_to: EmitTo) -> ArrayRef { + let hashes = match emit_to { + EmitTo::All => std::mem::take(&mut self.values), + EmitTo::First(n) => split_vec_min_alloc(&mut self.values, n), + }; + Arc::new(UInt64Array::from(hashes)) + } + + fn allocated_size(&self) -> usize { + self.values.allocated_size() + } + + fn clear_shrink(&mut self, capacity: usize) { + self.values.clear(); + self.values.shrink_to(capacity); + } +} /// Whether an aggregate stage consumes raw input data or intermediate /// accumulator state from a previous aggregation stage. @@ -2062,7 +2117,11 @@ fn create_schema( aggr_expr: &[Arc<AggregateFunctionExpr>], mode: AggregateMode, ) -> Result<Schema> { - let mut fields = Vec::with_capacity(group_by.num_output_exprs() + aggr_expr.len()); + let mut fields = Vec::with_capacity( + group_by.num_output_exprs() + + aggr_expr.len() + + outputs_group_hashes(mode, group_by) as usize, + ); fields.extend(group_by.output_fields(input_schema)?); match mode.output_mode() { @@ -2080,6 +2139,13 @@ fn create_schema( } } + if outputs_group_hashes(mode, group_by) { + let hasher = ExpressionHasher::new(group_by.output_exprs()); + fields.push( + Field::new(hasher.internal_hash_col_name(), DataType::UInt64, false).into(), + ); + } + Ok(Schema::new_with_metadata( fields, input_schema.metadata().clone(), @@ -2588,6 +2654,7 @@ mod tests { use datafusion_physical_expr::expressions::Literal; use crate::projection::ProjectionExec; + use crate::repartition::{HASH_ROWS_COMPUTED, HASH_ROWS_REUSED, RepartitionExec}; use datafusion_physical_expr::projection::ProjectionExpr; use futures::{FutureExt, Stream, StreamExt}; use insta::{allow_duplicates, assert_snapshot}; @@ -2604,6 +2671,81 @@ mod tests { Ok(schema) } + fn assert_and_strip_group_hashes( + batches: &[RecordBatch], + hash_exprs: &[Arc<dyn PhysicalExpr>], + ) -> Result<Vec<RecordBatch>> { + let mut hasher = ExpressionHasher::new(hash_exprs.to_vec()); + let hash_name = hasher.internal_hash_col_name(); + batches + .iter() + .map(|batch| { + let actual = hasher + .precomputed(batch) + .expect("partial aggregate output should contain group hashes"); + let arrays = evaluate_expressions_to_arrays(hash_exprs, batch)?; + let expected = hasher.compute_hashes(&arrays)?; + assert_eq!(actual, expected); + + let hash_index = batch.schema().index_of(&hash_name)?; + let projection = (0..batch.num_columns()) + .filter(|&index| index != hash_index) + .collect::<Vec<_>>(); + Ok(batch.project(&projection)?) + }) + .collect() + } + + fn batch_with_group_hashes( + schema: SchemaRef, + group_by: &PhysicalGroupBy, + columns: Vec<ArrayRef>, + ) -> Result<RecordBatch> { + let hash_exprs = group_by.output_exprs(); + let mut hasher = ExpressionHasher::new(hash_exprs.clone()); + let hash_name = hasher.internal_hash_col_name(); + let fields = schema + .fields() + .iter() + .filter(|field| field.name() != &hash_name) + .cloned() + .collect::<Vec<_>>(); + let visible_schema = + Arc::new(Schema::new(fields).with_metadata(schema.metadata().clone())); + let batch = RecordBatch::try_new(visible_schema, columns)?; + let arrays = evaluate_expressions_to_arrays(&hash_exprs, &batch)?; + let hashes = hasher.compute_hashes(&arrays)?; + let mut columns = batch.columns().to_vec(); + columns.push(Arc::new(UInt64Array::from(hashes.to_vec()))); + Ok(RecordBatch::try_new(schema, columns)?) + } + + #[test] + fn uses_precomputed_group_hashes_from_batch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let hash_exprs = group_by.input_exprs(); + let hasher = ExpressionHasher::new(hash_exprs); + let hash_schema = Arc::new(Schema::new(vec![ + schema.field(0).clone(), + Field::new(hasher.internal_hash_col_name(), DataType::UInt64, false), + ])); + let batch = RecordBatch::try_new( + hash_schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(UInt64Array::from(vec![11, 22])), + ], + )?; + + let hashes = hasher + .precomputed_group_by(&group_by, &batch) + .expect("batch should contain precomputed group hashes"); + assert_eq!(hashes, &[11, 22]); + Ok(()) + } + /// some mock data to aggregates fn some_data() -> (Arc<Schema>, Vec<RecordBatch>) { // define a schema. @@ -2727,6 +2869,12 @@ mod tests { )) } + fn metric_value(plan: &dyn ExecutionPlan, name: &str) -> usize { + plan.metrics() + .and_then(|metrics| metrics.sum_by_name(name)) + .map_or(0, |value| value.as_usize()) + } + async fn check_grouping_sets( input: Arc<dyn ExecutionPlan>, spill: bool, @@ -2775,6 +2923,8 @@ mod tests { let result = collect(partial_aggregate.execute(0, Arc::clone(&task_ctx))?).await?; + let result = + assert_and_strip_group_hashes(&result, &grouping_set.output_exprs())?; if spill { // In spill mode, we test with the limited memory, if the mem usage exceeds, @@ -2925,6 +3075,8 @@ mod tests { let result = collect(partial_aggregate.execute(0, Arc::clone(&task_ctx))?).await?; + let result = + assert_and_strip_group_hashes(&result, &grouping_set.output_exprs())?; if spill { allow_duplicates! { @@ -3390,6 +3542,125 @@ mod tests { Ok(()) } + #[tokio::test] + async fn partial_repartition_final_preserves_group_hashes() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ])); + let partitions = vec![ + vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 1])), + Arc::new(Int32Array::from(vec![1, 1, 1])), + ], + )?], + vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 3])), + Arc::new(Int32Array::from(vec![1, 1])), + ], + )?], + ]; + let input: Arc<dyn ExecutionPlan> = + TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?; + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let aggregate_exprs = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("COUNT(value)") + .build()?, + )]; + let partial = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggregate_exprs.clone(), + vec![None], + Arc::clone(&input), + Arc::clone(&schema), + )?); + let hash_exprs = group_by.output_exprs(); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::clone(&partial) as Arc<dyn ExecutionPlan>, + Partitioning::Hash(hash_exprs.clone(), 2), + )?); + let task_ctx = new_migrated_hash_ctx(1024); + + let repartition_output = crate::collect_partitioned( + Arc::clone(&repartition) as Arc<dyn ExecutionPlan>, + Arc::clone(&task_ctx), + ) + .await? + .into_iter() + .flatten() + .collect::<Vec<_>>(); + assert_and_strip_group_hashes(&repartition_output, &hash_exprs)?; + assert_eq!(metric_value(partial.as_ref(), HASH_ROWS_COMPUTED), 5); + assert_eq!(metric_value(partial.as_ref(), HASH_ROWS_REUSED), 0); + assert_eq!(metric_value(repartition.as_ref(), HASH_ROWS_COMPUTED), 0); + assert_eq!( + metric_value(repartition.as_ref(), HASH_ROWS_REUSED), + repartition_output + .iter() + .map(RecordBatch::num_rows) + .sum::<usize>() + ); + + // A RepartitionExec can only be consumed once, so build an equivalent + // pipeline for the end-to-end final aggregation assertion. + let partial = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggregate_exprs.clone(), + vec![None], + input, + Arc::clone(&schema), + )?); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::clone(&partial) as Arc<dyn ExecutionPlan>, + Partitioning::Hash(hash_exprs, 2), + )?); + let final_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::FinalPartitioned, + group_by.as_final(), + aggregate_exprs, + vec![None], + Arc::clone(&repartition) as Arc<dyn ExecutionPlan>, + schema, + )?); + let output = crate::collect_partitioned( + Arc::clone(&final_aggregate) as Arc<dyn ExecutionPlan>, + task_ctx, + ) + .await? + .into_iter() + .flatten() + .collect::<Vec<_>>(); + assert_eq!(metric_value(partial.as_ref(), HASH_ROWS_COMPUTED), 5); + assert_eq!(metric_value(partial.as_ref(), HASH_ROWS_REUSED), 0); + assert_eq!(metric_value(repartition.as_ref(), HASH_ROWS_COMPUTED), 0); + assert!(metric_value(repartition.as_ref(), HASH_ROWS_REUSED) > 0); + assert_eq!( + metric_value(final_aggregate.as_ref(), HASH_ROWS_COMPUTED), + 0 + ); + assert!(metric_value(final_aggregate.as_ref(), HASH_ROWS_REUSED) > 0); + assert_snapshot!(batches_to_sort_string(&output), @r" ++-----+--------------+ +| key | COUNT(value) | ++-----+--------------+ +| 1 | 2 | +| 2 | 2 | +| 3 | 1 | ++-----+--------------+ +"); + + Ok(()) + } + #[tokio::test] async fn partial_grouped_aggregate_materializes_before_slicing() -> Result<()> { let schema = Arc::new(Schema::new(vec![ @@ -3407,6 +3678,7 @@ mod tests { TestMemoryExec::try_new_exec(&[input_batches], Arc::clone(&schema), None)?; let group_by = PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let output_hash_exprs = group_by.output_exprs(); let udaf = Arc::new(AggregateUDF::from(NoFirstEmitUdaf::new())); let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new( AggregateExprBuilder::new(udaf, vec![col("value", &schema)?]) @@ -3447,6 +3719,7 @@ mod tests { vec![2, 1] ); assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3); + let batches = assert_and_strip_group_hashes(&batches, &output_hash_exprs)?; assert_snapshot!(batches_to_sort_string(&batches), @r" +-----+-----------------------------+ | key | no_first_emit(value)[count] | @@ -3511,6 +3784,8 @@ mod tests { .sum::<usize>(), 2 ); + let partial_output = + assert_and_strip_group_hashes(&partial_output, &group_by.output_exprs())?; assert_snapshot!(batches_to_sort_string(&partial_output), @r" +---+ | a | @@ -3656,8 +3931,9 @@ mod tests { Arc::clone(&schema), )?; let partial_schema = partial.schema(); - let partial_state_batch = RecordBatch::try_new( + let partial_state_batch = batch_with_group_hashes( Arc::clone(&partial_schema), + &group_by, vec![ Arc::new(UInt32Array::from(vec![1, 2, 1, 3])), Arc::new(Float64Array::from(vec![10.0, 20.0, 40.0, 30.0])), @@ -3763,6 +4039,7 @@ mod tests { (col("sort_col", &schema)?, "sort_col".to_string()), (col("group_col", &schema)?, "group_col".to_string()), ]); + let output_hash_exprs = group_by.output_exprs(); let aggr_expr = vec![Arc::new( AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) .schema(Arc::clone(&schema)) @@ -3788,6 +4065,7 @@ mod tests { let stream: SendableRecordBatchStream = stream.into(); let output = collect(stream).await?; + let output = assert_and_strip_group_hashes(&output, &output_hash_exprs)?; assert_snapshot!(batches_to_sort_string(&output), @r" +----------+-----------+-------------------------+ | sort_col | group_col | COUNT(value_col)[count] | @@ -3834,8 +4112,9 @@ mod tests { Arc::clone(&schema), )?; let partial_schema = partial_aggregate.schema(); - let partial_state_batch = RecordBatch::try_new( + let partial_state_batch = batch_with_group_hashes( Arc::clone(&partial_schema), + &group_by, vec![ Arc::new(Int32Array::from(vec![1, 1, 2, 3])), Arc::new(Int64Array::from(vec![2, 3, 5, 7])), @@ -4645,6 +4924,7 @@ mod tests { let group_by = PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let output_hash_exprs = group_by.output_exprs(); let aggr_expr = vec![ AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?]) @@ -4699,6 +4979,7 @@ mod tests { GroupedHashAggregateStream::new(aggregate_exec.as_ref(), &ctx, 0)?, ); let output = collect(stream).await?; + let output = assert_and_strip_group_hashes(&output, &output_hash_exprs)?; allow_duplicates! { assert_snapshot!(batches_to_string(&output), @r" @@ -4729,6 +5010,7 @@ mod tests { let group_by = PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let output_hash_exprs = group_by.output_exprs(); let aggr_expr = vec![ AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?]) @@ -4791,6 +5073,7 @@ mod tests { GroupedHashAggregateStream::new(aggregate_exec.as_ref(), &ctx, 0)?, ); let output = collect(stream).await?; + let output = assert_and_strip_group_hashes(&output, &output_hash_exprs)?; allow_duplicates! { assert_snapshot!(batches_to_string(&output), @r" @@ -4820,6 +5103,7 @@ mod tests { let group_by = PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let output_hash_exprs = group_by.output_exprs(); let aggr_expr = vec![ AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?]) @@ -4872,6 +5156,7 @@ mod tests { let ctx = Arc::new(TaskContext::default().with_session_config(session_config)); let output = collect(aggregate_exec.execute(0, Arc::clone(&ctx))?).await?; + let output = assert_and_strip_group_hashes(&output, &output_hash_exprs)?; allow_duplicates! { assert_snapshot!(batches_to_sort_string(&output), @r" @@ -4907,6 +5192,7 @@ mod tests { let group_by = PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let output_hash_exprs = group_by.output_exprs(); let aggr_expr = vec![ AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?]) @@ -4967,6 +5253,7 @@ mod tests { let ctx = Arc::new(TaskContext::default().with_session_config(session_config)); let output = collect(aggregate_exec.execute(0, Arc::clone(&ctx))?).await?; + let output = assert_and_strip_group_hashes(&output, &output_hash_exprs)?; allow_duplicates! { assert_snapshot!(batches_to_sort_string(&output), @r" diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 00df227cb8..20fbe18694 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -29,6 +29,7 @@ use crate::execution_plan::{Boundedness, EmissionType, reset_plan_states}; use crate::metrics::{ BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput, }; +use crate::repartition::ExpressionHasher; use crate::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, SendableRecordBatchStream, @@ -43,7 +44,8 @@ use datafusion_common::{ }; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; use futures::{Stream, StreamExt, ready}; @@ -440,10 +442,19 @@ struct DistinctDeduplicator { group_values: Box<dyn GroupValues>, reservation: MemoryReservation, intern_output_buffer: Vec<usize>, + hasher: ExpressionHasher, } impl DistinctDeduplicator { fn new(schema: SchemaRef, task_context: &TaskContext) -> Result<Self> { + let hash_exprs = schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| { + Arc::new(Column::new(field.name(), index)) as Arc<dyn PhysicalExpr> + }) + .collect(); let group_values = new_group_values(schema, &GroupOrdering::None)?; let reservation = MemoryConsumer::new("RecursiveQueryHashTable") .register(task_context.memory_pool()); @@ -451,6 +462,7 @@ impl DistinctDeduplicator { group_values, reservation, intern_output_buffer: Vec::new(), + hasher: ExpressionHasher::new(hash_exprs), }) } @@ -470,8 +482,12 @@ impl DistinctDeduplicator { "failed to reserve {additional} recursive query group ids: {e}" ) })?; - self.group_values - .intern(batch.columns(), &mut self.intern_output_buffer)?; + let hashes = self.hasher.compute_hashes(batch.columns())?; + 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/hash.rs b/datafusion/physical-plan/src/repartition/hash.rs new file mode 100644 index 0000000000..a50c2eb90f --- /dev/null +++ b/datafusion/physical-plan/src/repartition/hash.rs @@ -0,0 +1,177 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::PhysicalGroupBy; +use crate::joins::SeededRandomState; +use crate::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory}; +use arrow::array::{ArrayRef, RecordBatch}; +use datafusion_common::Result; +use datafusion_common::cast::as_uint64_array; +use datafusion_common::hash_utils::create_hashes; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::VecAllocExt; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +const INTERNAL_HASH_COL_PREFIX: &str = "__datafusion_internal_hash"; +pub(crate) const HASH_ROWS_COMPUTED: &str = "hash_rows_computed"; +pub(crate) const HASH_ROWS_REUSED: &str = "hash_rows_reused"; + +/// Execution metrics for expression hashing. +/// +/// Clones update the same underlying counters, allowing a replacement hasher +/// (for example while merging spilled aggregation state) to preserve the +/// operator's metrics. +#[derive(Debug, Clone)] +pub(crate) struct HashMetrics { + rows_computed: Count, + rows_reused: Count, +} + +impl HashMetrics { + pub(crate) fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + Self { + rows_computed: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter(HASH_ROWS_COMPUTED, partition), + rows_reused: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter(HASH_ROWS_REUSED, partition), + } + } + + fn record_computed(&self, rows: usize) { + self.rows_computed.add(rows); + } + + fn record_reused(&self, rows: usize) { + self.rows_reused.add(rows); + } +} + +pub(crate) struct ExpressionHasher { + hash_exprs: Vec<Arc<dyn PhysicalExpr>>, + hash_buffer: Vec<u64>, + random_state: SeededRandomState, + metrics: Option<HashMetrics>, +} + +impl ExpressionHasher { + pub(crate) fn new(hash_exprs: Vec<Arc<dyn PhysicalExpr>>) -> Self { + Self { + hash_exprs, + hash_buffer: vec![], + random_state: SeededRandomState::with_seed(0), + metrics: None, + } + } + + pub(crate) fn new_with_metrics( + hash_exprs: Vec<Arc<dyn PhysicalExpr>>, + metrics: HashMetrics, + ) -> Self { + Self { + metrics: Some(metrics), + ..Self::new(hash_exprs) + } + } + + pub(crate) fn set_metrics(&mut self, metrics: HashMetrics) { + self.metrics = Some(metrics); + } + + /// Creates a hasher for another expression list while preserving metrics. + pub(crate) fn new_for_exprs(&self, hash_exprs: Vec<Arc<dyn PhysicalExpr>>) -> Self { + Self { + metrics: self.metrics.clone(), + ..Self::new(hash_exprs) + } + } + + /// Builds the name for the column that will carry hashes across [`ExecutionPlan`]s. + /// + /// In order to avoid recomputation of hashes, some nodes have the capability of computing the + /// hash once, and propagate it through the plan so that future nodes can reuse them. + /// Equal expression lists produce the same name, while the expression order is part of the + /// identity. The name is an identifier within a physical plan and is not stable across versions. + /// + /// [`ExecutionPlan`]: crate::ExecutionPlan + pub(crate) fn internal_hash_col_name(&self) -> String { + let mut hasher = DefaultHasher::new(); + self.hash_exprs.hash(&mut hasher); + format!("{INTERNAL_HASH_COL_PREFIX}_{:016x}", hasher.finish()) + } + + pub(crate) fn precomputed<'a>(&self, batch: &'a RecordBatch) -> Option<&'a [u64]> { + let internal_hash_col_name = self.internal_hash_col_name(); + let hash_column = batch.column_by_name(&internal_hash_col_name)?; + let hash_array = as_uint64_array(hash_column.as_ref()).ok()?; + if let Some(metrics) = &self.metrics { + metrics.record_reused(hash_array.len()); + } + Some(hash_array.values()) + } + + pub(crate) fn precomputed_group_by<'a>( + &self, + group_by: &PhysicalGroupBy, + batch: &'a RecordBatch, + ) -> Option<&'a [u64]> { + if !group_by.is_single() { + return None; + } + + self.precomputed(batch) + } + + pub(crate) fn compute_hashes(&mut self, arrays: &[ArrayRef]) -> Result<&[u64]> { + let num_rows = arrays.first().map(|array| array.len()).unwrap_or(0); + self.hash_buffer.clear(); + self.hash_buffer.resize(num_rows, 0); + + create_hashes( + arrays, + self.random_state.random_state(), + &mut self.hash_buffer, + )?; + + if let Some(metrics) = &self.metrics { + metrics.record_computed(num_rows); + } + + Ok(&self.hash_buffer) + } + + pub(crate) fn compute_hashes_for_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<&[u64]> { + let arrays = evaluate_expressions_to_arrays(&self.hash_exprs, batch)?; + self.compute_hashes(&arrays) + } + + pub(crate) fn allocated_size(&self) -> usize { + self.hash_exprs.allocated_size() + self.hash_buffer.allocated_size() + } + + pub(crate) fn clear_shrink(&mut self, capacity: usize) { + self.hash_buffer.clear(); + self.hash_buffer.shrink_to(capacity); + } +} diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 6587157421..5b1f9256b8 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -34,7 +34,6 @@ use super::{ }; use crate::coalesce::LimitedBatchCoalescer; use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; -use crate::hash_utils::create_hashes; use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::projection::{ProjectionExec, all_columns, make_with_child, update_expr}; use crate::sorts::streaming_merge::StreamingMergeBuilder; @@ -79,7 +78,12 @@ use log::trace; use parking_lot::Mutex; mod distributor_channels; +mod hash; + use crate::repartition::distributor_channels::SendError; +pub(crate) use crate::repartition::hash::{ExpressionHasher, HashMetrics}; +#[cfg(test)] +pub(crate) use crate::repartition::hash::{HASH_ROWS_COMPUTED, HASH_ROWS_REUSED}; use distributor_channels::{ DistributionReceiver, DistributionSender, channels, partition_aware_channels, }; @@ -366,7 +370,7 @@ impl RepartitionExecState { &mut self, input: &Arc<dyn ExecutionPlan>, metrics: &ExecutionPlanMetricsSet, - output_partitions: usize, + partitioning: &Partitioning, ctx: &Arc<TaskContext>, ) -> Result<()> { if !matches!(self, RepartitionExecState::NotInitialized) { @@ -374,10 +378,16 @@ impl RepartitionExecState { } let num_input_partitions = input.output_partitioning().partition_count(); + let output_partitions = partitioning.partition_count(); let mut streams_and_metrics = Vec::with_capacity(num_input_partitions); for i in 0..num_input_partitions { - let metrics = RepartitionMetrics::new(i, output_partitions, metrics); + let metrics = RepartitionMetrics::new( + i, + output_partitions, + matches!(partitioning, Partitioning::Hash(..)), + metrics, + ); let timer = metrics.fetch_time.timer(); let stream = input.execute(i, Arc::clone(ctx))?; @@ -405,7 +415,7 @@ impl RepartitionExecState { self.ensure_input_streams_initialized( input, metrics, - partitioning.partition_count(), + partitioning, context, )?; let RepartitionExecState::InputStreamsInitialized(value) = self else { @@ -564,9 +574,8 @@ pub struct BatchPartitioner { enum BatchPartitionerState { Hash { - exprs: Vec<Arc<dyn PhysicalExpr>>, partition_reducer: StrengthReducedU64, - hash_buffer: Vec<u64>, + hasher: ExpressionHasher, indices: Vec<Vec<u32>>, }, RoundRobin { @@ -664,6 +673,12 @@ impl StrengthReducedU64 { } impl BatchPartitioner { + fn set_hash_metrics(&mut self, metrics: HashMetrics) { + if let BatchPartitionerState::Hash { hasher, .. } = &mut self.state { + hasher.set_metrics(metrics); + } + } + /// Create a new [`BatchPartitioner`] for hash-based repartitioning. /// /// # Parameters @@ -687,9 +702,8 @@ impl BatchPartitioner { Ok(Self { state: BatchPartitionerState::Hash { - exprs, partition_reducer: StrengthReducedU64::new(num_partitions as u64), - hash_buffer: vec![], + hasher: ExpressionHasher::new(exprs), indices: vec![vec![]; num_partitions], }, timer, @@ -837,28 +851,20 @@ impl BatchPartitioner { Box::new(std::iter::once(Ok((idx, batch)))) } BatchPartitionerState::Hash { - exprs, partition_reducer, - hash_buffer, + hasher, indices, } => { // 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, - )?; - indices.iter_mut().for_each(|v| v.clear()); + let hash_buffer = match hasher.precomputed(&batch) { + Some(hashes) => hashes, + None => hasher.compute_hashes_for_batch(&batch)?, + }; + partition_reducer.partition_indices(hash_buffer, indices); // Finished building index-arrays for output partitions @@ -1172,12 +1178,15 @@ struct RepartitionMetrics { /// /// One metric per output partition. send_time: Vec<metrics::Time>, + /// Rows whose hash values were computed or reused from an internal column. + hash_metrics: Option<HashMetrics>, } impl RepartitionMetrics { pub fn new( input_partition: usize, num_output_partitions: usize, + track_hashes: bool, metrics: &ExecutionPlanMetricsSet, ) -> Self { // Time in nanos to execute child operator and fetch batches @@ -1203,6 +1212,8 @@ impl RepartitionMetrics { fetch_time, repartition_time, send_time, + hash_metrics: track_hashes + .then(|| HashMetrics::new(metrics, input_partition)), } } } @@ -1361,7 +1372,7 @@ impl ExecutionPlan for RepartitionExec { state.ensure_input_streams_initialized( &input, &metrics, - partitioning.partition_count(), + &partitioning, &context, )?; } @@ -1758,6 +1769,9 @@ impl RepartitionExec { input_partition, num_input_partitions, )?; + if let Some(hash_metrics) = &metrics.hash_metrics { + partitioner.set_hash_metrics(hash_metrics.clone()); + } // While there are still outputs to send to, keep pulling inputs let mut batches_until_yield = partitioner.num_partitions(); diff --git a/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt b/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt index e879947e32..552c4dc3b8 100644 --- a/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt +++ b/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt @@ -742,9 +742,9 @@ Plan with Metrics 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, a@0)], projection=[a@0, min_value@2], metrics=[output_rows=2, output_batches=2, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=2, input_rows=2, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/join_agg_build.parquet]]}, projection=[a], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_ [...] 03)--ProjectionExec: expr=[a@0 as a, min(join_agg_probe.value)@1 as min_value], metrics=[output_rows=2, output_batches=2] -04)----AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[min(join_agg_probe.value)], metrics=[output_rows=2, output_batches=2, spill_count=0, spilled_rows=0] -05)------RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=1, metrics=[output_rows=2, output_batches=2, spill_count=0, spilled_rows=0] -06)--------AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[min(join_agg_probe.value)], metrics=[output_rows=2, output_batches=1, spill_count=0, spilled_rows=0, skipped_aggregation_rows=0, reduction_factor=100% (2/2)] +04)----AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[min(join_agg_probe.value)], metrics=[output_rows=2, output_batches=2, spill_count=0, spilled_rows=0, hash_rows_computed=0, hash_rows_reused=2] +05)------RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=1, metrics=[output_rows=2, output_batches=2, spill_count=0, spilled_rows=0, hash_rows_computed=0, hash_rows_reused=2] +06)--------AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[min(join_agg_probe.value)], metrics=[output_rows=2, output_batches=1, spill_count=0, spilled_rows=0, hash_rows_computed=2, hash_rows_reused=0, skipped_aggregation_rows=0, reduction_factor=100% (2/2)] 07)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/join_agg_probe.parquet]]}, projection=[a, value], file_type=parquet, predicate=DynamicFilter [ a@0 >= h1 AND a@0 <= h2 AND a@0 IN (SET) ([h1, h2]) ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= h1 AND a_null_count@1 != row_count@2 AND a_min@3 <= h2 AND (a_null_count@1 != row_count@2 AND a_min@3 <= h1 AND h1 [...] statement ok --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
