From 8b687239c1c553e17f62b1af011709b1a7d97c2a Mon Sep 17 00:00:00 2001 From: Gabriel Date: Sun, 19 Jul 2026 09:53:51 +0200 Subject: [PATCH 1/6] Add name filter to metrics --- .../physical-expr-common/src/metrics/mod.rs | 31 ++++++++++++ datafusion/physical-plan/src/display.rs | 50 +++++++++++++++++++ 2 files changed, 81 insertions(+) diff --git a/datafusion/physical-expr-common/src/metrics/mod.rs b/datafusion/physical-expr-common/src/metrics/mod.rs index d6048a0fcd338..f510714dae2a4 100644 --- a/datafusion/physical-expr-common/src/metrics/mod.rs +++ b/datafusion/physical-expr-common/src/metrics/mod.rs @@ -418,6 +418,21 @@ impl MetricsSet { .collect::>(); Self { metrics } } + + /// Returns a new `MetricsSet` filtered by metric name. + /// Only metrics with the names appearing the list will be kept. + pub fn filter_by_names(self, names: &[String]) -> Self { + if names.is_empty() { + return Self { metrics: vec![] }; + } + + let metrics = self + .metrics + .into_iter() + .filter(|metric| names.iter().any(|name| name == metric.value().name())) + .collect::>(); + Self { metrics } + } } impl Display for MetricsSet { @@ -966,4 +981,20 @@ mod tests { metric_names(&metrics) ); } + + #[test] + fn test_filter_by_names() { + let metrics = ExecutionPlanMetricsSet::new(); + MetricBuilder::new(&metrics).output_rows(0); + MetricBuilder::new(&metrics).counter("custom_counter", 0); + + let names = vec!["output_rows".to_string()]; + let filtered = metrics.clone_inner().filter_by_names(&names); + + assert_eq!(filtered.iter().count(), 1); + assert_eq!( + filtered.iter().next().unwrap().value().name(), + "output_rows" + ); + } } diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 6a4d09057bec9..d14be3575a361 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -129,6 +129,9 @@ pub struct DisplayableExecutionPlan<'a> { /// Optional filter by semantic category (rows / bytes / timing). /// `None` means show all categories; `Some(vec![])` means plan-only. metric_categories: Option>, + /// Optional filter by metric names. Only metric names in this list + /// will be rendered. + metric_names: Option>, // (TreeRender) Maximum total width of the rendered tree tree_maximum_render_width: usize, /// Optional summary totals (currently only used by `pgjson`) — the total @@ -159,6 +162,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: false, metric_types: Self::default_metric_types(), metric_categories: None, + metric_names: None, tree_maximum_render_width: 240, summary: None, } @@ -175,6 +179,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: false, metric_types: Self::default_metric_types(), metric_categories: None, + metric_names: None, tree_maximum_render_width: 240, summary: None, } @@ -191,6 +196,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: false, metric_types: Self::default_metric_types(), metric_categories: None, + metric_names: None, tree_maximum_render_width: 240, summary: None, } @@ -234,6 +240,15 @@ impl<'a> DisplayableExecutionPlan<'a> { self } + /// Specify which metric names to include. + /// + /// - An empty vector means plan-only — suppress all metrics. + /// - `vec!["metric_1"]` means show only the metric named `metric_1`. + pub fn set_metric_names(mut self, metric_names: Vec) -> Self { + self.metric_names = Some(metric_names); + self + } + /// Set the maximum render width for the tree format pub fn set_tree_maximum_render_width(mut self, width: usize) -> Self { self.tree_maximum_render_width = width; @@ -279,6 +294,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: bool, metric_types: Vec, metric_categories: Option>, + metric_names: Option>, } impl fmt::Display for Wrapper<'_> { fn fmt(&self, f: &mut Formatter) -> fmt::Result { @@ -291,6 +307,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), }; accept(self.plan, &mut visitor) } @@ -303,6 +320,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), } } @@ -324,6 +342,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_statistics: bool, metric_types: Vec, metric_categories: Option>, + metric_names: Option>, } impl fmt::Display for Wrapper<'_> { fn fmt(&self, f: &mut Formatter) -> fmt::Result { @@ -336,6 +355,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_statistics: self.show_statistics, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), graphviz_builder: GraphvizBuilder::default(), parents: Vec::new(), }; @@ -355,6 +375,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_statistics: self.show_statistics, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), } } @@ -403,6 +424,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: bool, metric_types: Vec, metric_categories: Option>, + metric_names: Option>, summary: Option, } impl fmt::Display for Wrapper<'_> { @@ -413,6 +435,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), objects: HashMap::new(), parent_ids: Vec::new(), next_id: 0, @@ -446,6 +469,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), summary: self.summary, } } @@ -460,6 +484,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: bool, metric_types: Vec, metric_categories: Option>, + metric_names: Option>, } impl fmt::Display for Wrapper<'_> { @@ -473,6 +498,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), }; visitor.pre_visit(self.plan)?; Ok(()) @@ -486,6 +512,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), } } @@ -544,6 +571,8 @@ struct IndentVisitor<'a, 'b> { metric_types: &'a [MetricType], /// Optional filter by semantic category (rows / bytes / timing). metric_categories: Option<&'a [MetricCategory]>, + /// Optional filter by metric name. + metric_names: Option<&'a [String]>, } impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { @@ -563,6 +592,9 @@ impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { if let Some(cats) = self.metric_categories { metrics = metrics.filter_by_categories(cats); } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } write!(self.f, ", metrics=[{metrics}]")?; } else { write!(self.f, ", metrics=[]")?; @@ -574,6 +606,9 @@ impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { if let Some(cats) = self.metric_categories { metrics = metrics.filter_by_categories(cats); } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } write!(self.f, ", metrics=[{metrics}]")?; } else { write!(self.f, ", metrics=[]")?; @@ -616,6 +651,8 @@ struct GraphvizVisitor<'a, 'b> { metric_types: &'a [MetricType], /// Optional filter by semantic category metric_categories: Option<&'a [MetricCategory]>, + /// Optional filter by metric name. + metric_names: Option<&'a [String]>, graphviz_builder: GraphvizBuilder, /// Used to record parent node ids when visiting a plan. @@ -660,6 +697,9 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> { if let Some(cats) = self.metric_categories { metrics = metrics.filter_by_categories(cats); } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } format!("metrics=[{metrics}]") } else { "metrics=[]".to_string() @@ -671,6 +711,9 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> { if let Some(cats) = self.metric_categories { metrics = metrics.filter_by_categories(cats); } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } format!("metrics=[{metrics}]") } else { "metrics=[]".to_string() @@ -729,6 +772,7 @@ struct PgJsonExecutionPlanVisitor<'a> { show_schema: bool, metric_types: &'a [MetricType], metric_categories: Option<&'a [MetricCategory]>, + metric_names: Option<&'a [String]>, objects: HashMap, parent_ids: Vec, next_id: u32, @@ -813,6 +857,12 @@ impl PgJsonExecutionPlanVisitor<'_> { metrics }; + let metrics = if let Some(names) = self.metric_names { + metrics.filter_by_names(names) + } else { + metrics + }; + // Build the Extras bucket, while extracting PG-canonical keys to the // top level. let mut extras = serde_json::Map::new(); From adaf98e2e2f26025d7d80a713871925e5bdc00d5 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Sat, 18 Jul 2026 18:07:31 +0200 Subject: [PATCH 2/6] Reuse hashes in aggregations --- datafusion/core/tests/dataframe/mod.rs | 4 + .../partition_statistics.rs | 20 +- .../physical-expr-common/src/binary_map.rs | 45 ++- .../src/binary_view_map.rs | 45 ++- .../benches/dictionary_group_values.rs | 27 +- .../physical-plan/benches/multi_group_by.rs | 27 +- .../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 +- .../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 +- .../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 +- .../physical-plan/src/aggregates/mod.rs | 303 +++++++++++++++++- .../physical-plan/src/recursive_query.rs | 22 +- .../physical-plan/src/repartition/hash.rs | 177 ++++++++++ .../physical-plan/src/repartition/mod.rs | 60 ++-- .../test_files/push_down_filter_parquet.slt | 6 +- 23 files changed, 971 insertions(+), 177 deletions(-) create mode 100644 datafusion/physical-plan/src/repartition/hash.rs diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index b9fecb5fdd732..10ecc2e683fa2 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -7183,6 +7183,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 6cabcdb710393..2aaa3032685ce 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 ad184d6500d56..77bce64eed95e 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( + &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::>( values, + hashes, make_payload_fn, observe_payload_fn, ) @@ -319,6 +348,7 @@ where )); self.insert_if_new_inner::>( values, + hashes, make_payload_fn, observe_payload_fn, ) @@ -338,6 +368,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 +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::(); // 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 9d4b556393a24..34b4df8b25200 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( + &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::( 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::( values, + hashes, make_payload_fn, observe_payload_fn, ) @@ -245,6 +275,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 +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::(); // 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..756860da9fe18 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) { + 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::::with_capacity(size), + Vec::::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::::with_capacity(size), + Vec::::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 11c2800864316..388cb4be9d564 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -32,6 +32,7 @@ use arrow::compute::take; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::util::bench_util::create_fsb_array; 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; @@ -39,6 +40,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 = (0..num_cols) @@ -98,10 +100,15 @@ fn bench_intern( gv: &mut Box, batches: &[Vec], groups: &mut Vec, + hashes: &mut Vec, ) { + 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); } @@ -135,9 +142,10 @@ fn bench_issue_17850_regression(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(DEFAULT_BATCH_SIZE), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), ) }, - |(gv, groups)| bench_intern(gv, batches, groups), + |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), criterion::BatchSize::LargeInput, ); }, @@ -178,9 +186,10 @@ fn bench_low_cardinality(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(DEFAULT_BATCH_SIZE), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), ) }, - |(gv, groups)| bench_intern(gv, batches, groups), + |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), criterion::BatchSize::LargeInput, ); }, @@ -217,9 +226,10 @@ fn bench_batch_size_sensitivity(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(batch_size), + Vec::::with_capacity(batch_size), ) }, - |(gv, groups)| bench_intern(gv, batches, groups), + |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), criterion::BatchSize::LargeInput, ); }, @@ -257,9 +267,10 @@ fn bench_column_scaling(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(DEFAULT_BATCH_SIZE), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), ) }, - |(gv, groups)| bench_intern(gv, batches, groups), + |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), criterion::BatchSize::LargeInput, ); }, @@ -295,9 +306,10 @@ fn bench_high_cardinality_scaling(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(DEFAULT_BATCH_SIZE), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), ) }, - |(gv, groups)| bench_intern(gv, batches, groups), + |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), criterion::BatchSize::LargeInput, ); }, @@ -336,9 +348,10 @@ fn bench_group_count_sweep(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(DEFAULT_BATCH_SIZE), + Vec::::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 eaf39929ced62..71dfb1571d141 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 AggregateHashTable { 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 AggregateHashTable { } /// See comments in [`EvaluatedAggregateBatch`] - pub(super) fn evaluate_batch( + pub(super) fn evaluate_batch<'a>( &self, - batch: &RecordBatch, - ) -> Result { + batch: &'a RecordBatch, + ) -> Result> { 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 AggregateHashTable { Ok(EvaluatedAggregateBatch { grouping_set_args, accumulator_args, + precomputed_group_hashes: state + .hasher + .precomputed_group_by(&state.group_by, batch), }) } @@ -189,9 +200,24 @@ 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 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 AggregateHashTable { 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 AggregateHashTable { 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>, /// Evaluated arguments and filters, one entry per aggregate expression. pub(super) accumulator_args: Vec, + + /// 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, + /// Hashes retained for the internal hash column when output requires it. + pub(super) group_hash_tracker: Option, + + /// 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 c83303c51d6e8..f25c0f9bb0ff8 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, + /// Hashes retained for the internal hash column when output requires it. + pub(super) group_hash_tracker: Option, + + /// 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 OrderedAggregateTable { 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 OrderedAggregateTable { /// /// 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 { + batch: &'a RecordBatch, + ) -> Result> { 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 OrderedAggregateTable { 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 OrderedAggregateTable { + 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 OrderedAggregateTable { /// - `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 OrderedAggregateTable { 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 ffac42feaa3b3..41949dc67ab0a 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 { 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 { .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 { 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 { { 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 ee253e5d7afdd..d6024eea4a1a2 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) -> 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..7e99ee184c326 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,15 @@ impl GroupValuesColumn { &mut self, cols: &[ArrayRef], groups: &mut Vec, + 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 GroupValuesColumn { &mut self, cols: &[ArrayRef], groups: &mut Vec, + 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 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 +475,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 +1059,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 +1210,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 { @@ -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 4976a098ecee5..0be654bdbf05c 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. @@ -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 e993c0c53d199..ac94c8d4acb7b 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) -> Result<()> { + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec, + 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 b881a51b25474..71bcb98509699 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,10 @@ impl GroupValues for GroupValuesBytes { 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 7a56f7c52c11a..e4a2397b06aae 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,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 e254aebcfd7ce..b8f1356d95bb7 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,17 @@ 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); + assert_eq!(cols[0].len(), hashes.len()); groups.clear(); - for v in cols[0].as_primitive::() { + for (v, &hash) in cols[0].as_primitive::().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 0d00e5c4d0d86..ea365bdb14f68 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, + /// Hashes retained for the internal hash column when output requires it. + group_hash_tracker: Option, + + /// 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 { + fn transform_to_states(&mut self, batch: &RecordBatch) -> Result { 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 97f4662c11342..1d512fdf53b5b 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>>, group_by: Arc, 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 732da32ab0391..73c9395e31a80 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -162,6 +162,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, @@ -177,13 +178,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}; @@ -196,6 +198,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; @@ -226,10 +229,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, +} + +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. @@ -2193,7 +2248,11 @@ fn create_schema( aggr_expr: &[Arc], mode: AggregateMode, ) -> Result { - 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() { @@ -2211,6 +2270,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(), @@ -2719,6 +2785,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}; @@ -2735,6 +2802,81 @@ mod tests { Ok(schema) } + fn assert_and_strip_group_hashes( + batches: &[RecordBatch], + hash_exprs: &[Arc], + ) -> Result> { + 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::>(); + Ok(batch.project(&projection)?) + }) + .collect() + } + + fn batch_with_group_hashes( + schema: SchemaRef, + group_by: &PhysicalGroupBy, + columns: Vec, + ) -> Result { + 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::>(); + 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, Vec) { // define a schema. @@ -2858,6 +3000,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, spill: bool, @@ -2906,6 +3054,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, @@ -3056,6 +3206,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! { @@ -3521,6 +3673,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 = + 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, + 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, + Arc::clone(&task_ctx), + ) + .await? + .into_iter() + .flatten() + .collect::>(); + 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::() + ); + + // 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, + 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, + schema, + )?); + let output = crate::collect_partitioned( + Arc::clone(&final_aggregate) as Arc, + task_ctx, + ) + .await? + .into_iter() + .flatten() + .collect::>(); + 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![ @@ -3538,6 +3809,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> = vec![Arc::new( AggregateExprBuilder::new(udaf, vec![col("value", &schema)?]) @@ -3578,6 +3850,7 @@ mod tests { vec![2, 1] ); assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 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] | @@ -3642,6 +3915,8 @@ mod tests { .sum::(), 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 | @@ -3787,8 +4062,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])), @@ -3894,6 +4170,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)) @@ -3919,6 +4196,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] | @@ -3965,8 +4243,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])), @@ -4776,6 +5055,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)?]) @@ -4830,6 +5110,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" @@ -4860,6 +5141,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)?]) @@ -4922,6 +5204,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" @@ -4951,6 +5234,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)?]) @@ -5003,6 +5287,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" @@ -5038,6 +5323,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)?]) @@ -5098,6 +5384,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 00df227cb87db..20fbe1869460b 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, reservation: MemoryReservation, intern_output_buffer: Vec, + hasher: ExpressionHasher, } impl DistinctDeduplicator { fn new(schema: SchemaRef, task_context: &TaskContext) -> Result { + let hash_exprs = schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| { + Arc::new(Column::new(field.name(), index)) as Arc + }) + .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 0000000000000..a50c2eb90f415 --- /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>, + hash_buffer: Vec, + random_state: SeededRandomState, + metrics: Option, +} + +impl ExpressionHasher { + pub(crate) fn new(hash_exprs: Vec>) -> Self { + Self { + hash_exprs, + hash_buffer: vec![], + random_state: SeededRandomState::with_seed(0), + metrics: None, + } + } + + pub(crate) fn new_with_metrics( + hash_exprs: Vec>, + 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>) -> 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 12229c26b7d98..fe3300ac551e8 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, }; @@ -398,7 +402,7 @@ impl RepartitionExecState { &mut self, input: &Arc, metrics: &ExecutionPlanMetricsSet, - output_partitions: usize, + partitioning: &Partitioning, ctx: &Arc, ) -> Result<()> { if !matches!(self, RepartitionExecState::NotInitialized) { @@ -406,10 +410,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))?; @@ -437,7 +447,7 @@ impl RepartitionExecState { self.ensure_input_streams_initialized( input, metrics, - partitioning.partition_count(), + partitioning, context, )?; let RepartitionExecState::InputStreamsInitialized(value) = self else { @@ -610,9 +620,8 @@ pub struct BatchPartitioner { enum BatchPartitionerState { Hash { - exprs: Vec>, partition_reducer: StrengthReducedU64, - hash_buffer: Vec, + hasher: ExpressionHasher, indices: Vec>, }, RoundRobin { @@ -710,6 +719,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 @@ -733,9 +748,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, @@ -883,28 +897,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 @@ -1218,12 +1224,15 @@ struct RepartitionMetrics { /// /// One metric per output partition. send_time: Vec, + /// Rows whose hash values were computed or reused from an internal column. + hash_metrics: Option, } 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 @@ -1249,6 +1258,8 @@ impl RepartitionMetrics { fetch_time, repartition_time, send_time, + hash_metrics: track_hashes + .then(|| HashMetrics::new(metrics, input_partition)), } } } @@ -1407,7 +1418,7 @@ impl ExecutionPlan for RepartitionExec { state.ensure_input_streams_initialized( &input, &metrics, - partitioning.partition_count(), + &partitioning, &context, )?; } @@ -1812,6 +1823,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 e879947e324bb..552c4dc3b8179 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_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=14.45% (64/443)] 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 <= a_max@0 OR a_null_count@1 != row_count@2 AND a_min@3 <= h2 AND h2 <= a_max@0), required_guarantees=[a in (h1, h2)], 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=1 total → 1 matched, page_index_rows_pruned=4 total → 4 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=4, predicate_cache_records=2, scan_efficiency_ratio=19.07% (151/792)] statement ok From 4cde15b1049251730a355ed0f4f5a794da28bdc4 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Sun, 19 Jul 2026 10:18:49 +0200 Subject: [PATCH 3/6] Add hash_reuse integration test --- datafusion/core/tests/execution/hash_reuse.rs | 254 ++++++++++++++++++ datafusion/core/tests/execution/mod.rs | 1 + .../aggregates/aggregate_hash_table/common.rs | 2 + 3 files changed, 257 insertions(+) create mode 100644 datafusion/core/tests/execution/hash_reuse.rs diff --git a/datafusion/core/tests/execution/hash_reuse.rs b/datafusion/core/tests/execution/hash_reuse.rs new file mode 100644 index 0000000000000..d2e2c01f4ade0 --- /dev/null +++ b/datafusion/core/tests/execution/hash_reuse.rs @@ -0,0 +1,254 @@ +// 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 datafusion::error::Result; +use datafusion::physical_plan::collect; +use datafusion::physical_plan::display::DisplayableExecutionPlan; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use insta::assert_snapshot; + +#[tokio::test] +async fn grouped_tpch_aggregate_reuses_hashes_after_repartition() -> Result<()> { + let plan_with_hash_metrics = HashReuseTest::new( + r"SELECT l_returnflag, l_linestatus, COUNT(*) AS row_count + FROM lineitem + GROUP BY l_returnflag, l_linestatus", + ) + .run() + .await?; + + assert_snapshot!(plan_with_hash_metrics, @r" + ProjectionExec: expr=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus, count(Int64(1))@2 as row_count], metrics=[] + AggregateExec: mode=FinalPartitioned, gby=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus], aggr=[count(Int64(1))], metrics=[hash_rows_computed=0, hash_rows_reused=3] + RepartitionExec: partitioning=Hash([l_returnflag@0, l_linestatus@1], 3), input_partitions=1, metrics=[hash_rows_computed=0, hash_rows_reused=3] + AggregateExec: mode=Partial, gby=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus], aggr=[count(Int64(1))], metrics=[hash_rows_computed=20, hash_rows_reused=0] + DataSourceExec: file_groups={1 group: [[$DATAFUSION_CORE/tests/data/tpch_lineitem_small.parquet]]}, projection=[l_returnflag, l_linestatus], file_type=parquet, metrics=[] + "); + + Ok(()) +} + +#[tokio::test] +async fn legacy_grouped_tpch_aggregate_reuses_hashes_after_repartition() -> Result<()> { + let plan_with_hash_metrics = HashReuseTest::new( + r"SELECT l_returnflag, l_linestatus, COUNT(*) AS row_count + FROM lineitem + GROUP BY l_returnflag, l_linestatus", + ) + .with_migration_aggregate(false) + .run() + .await?; + + assert_snapshot!(plan_with_hash_metrics, @r" + ProjectionExec: expr=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus, count(Int64(1))@2 as row_count], metrics=[] + AggregateExec: mode=FinalPartitioned, gby=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus], aggr=[count(Int64(1))], metrics=[hash_rows_computed=0, hash_rows_reused=3] + RepartitionExec: partitioning=Hash([l_returnflag@0, l_linestatus@1], 3), input_partitions=1, metrics=[hash_rows_computed=0, hash_rows_reused=3] + AggregateExec: mode=Partial, gby=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus], aggr=[count(Int64(1))], metrics=[hash_rows_computed=20, hash_rows_reused=0] + DataSourceExec: file_groups={1 group: [[$DATAFUSION_CORE/tests/data/tpch_lineitem_small.parquet]]}, projection=[l_returnflag, l_linestatus], file_type=parquet, metrics=[] + "); + + Ok(()) +} + +#[tokio::test] +async fn chunked_partial_tpch_aggregate_reuses_hashes_after_repartition() -> Result<()> { + let plan_with_hash_metrics = HashReuseTest::new( + r"SELECT l_returnflag, l_linestatus, COUNT(*) AS row_count + FROM lineitem + GROUP BY l_returnflag, l_linestatus", + ) + .with_batch_size(1) + .run() + .await?; + + assert_snapshot!(plan_with_hash_metrics, @r" + ProjectionExec: expr=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus, count(Int64(1))@2 as row_count], metrics=[] + AggregateExec: mode=FinalPartitioned, gby=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus], aggr=[count(Int64(1))], metrics=[hash_rows_computed=0, hash_rows_reused=9] + RepartitionExec: partitioning=Hash([l_returnflag@0, l_linestatus@1], 3), input_partitions=3, metrics=[hash_rows_computed=0, hash_rows_reused=9] + AggregateExec: mode=Partial, gby=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus], aggr=[count(Int64(1))], metrics=[hash_rows_computed=20, hash_rows_reused=0] + RepartitionExec: partitioning=RoundRobinBatch(3), input_partitions=1, metrics=[] + DataSourceExec: file_groups={1 group: [[$DATAFUSION_CORE/tests/data/tpch_lineitem_small.parquet]]}, projection=[l_returnflag, l_linestatus], file_type=parquet, metrics=[] + "); + + Ok(()) +} + +#[tokio::test] +async fn expression_grouped_tpch_aggregate_reuses_hashes_after_repartition() -> Result<()> +{ + let plan_with_hash_metrics = HashReuseTest::new( + r"SELECT upper(l_returnflag) AS returnflag, COUNT(*) AS row_count + FROM lineitem + GROUP BY upper(l_returnflag)", + ) + .run() + .await?; + + assert_snapshot!(plan_with_hash_metrics, @r" + ProjectionExec: expr=[upper(lineitem.l_returnflag)@0 as returnflag, count(Int64(1))@1 as row_count], metrics=[] + AggregateExec: mode=FinalPartitioned, gby=[upper(lineitem.l_returnflag)@0 as upper(lineitem.l_returnflag)], aggr=[count(Int64(1))], metrics=[hash_rows_computed=0, hash_rows_reused=3] + RepartitionExec: partitioning=Hash([upper(lineitem.l_returnflag)@0], 3), input_partitions=1, metrics=[hash_rows_computed=0, hash_rows_reused=3] + AggregateExec: mode=Partial, gby=[upper(l_returnflag@0) as upper(lineitem.l_returnflag)], aggr=[count(Int64(1))], metrics=[hash_rows_computed=20, hash_rows_reused=0] + DataSourceExec: file_groups={1 group: [[$DATAFUSION_CORE/tests/data/tpch_lineitem_small.parquet]]}, projection=[l_returnflag], file_type=parquet, metrics=[] + "); + + Ok(()) +} + +#[tokio::test] +async fn grouping_sets_tpch_aggregate_reuses_hashes_after_repartition() -> Result<()> { + let plan_with_hash_metrics = HashReuseTest::new( + r"SELECT l_returnflag, l_linestatus, COUNT(*) AS row_count + FROM lineitem + GROUP BY GROUPING SETS ((l_returnflag), (l_linestatus))", + ) + .run() + .await?; + + assert_snapshot!(plan_with_hash_metrics, @r" + ProjectionExec: expr=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus, count(Int64(1))@3 as row_count], metrics=[] + AggregateExec: mode=FinalPartitioned, gby=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus, __grouping_id@2 as __grouping_id], aggr=[count(Int64(1))], metrics=[hash_rows_computed=0, hash_rows_reused=5] + RepartitionExec: partitioning=Hash([l_returnflag@0, l_linestatus@1, __grouping_id@2], 3), input_partitions=1, metrics=[hash_rows_computed=0, hash_rows_reused=5] + AggregateExec: mode=Partial, gby=[(l_returnflag@0 as l_returnflag, NULL as l_linestatus), (NULL as l_returnflag, l_linestatus@1 as l_linestatus)], aggr=[count(Int64(1))], metrics=[hash_rows_computed=40, hash_rows_reused=0] + DataSourceExec: file_groups={1 group: [[$DATAFUSION_CORE/tests/data/tpch_lineitem_small.parquet]]}, projection=[l_returnflag, l_linestatus], file_type=parquet, metrics=[] + "); + + Ok(()) +} + +#[tokio::test] +async fn distinct_tpch_aggregate_reuses_hashes_after_repartition() -> Result<()> { + let plan_with_hash_metrics = + HashReuseTest::new("SELECT DISTINCT l_returnflag, l_linestatus FROM lineitem") + .run() + .await?; + + assert_snapshot!(plan_with_hash_metrics, @r" + AggregateExec: mode=FinalPartitioned, gby=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus], aggr=[], metrics=[hash_rows_computed=0, hash_rows_reused=3] + RepartitionExec: partitioning=Hash([l_returnflag@0, l_linestatus@1], 3), input_partitions=1, metrics=[hash_rows_computed=0, hash_rows_reused=3] + AggregateExec: mode=Partial, gby=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus], aggr=[], metrics=[hash_rows_computed=20, hash_rows_reused=0] + DataSourceExec: file_groups={1 group: [[$DATAFUSION_CORE/tests/data/tpch_lineitem_small.parquet]]}, projection=[l_returnflag, l_linestatus], file_type=parquet, metrics=[] + "); + + Ok(()) +} + +#[tokio::test] +async fn ordered_tpch_aggregate_reuses_hashes_after_repartition() -> Result<()> { + let plan_with_hash_metrics = HashReuseTest::new( + r"SELECT l_returnflag, max(row_num) AS max_row_num + FROM ( + SELECT l_returnflag, row_number() OVER (ORDER BY l_returnflag) AS row_num + FROM lineitem + ) + GROUP BY l_returnflag", + ) + .run() + .await?; + + assert_snapshot!(plan_with_hash_metrics, @r##" + ProjectionExec: expr=[l_returnflag@0 as l_returnflag, max(row_num)@1 as max_row_num], metrics=[] + AggregateExec: mode=FinalPartitioned, gby=[l_returnflag@0 as l_returnflag], aggr=[max(row_num)], ordering_mode=Sorted, metrics=[hash_rows_computed=0, hash_rows_reused=3] + RepartitionExec: partitioning=Hash([l_returnflag@0], 3), input_partitions=1, maintains_sort_order=true, metrics=[hash_rows_computed=0, hash_rows_reused=3] + AggregateExec: mode=Partial, gby=[l_returnflag@0 as l_returnflag], aggr=[max(row_num)], ordering_mode=Sorted, metrics=[hash_rows_computed=20, hash_rows_reused=0] + ProjectionExec: expr=[l_returnflag@0 as l_returnflag, row_number() ORDER BY [lineitem.l_returnflag ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@1 as row_num], metrics=[] + BoundedWindowAggExec: wdw=[row_number() ORDER BY [lineitem.l_returnflag ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() ORDER BY [lineitem.l_returnflag ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted], metrics=[] + SortExec: expr=[l_returnflag@0 ASC NULLS LAST], preserve_partitioning=[false], metrics=[] + DataSourceExec: file_groups={1 group: [[$DATAFUSION_CORE/tests/data/tpch_lineitem_small.parquet]]}, projection=[l_returnflag], file_type=parquet, sort_order_for_reorder=[l_returnflag@0 ASC NULLS LAST], metrics=[] + "##); + + Ok(()) +} + +#[tokio::test] +async fn summed_tpch_aggregate_reuses_hashes_after_repartition() -> Result<()> { + let plan_with_hash_metrics = HashReuseTest::new( + r"SELECT l_returnflag, sum(l_quantity) AS total_quantity + FROM lineitem + GROUP BY l_returnflag", + ) + .run() + .await?; + + assert_snapshot!(plan_with_hash_metrics, @r" + ProjectionExec: expr=[l_returnflag@0 as l_returnflag, sum(lineitem.l_quantity)@1 as total_quantity], metrics=[] + AggregateExec: mode=FinalPartitioned, gby=[l_returnflag@0 as l_returnflag], aggr=[sum(lineitem.l_quantity)], metrics=[hash_rows_computed=0, hash_rows_reused=3] + RepartitionExec: partitioning=Hash([l_returnflag@0], 3), input_partitions=1, metrics=[hash_rows_computed=0, hash_rows_reused=3] + AggregateExec: mode=Partial, gby=[l_returnflag@1 as l_returnflag], aggr=[sum(lineitem.l_quantity)], metrics=[hash_rows_computed=20, hash_rows_reused=0] + DataSourceExec: file_groups={1 group: [[$DATAFUSION_CORE/tests/data/tpch_lineitem_small.parquet]]}, projection=[l_quantity, l_returnflag], file_type=parquet, metrics=[] + "); + + Ok(()) +} + +struct HashReuseTest<'a> { + sql: &'a str, + batch_size: usize, + enable_migration_aggregate: bool, +} + +impl<'a> HashReuseTest<'a> { + fn new(sql: &'a str) -> Self { + Self { + sql, + batch_size: 64, + enable_migration_aggregate: true, + } + } + + fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = batch_size; + self + } + + fn with_migration_aggregate(mut self, enable: bool) -> Self { + self.enable_migration_aggregate = enable; + self + } + + async fn run(self) -> Result { + let config = SessionConfig::new() + .with_target_partitions(3) + .with_batch_size(self.batch_size) + .set_bool( + "datafusion.execution.enable_migration_aggregate", + self.enable_migration_aggregate, + ); + let ctx = SessionContext::new_with_config(config); + ctx.register_parquet( + "lineitem", + "tests/data/tpch_lineitem_small.parquet", + ParquetReadOptions::default(), + ) + .await?; + let dataframe = ctx.sql(self.sql).await?; + let plan = dataframe.create_physical_plan().await?; + let output = collect(plan.clone(), ctx.task_ctx()).await?; + assert!(!output.is_empty()); + + let manifest_dir = env!("CARGO_MANIFEST_DIR"); + Ok(DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .set_metric_names(vec![ + "hash_rows_reused".into(), + "hash_rows_computed".into(), + ]) + .indent(true) + .to_string() + .replace(manifest_dir, "$DATAFUSION_CORE") + .replace(manifest_dir.trim_start_matches('/'), "$DATAFUSION_CORE")) + } +} diff --git a/datafusion/core/tests/execution/mod.rs b/datafusion/core/tests/execution/mod.rs index f33ef87aa3023..d55457fcd2614 100644 --- a/datafusion/core/tests/execution/mod.rs +++ b/datafusion/core/tests/execution/mod.rs @@ -17,5 +17,6 @@ mod coop; mod datasource_split; +mod hash_reuse; mod logical_plan; mod register_arrow; 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 71dfb1571d141..397a3ebbf72a0 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -441,6 +441,8 @@ pub(super) struct AggregateHashTableBuffer { pub(super) batch_group_indices: Vec, /// Hashes retained for the internal hash column when output requires it. + /// If this field is `None`, it means that the output does not require to emit + /// intermediate hashes. pub(super) group_hash_tracker: Option, /// Computes group hashes when the input has no reusable hash column. From 8c19918666fe421ce79c583ee7ba674699e3b85f Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 20 Jul 2026 14:00:39 +0200 Subject: [PATCH 4/6] Address clippy::large_enum_variant --- .../physical-plan/src/aggregates/hash_stream.rs | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 62b92965030ae..9270fe0f42cda 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -148,7 +148,7 @@ enum PartialHashAggregateState { /// finish in `Done`. If `Some`, partial skip has triggered and the /// stream will move to `SkippingAggregation` after these accumulated /// groups are emitted. - skip_hash_table: Option>, + skip_hash_table: Option>>, }, SkippingAggregation { hash_table: AggregateHashTable, @@ -457,7 +457,7 @@ impl PartialHashAggregateStream { return ControlFlow::Continue( PartialHashAggregateState::ProducingOutput { hash_table, - skip_hash_table: Some(skip_hash_table), + skip_hash_table: Some(Box::new(skip_hash_table)), }, ); } @@ -547,9 +547,9 @@ impl PartialHashAggregateStream { PartialHashAggregateState::ProducingOutput { skip_hash_table: Some(hash_table), .. - } => { - PartialHashAggregateState::SkippingAggregation { hash_table } - } + } => PartialHashAggregateState::SkippingAggregation { + hash_table: *hash_table, + }, PartialHashAggregateState::ProducingOutput { skip_hash_table: None, .. @@ -573,7 +573,9 @@ impl PartialHashAggregateStream { PartialHashAggregateState::ProducingOutput { skip_hash_table: Some(hash_table), .. - } => PartialHashAggregateState::SkippingAggregation { hash_table }, + } => PartialHashAggregateState::SkippingAggregation { + hash_table: *hash_table, + }, PartialHashAggregateState::ProducingOutput { skip_hash_table: None, .. From 39284b1f1b39a8e6f925025d490052869357a4fd Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 21 Jul 2026 11:22:05 +0200 Subject: [PATCH 5/6] Only propagate hashes if expressions evaluate to non-primitive --- datafusion/core/tests/execution/hash_reuse.rs | 21 ++++++++++++ .../partition_statistics.rs | 23 +++---------- .../aggregates/aggregate_hash_table/common.rs | 21 +++++++----- .../aggregate_hash_table/common_ordered.rs | 19 ++++++----- .../src/aggregates/grouped_hash_stream.rs | 11 +++--- .../src/aggregates/grouped_topk_stream.rs | 8 +++-- .../physical-plan/src/aggregates/mod.rs | 34 ++++++++++--------- .../physical-plan/src/repartition/hash.rs | 19 +++++++++++ 8 files changed, 97 insertions(+), 59 deletions(-) diff --git a/datafusion/core/tests/execution/hash_reuse.rs b/datafusion/core/tests/execution/hash_reuse.rs index d2e2c01f4ade0..428b44d418ce8 100644 --- a/datafusion/core/tests/execution/hash_reuse.rs +++ b/datafusion/core/tests/execution/hash_reuse.rs @@ -195,6 +195,27 @@ async fn summed_tpch_aggregate_reuses_hashes_after_repartition() -> Result<()> { Ok(()) } +#[tokio::test] +async fn primitive_grouped_tpch_aggregate_does_not_propagate_hashes() -> Result<()> { + let plan_with_hash_metrics = HashReuseTest::new( + r"SELECT l_linenumber, COUNT(*) AS row_count + FROM lineitem + GROUP BY l_linenumber", + ) + .run() + .await?; + + assert_snapshot!(plan_with_hash_metrics, @r" + ProjectionExec: expr=[l_linenumber@0 as l_linenumber, count(Int64(1))@1 as row_count], metrics=[] + AggregateExec: mode=FinalPartitioned, gby=[l_linenumber@0 as l_linenumber], aggr=[count(Int64(1))], metrics=[hash_rows_computed=6, hash_rows_reused=0] + RepartitionExec: partitioning=Hash([l_linenumber@0], 3), input_partitions=1, metrics=[hash_rows_computed=6, hash_rows_reused=0] + AggregateExec: mode=Partial, gby=[l_linenumber@0 as l_linenumber], aggr=[count(Int64(1))], metrics=[hash_rows_computed=20, hash_rows_reused=0] + DataSourceExec: file_groups={1 group: [[$DATAFUSION_CORE/tests/data/tpch_lineitem_small.parquet]]}, projection=[l_linenumber], file_type=parquet, metrics=[] + "); + + Ok(()) +} + struct HashReuseTest<'a> { sql: &'a str, batch_size: usize, diff --git a/datafusion/core/tests/physical_optimizer/partition_statistics.rs b/datafusion/core/tests/physical_optimizer/partition_statistics.rs index 2aaa3032685ce..5d6672c9cd3b3 100644 --- a/datafusion/core/tests/physical_optimizer/partition_statistics.rs +++ b/datafusion/core/tests/physical_optimizer/partition_statistics.rs @@ -911,12 +911,7 @@ 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_partial_p0_statistics); + assert_eq!(*p0_statistics, expected_p0_statistics); let expected_p1_statistics = Statistics { num_rows: Precision::Inexact(2), @@ -934,16 +929,11 @@ 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_partial_p1_statistics); + assert_eq!(*p1_statistics, expected_p1_statistics); validate_statistics_with_data( aggregate_exec_partial.clone(), @@ -1017,20 +1007,15 @@ 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_partial_stat, + empty_stat, *StatisticsContext::new().compute( agg_partial.as_ref(), &StatisticsArgs::new().with_partition(Some(0)) )? ); assert_eq!( - empty_partial_stat, + empty_stat, *StatisticsContext::new().compute( agg_partial.as_ref(), &StatisticsArgs::new().with_partition(Some(1)) 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 397a3ebbf72a0..74556f74c9f51 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,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, GroupHashTracker, PhysicalGroupBy, aggregate_expressions, - evaluate_group_by, outputs_group_hashes, + AggregateExec, AggregateOutputMode, GroupHashTracker, PhysicalGroupBy, + aggregate_expressions, evaluate_group_by, }; use crate::repartition::{ExpressionHasher, HashMetrics}; @@ -131,21 +131,24 @@ impl AggregateHashTable { let group_schema = agg.group_by.group_schema(&input_schema)?; let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + let hasher = ExpressionHasher::new_with_metrics( + agg.group_by.input_exprs(), + HashMetrics::new(&agg.metrics, partition), + ); + let should_output_hashes = hasher.should_output_hashes(&input_schema)? + && agg.mode.output_mode() == AggregateOutputMode::Partial; + Ok(Self { group_by_metrics: GroupByMetrics::new(&agg.metrics, partition), - input_schema, + input_schema: Arc::clone(&input_schema), output_schema, batch_size, state: AggregateHashTableState::Building(AggregateHashTableBuffer { 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), - ), + group_hash_tracker: should_output_hashes.then(GroupHashTracker::default), + hasher, accumulators, }), _mode: PhantomData, 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 f25c0f9bb0ff8..902dfebfd6d6a 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,8 +34,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, AggregateMode, GroupHashTracker, PhysicalGroupBy, - aggregate_expressions, evaluate_group_by, outputs_group_hashes, + AggregateExec, AggregateMode, AggregateOutputMode, GroupHashTracker, PhysicalGroupBy, + aggregate_expressions, evaluate_group_by, }; use crate::repartition::{ExpressionHasher, HashMetrics}; @@ -173,6 +173,13 @@ impl OrderedAggregateTable { }) .collect::>()?; + let hasher = ExpressionHasher::new_with_metrics( + agg.group_by.input_exprs(), + HashMetrics::new(&agg.metrics, partition), + ); + let should_output_hashes = hasher.should_output_hashes(input_schema)? + && aggregate_mode.output_mode() == AggregateOutputMode::Partial; + Ok(Self { output_schema, batch_size, @@ -182,12 +189,8 @@ impl OrderedAggregateTable { 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), - ), + group_hash_tracker: should_output_hashes.then(GroupHashTracker::default), + hasher, accumulators, }, _mode: PhantomData, diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index ea365bdb14f68..fa0e591f0c6b3 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -594,13 +594,16 @@ 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), ); + // Spilled batches contain partial state that a later aggregate merges, + // so they use the same hash-output policy as regular partial output. + let should_output_hashes = hasher.should_output_hashes(&agg.input().schema())? + && (agg.mode.output_mode() == AggregateOutputMode::Partial + || oom_mode == OutOfMemoryMode::Spill); Ok(GroupedHashAggregateStream { schema: agg_schema, @@ -615,7 +618,7 @@ impl GroupedHashAggregateStream { oom_mode, group_values, current_group_indices: Default::default(), - group_hash_tracker: track_group_hashes.then(GroupHashTracker::default), + group_hash_tracker: should_output_hashes.then(GroupHashTracker::default), hasher, exec_state, baseline_metrics, diff --git a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs index 1d512fdf53b5b..6323f2e87dffe 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs @@ -22,8 +22,8 @@ use crate::aggregates::topk::priority_map::PriorityMap; #[cfg(debug_assertions)] use crate::aggregates::topk_types_supported; use crate::aggregates::{ - AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, - evaluate_many, outputs_group_hashes, + AggregateExec, AggregateOutputMode, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, evaluate_many, }; use crate::metrics::BaselineMetrics; use crate::repartition::{ExpressionHasher, HashMetrics}; @@ -112,7 +112,9 @@ 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_group_hashes = ExpressionHasher::new(group_by.input_exprs()) + .should_output_hashes(&aggr.input().schema())? + && aggr.mode.output_mode() == AggregateOutputMode::Partial; let output_hasher = ExpressionHasher::new_with_metrics( group_by.output_exprs(), HashMetrics::new(&aggr.metrics, partition), diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 73c9395e31a80..abb1059f0de0e 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -229,10 +229,6 @@ 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) } -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 @@ -2248,10 +2244,11 @@ fn create_schema( aggr_expr: &[Arc], mode: AggregateMode, ) -> Result { + let input_hasher = ExpressionHasher::new(group_by.input_exprs()); + let output_hashes = input_hasher.should_output_hashes(input_schema)? + && mode.output_mode() == AggregateOutputMode::Partial; let mut fields = Vec::with_capacity( - group_by.num_output_exprs() - + aggr_expr.len() - + outputs_group_hashes(mode, group_by) as usize, + group_by.num_output_exprs() + aggr_expr.len() + output_hashes as usize, ); fields.extend(group_by.output_fields(input_schema)?); @@ -2270,7 +2267,7 @@ fn create_schema( } } - if outputs_group_hashes(mode, group_by) { + if output_hashes { let hasher = ExpressionHasher::new(group_by.output_exprs()); fields.push( Field::new(hasher.internal_hash_col_name(), DataType::UInt64, false).into(), @@ -2757,7 +2754,7 @@ mod tests { use arrow::array::{ BooleanArray, DictionaryArray, Float32Array, Float64Array, Int32Array, - Int64Array, StructArray, UInt32Array, UInt64Array, + Int64Array, StringArray, StructArray, UInt32Array, UInt64Array, }; use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::Int32Type; @@ -2811,6 +2808,9 @@ mod tests { batches .iter() .map(|batch| { + let Ok(hash_index) = batch.schema().index_of(&hash_name) else { + return Ok(batch.clone()); + }; let actual = hasher .precomputed(batch) .expect("partial aggregate output should contain group hashes"); @@ -2818,7 +2818,6 @@ mod tests { 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::>(); @@ -2835,6 +2834,9 @@ mod tests { let hash_exprs = group_by.output_exprs(); let mut hasher = ExpressionHasher::new(hash_exprs.clone()); let hash_name = hasher.internal_hash_col_name(); + if schema.index_of(&hash_name).is_err() { + return RecordBatch::try_new(schema, columns).map_err(Into::into); + } let fields = schema .fields() .iter() @@ -3676,21 +3678,21 @@ mod tests { #[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("key", DataType::Utf8, 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(StringArray::from(vec!["a", "b", "a"])), 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(StringArray::from(vec!["b", "c"])), Arc::new(Int32Array::from(vec![1, 1])), ], )?], @@ -3783,9 +3785,9 @@ mod tests { +-----+--------------+ | key | COUNT(value) | +-----+--------------+ -| 1 | 2 | -| 2 | 2 | -| 3 | 1 | +| a | 2 | +| b | 2 | +| c | 1 | +-----+--------------+ "); diff --git a/datafusion/physical-plan/src/repartition/hash.rs b/datafusion/physical-plan/src/repartition/hash.rs index a50c2eb90f415..8a0d1c5d71c7c 100644 --- a/datafusion/physical-plan/src/repartition/hash.rs +++ b/datafusion/physical-plan/src/repartition/hash.rs @@ -19,6 +19,7 @@ use crate::aggregates::PhysicalGroupBy; use crate::joins::SeededRandomState; use crate::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory}; use arrow::array::{ArrayRef, RecordBatch}; +use arrow::datatypes::Schema; use datafusion_common::Result; use datafusion_common::cast::as_uint64_array; use datafusion_common::hash_utils::create_hashes; @@ -104,6 +105,24 @@ impl ExpressionHasher { } } + /// Returns whether propagating hashes is worthwhile for these expressions. + /// + /// Hashing primitive values is inexpensive, so the cost of propagating an + /// additional `UInt64` column outweighs avoiding a later hash computation. + /// Variable-width and nested values are more expensive to hash, and retain + /// the propagated hashes when at least one expression has a non-primitive + /// type. Callers combine this policy with their output mode, such as + /// whether they are emitting partial aggregation state. + pub(crate) fn should_output_hashes(&self, input_schema: &Schema) -> Result { + for expr in &self.hash_exprs { + if !expr.data_type(input_schema)?.is_primitive() { + return Ok(true); + } + } + + Ok(false) + } + /// 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 From 9bcc60a5097dfebff0167a255884ab979048870a Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 21 Jul 2026 18:11:40 +0200 Subject: [PATCH 6/6] Represent hash buffers as UInt64Array everywhere --- .../physical-expr-common/src/binary_map.rs | 24 +++++--- .../src/binary_view_map.rs | 21 +++---- .../benches/dictionary_group_values.rs | 18 +++--- .../physical-plan/benches/multi_group_by.rs | 41 ++++++++++---- .../aggregates/aggregate_hash_table/common.rs | 4 +- .../aggregate_hash_table/partial_table.rs | 4 +- .../src/aggregates/group_values/mod.rs | 4 +- .../group_values/multi_group_by/mod.rs | 18 +++--- .../src/aggregates/group_values/row.rs | 6 +- .../group_values/single_group_by/boolean.rs | 3 +- .../group_values/single_group_by/bytes.rs | 4 +- .../single_group_by/bytes_view.rs | 4 +- .../group_values/single_group_by/primitive.rs | 6 +- .../src/aggregates/grouped_hash_stream.rs | 2 +- .../src/aggregates/grouped_topk_stream.rs | 4 +- .../physical-plan/src/aggregates/mod.rs | 55 ++++++++++++------- .../physical-plan/src/repartition/hash.rs | 46 ++++++++++------ .../physical-plan/src/repartition/mod.rs | 12 ++-- 18 files changed, 168 insertions(+), 108 deletions(-) diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 77bce64eed95e..6d9ceee384dfd 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -20,7 +20,7 @@ use arrow::array::{ Array, ArrayRef, BufferBuilder, GenericBinaryArray, GenericStringArray, - NullBufferBuilder, OffsetSizeTrait, + NullBufferBuilder, OffsetSizeTrait, UInt64Array, cast::AsArray, types::{ByteArrayType, GenericBinaryType, GenericStringType}, }; @@ -228,7 +228,7 @@ where /// random state used to generate hashes random_state: RandomState, /// buffer that stores hash values (reused across batches to save allocations) - hashes_buffer: Vec, + hashes_buffer: UInt64Array, /// `(payload, null_index)` for the 'null' value, if any /// NOTE null_index is the logical index in the final array, not the index /// in the buffer @@ -251,7 +251,7 @@ where buffer: BufferBuilder::new(INITIAL_BUFFER_CAPACITY), offsets: vec![O::default()], // first offset is always 0 random_state: RandomState::default(), - hashes_buffer: vec![], + hashes_buffer: UInt64Array::from(Vec::::new()), null: None, } } @@ -299,13 +299,14 @@ where MP: FnMut(Option<&[u8]>) -> V, OP: FnMut(V), { - let mut hashes = std::mem::take(&mut self.hashes_buffer); + let mut hashes = take_hash_buffer(&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(); + let hashes = UInt64Array::from(hashes); self.insert_if_new_with_hashes( values, &hashes, @@ -319,7 +320,7 @@ where pub fn insert_if_new_with_hashes( &mut self, values: &ArrayRef, - hashes: &[u64], + hashes: &UInt64Array, make_payload_fn: MP, observe_payload_fn: OP, ) where @@ -368,7 +369,7 @@ where fn insert_if_new_inner( &mut self, values: &ArrayRef, - hashes: &[u64], + hashes: &UInt64Array, mut make_payload_fn: MP, mut observe_payload_fn: OP, ) where @@ -381,7 +382,7 @@ where // Ensure lengths are equivalent assert_eq!(values.len(), hashes.len()); - for (value, &hash) in values.iter().zip(hashes) { + for (value, &hash) in values.iter().zip(hashes.values()) { // handle null value let Some(value) = value else { let payload = if let Some(&(payload, _offset)) = self.null.as_ref() { @@ -568,10 +569,17 @@ where self.map_size + self.buffer.capacity() * size_of::() + self.offsets.allocated_size() - + self.hashes_buffer.allocated_size() + + self.hashes_buffer.get_buffer_memory_size() } } +pub(crate) fn take_hash_buffer(hashes: &mut UInt64Array) -> Vec { + let empty = UInt64Array::from(Vec::::new()); + let (_, values, nulls) = std::mem::replace(hashes, empty).into_parts(); + debug_assert!(nulls.is_none()); + values.into() +} + /// Returns a `NullBuffer` with a single null value at the given index fn single_null_buffer(num_values: usize, null_index: usize) -> NullBuffer { let mut null_builder = NullBufferBuilder::new(num_values); diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 34b4df8b25200..1d9453a022fc8 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -17,15 +17,15 @@ //! [`ArrowBytesViewMap`] and [`ArrowBytesViewSet`] for storing maps/sets of values from //! `StringViewArray`/`BinaryViewArray`. -use crate::binary_map::OutputType; +use crate::binary_map::{OutputType, take_hash_buffer}; use arrow::array::NullBufferBuilder; use arrow::array::cast::AsArray; -use arrow::array::{Array, ArrayRef, BinaryViewArray, ByteView, make_view}; +use arrow::array::{Array, ArrayRef, BinaryViewArray, ByteView, UInt64Array, make_view}; use arrow::buffer::{Buffer, ScalarBuffer}; use arrow::datatypes::{BinaryViewType, ByteViewType, DataType, StringViewType}; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; -use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_common::utils::proxy::HashTableAllocExt; use std::fmt::Debug; use std::mem::size_of; use std::sync::Arc; @@ -140,7 +140,7 @@ where /// random state used to generate hashes random_state: RandomState, /// buffer that stores hash values (reused across batches to save allocations) - hashes_buffer: Vec, + hashes_buffer: UInt64Array, /// `(payload, null_index)` for the 'null' value, if any /// NOTE null_index is the logical index in the final array, not the index /// in the buffer @@ -164,7 +164,7 @@ where completed: Vec::new(), nulls: NullBufferBuilder::new(0), random_state: RandomState::default(), - hashes_buffer: vec![], + hashes_buffer: UInt64Array::from(Vec::::new()), null: None, } } @@ -212,13 +212,14 @@ where MP: FnMut(Option<&[u8]>) -> V, OP: FnMut(V), { - let mut hashes = std::mem::take(&mut self.hashes_buffer); + let mut hashes = take_hash_buffer(&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(); + let hashes = UInt64Array::from(hashes); self.insert_if_new_with_hashes( values, &hashes, @@ -232,7 +233,7 @@ where pub fn insert_if_new_with_hashes( &mut self, values: &ArrayRef, - hashes: &[u64], + hashes: &UInt64Array, make_payload_fn: MP, observe_payload_fn: OP, ) where @@ -275,7 +276,7 @@ where fn insert_if_new_inner( &mut self, values: &ArrayRef, - hashes: &[u64], + hashes: &UInt64Array, mut make_payload_fn: MP, mut observe_payload_fn: OP, ) where @@ -293,7 +294,7 @@ where for i in 0..values.len() { let view_u128 = input_views[i]; - let hash = hashes[i]; + let hash = hashes.value(i); // handle null value via validity bitmap check if values.is_null(i) { @@ -502,7 +503,7 @@ where + in_progress_size + completed_size + nulls_size - + self.hashes_buffer.allocated_size() + + self.hashes_buffer.get_buffer_memory_size() } } diff --git a/datafusion/physical-plan/benches/dictionary_group_values.rs b/datafusion/physical-plan/benches/dictionary_group_values.rs index 756860da9fe18..f0bf7ee3fbe6e 100644 --- a/datafusion/physical-plan/benches/dictionary_group_values.rs +++ b/datafusion/physical-plan/benches/dictionary_group_values.rs @@ -21,7 +21,7 @@ //! `new_group_values` is constructed in the setup closure of //! `iter_batched_ref` and is not included in the timing. -use arrow::array::{ArrayRef, DictionaryArray, PrimitiveArray, StringArray}; +use arrow::array::{ArrayRef, DictionaryArray, PrimitiveArray, StringArray, UInt64Array}; use arrow::buffer::{Buffer, NullBuffer}; use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef}; use criterion::{ @@ -44,15 +44,19 @@ const N_BATCHES: usize = 4; const SEED: u64 = 0xD1C7; const AGGREGATION_HASH_SEED: u64 = 15395726432021054657; -fn create_aggregation_hashes(array: &ArrayRef, hashes: &mut Vec) { - hashes.clear(); - hashes.resize(array.len(), 0); +fn create_aggregation_hashes(array: &ArrayRef, hashes: &mut UInt64Array) { + let empty = UInt64Array::from(Vec::::new()); + let (_, values, _) = std::mem::replace(hashes, empty).into_parts(); + let mut values: Vec = values.into(); + values.clear(); + values.resize(array.len(), 0); create_hashes( std::slice::from_ref(array), &RandomState::with_seed(AGGREGATION_HASH_SEED), - hashes, + &mut values, ) .unwrap(); + *hashes = UInt64Array::from(values); } fn dict_schema() -> SchemaRef { @@ -122,7 +126,7 @@ fn bench_intern_emit(c: &mut Criterion) { new_group_values(schema.clone(), &GroupOrdering::None) .unwrap(), Vec::::with_capacity(size), - Vec::::with_capacity(size), + UInt64Array::from(Vec::::with_capacity(size)), ) }, |(gv, groups, hashes)| { @@ -170,7 +174,7 @@ fn bench_repeated_intern_emit(c: &mut Criterion) { new_group_values(schema.clone(), &GroupOrdering::None) .unwrap(), Vec::::with_capacity(size), - Vec::::with_capacity(size), + UInt64Array::from(Vec::::with_capacity(size)), ) }, |(gv, groups, hashes)| { diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 388cb4be9d564..e6a3eac6eeb2d 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -27,7 +27,7 @@ //! covers a `(FixedSizeBinary, Int32)` key to exercise the //! `FixedSizeBinaryGroupValueBuilder`. -use arrow::array::{ArrayRef, Int32Array, UInt32Array}; +use arrow::array::{ArrayRef, Int32Array, UInt32Array, UInt64Array}; use arrow::compute::take; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::util::bench_util::create_fsb_array; @@ -100,14 +100,18 @@ fn bench_intern( gv: &mut Box, batches: &[Vec], groups: &mut Vec, - hashes: &mut Vec, + hashes: &mut UInt64Array, ) { let random_state = RandomState::with_seed(AGGREGATION_HASH_SEED); for batch in batches { groups.clear(); - hashes.clear(); - hashes.resize(batch[0].len(), 0); - create_hashes(batch, &random_state, hashes).unwrap(); + let empty = UInt64Array::from(Vec::::new()); + let (_, values, _) = std::mem::replace(hashes, empty).into_parts(); + let mut values: Vec = values.into(); + values.clear(); + values.resize(batch[0].len(), 0); + create_hashes(batch, &random_state, &mut values).unwrap(); + *hashes = UInt64Array::from(values); gv.intern(batch, groups, hashes).unwrap(); } black_box(&*groups); @@ -142,7 +146,9 @@ fn bench_issue_17850_regression(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(DEFAULT_BATCH_SIZE), - Vec::::with_capacity(DEFAULT_BATCH_SIZE), + UInt64Array::from(Vec::::with_capacity( + DEFAULT_BATCH_SIZE, + )), ) }, |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), @@ -186,7 +192,9 @@ fn bench_low_cardinality(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(DEFAULT_BATCH_SIZE), - Vec::::with_capacity(DEFAULT_BATCH_SIZE), + UInt64Array::from(Vec::::with_capacity( + DEFAULT_BATCH_SIZE, + )), ) }, |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), @@ -226,7 +234,7 @@ fn bench_batch_size_sensitivity(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(batch_size), - Vec::::with_capacity(batch_size), + UInt64Array::from(Vec::::with_capacity(batch_size)), ) }, |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), @@ -267,7 +275,9 @@ fn bench_column_scaling(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(DEFAULT_BATCH_SIZE), - Vec::::with_capacity(DEFAULT_BATCH_SIZE), + UInt64Array::from(Vec::::with_capacity( + DEFAULT_BATCH_SIZE, + )), ) }, |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), @@ -306,7 +316,9 @@ fn bench_high_cardinality_scaling(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(DEFAULT_BATCH_SIZE), - Vec::::with_capacity(DEFAULT_BATCH_SIZE), + UInt64Array::from(Vec::::with_capacity( + DEFAULT_BATCH_SIZE, + )), ) }, |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), @@ -348,7 +360,9 @@ fn bench_group_count_sweep(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(DEFAULT_BATCH_SIZE), - Vec::::with_capacity(DEFAULT_BATCH_SIZE), + UInt64Array::from(Vec::::with_capacity( + DEFAULT_BATCH_SIZE, + )), ) }, |(gv, groups, hashes)| bench_intern(gv, batches, groups, hashes), @@ -445,9 +459,12 @@ fn bench_fixed_size_binary(c: &mut Criterion) { ( create_group_values(&schema, vectorized), Vec::::with_capacity(DEFAULT_BATCH_SIZE), + UInt64Array::from(Vec::::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 74556f74c9f51..9049bfcd7fd2b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -18,7 +18,7 @@ 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::{Result, internal_err}; @@ -420,7 +420,7 @@ pub(super) struct EvaluatedAggregateBatch<'a> { pub(super) accumulator_args: Vec, /// Hashes carried by the input batch for a simple GROUP BY, when available. - pub(super) precomputed_group_hashes: Option<&'a [u64]>, + pub(super) precomputed_group_hashes: Option<&'a UInt64Array>, } /// Buffer for the aggregate hash table's group keys and accumulator states. 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 41949dc67ab0a..c1c838a684408 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -19,7 +19,7 @@ use std::collections::HashMap; use std::marker::PhantomData; use std::sync::Arc; -use arrow::array::{ArrayRef, BooleanArray, UInt64Array, new_null_array}; +use arrow::array::{ArrayRef, BooleanArray, new_null_array}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, assert_eq_or_internal_err}; @@ -227,7 +227,7 @@ impl AggregateHashTable { .hasher .compute_hashes(&evaluated_batch.grouping_set_args[0])?, }; - Some(Arc::new(UInt64Array::from(hashes.to_vec())) as ArrayRef) + Some(Arc::new(hashes.clone()) as ArrayRef) } else { None }; diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index d6024eea4a1a2..151fcf65aa0b7 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -22,7 +22,7 @@ use arrow::array::types::{ Time64MicrosecondType, Time64NanosecondType, TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, }; -use arrow::array::{ArrayRef, downcast_primitive}; +use arrow::array::{ArrayRef, UInt64Array, downcast_primitive}; use arrow::datatypes::{DataType, SchemaRef, TimeUnit}; use datafusion_common::Result; @@ -102,7 +102,7 @@ pub trait GroupValues: Send { &mut self, cols: &[ArrayRef], groups: &mut Vec, - hashes: &[u64], + hashes: &UInt64Array, ) -> Result<()>; /// Returns the number of bytes of memory used by this [`GroupValues`] 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 7e99ee184c326..c3122579ae857 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 @@ -29,7 +29,7 @@ use crate::aggregates::group_values::multi_group_by::{ boolean::BooleanGroupValueBuilder, bytes::ByteGroupValueBuilder, bytes_view::ByteViewGroupValueBuilder, primitive::PrimitiveGroupValueBuilder, }; -use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; +use arrow::array::{Array, ArrayRef, BooleanBufferBuilder, UInt64Array}; use arrow::compute::cast; use arrow::datatypes::{ BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, Field, Float32Type, @@ -338,7 +338,7 @@ impl GroupValuesColumn { &mut self, cols: &[ArrayRef], groups: &mut Vec, - hashes: &[u64], + hashes: &UInt64Array, ) -> Result<()> { let n_rows = cols[0].len(); assert_eq!(n_rows, hashes.len()); @@ -346,7 +346,7 @@ impl GroupValuesColumn { // tracks to which group each of the input rows belongs groups.clear(); - for (row, &target_hash) in hashes.iter().enumerate() { + for (row, &target_hash) in hashes.values().iter().enumerate() { let entry = self .map .find_mut(target_hash, |(exist_hash, group_idx_view)| { @@ -435,7 +435,7 @@ impl GroupValuesColumn { &mut self, cols: &[ArrayRef], groups: &mut Vec, - hashes: &[u64], + hashes: &UInt64Array, ) -> Result<()> { let n_rows = cols[0].len(); assert_eq!(n_rows, hashes.len()); @@ -495,7 +495,7 @@ impl GroupValuesColumn { /// Otherwise get all group indices from `group_index_lists`, and add them. fn collect_vectorized_process_context( &mut self, - batch_hashes: &[u64], + batch_hashes: &UInt64Array, groups: &mut [usize], ) { self.vectorized_operation_buffers.append_row_indices.clear(); @@ -506,7 +506,7 @@ impl GroupValuesColumn { .equal_to_group_indices .clear(); - for (row, &target_hash) in batch_hashes.iter().enumerate() { + for (row, &target_hash) in batch_hashes.values().iter().enumerate() { let entry = self .map .find(target_hash, |(exist_hash, _)| target_hash == *exist_hash); @@ -722,7 +722,7 @@ impl GroupValuesColumn { fn scalarized_intern_remaining( &mut self, cols: &[ArrayRef], - batch_hashes: &[u64], + batch_hashes: &UInt64Array, groups: &mut [usize], ) -> Result<()> { if self @@ -736,7 +736,7 @@ impl GroupValuesColumn { let mut map = mem::take(&mut self.map); for &row in &self.vectorized_operation_buffers.remaining_row_indices { - let target_hash = batch_hashes[row]; + let target_hash = batch_hashes.value(row); let entry = map.find_mut(target_hash, |(exist_hash, _)| { // Somewhat surprisingly, this closure can be called even if the // hash doesn't match, so check the hash first with an integer @@ -1063,7 +1063,7 @@ impl GroupValues for GroupValuesColumn { &mut self, cols: &[ArrayRef], groups: &mut Vec, - hashes: &[u64], + hashes: &UInt64Array, ) -> Result<()> { // `try_new` and the reset points in `emit` / `clear_shrink` keep // `self.group_values` populated with one builder per schema field, diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 0be654bdbf05c..8ebe396cf371a 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -17,7 +17,7 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::{ - Array, ArrayRef, ListArray, PrimitiveArray, RunArray, StructArray, + Array, ArrayRef, ListArray, PrimitiveArray, RunArray, StructArray, UInt64Array, downcast_run_end_index, }; use arrow::compute::cast; @@ -110,7 +110,7 @@ impl GroupValues for GroupValuesRows { &mut self, cols: &[ArrayRef], groups: &mut Vec, - hashes: &[u64], + hashes: &UInt64Array, ) -> Result<()> { // Normalize -0.0 → +0.0 so RowConverter (IEEE 754 totalOrder) and // primitive hashing both group ±0 together. No-op for non-float @@ -134,7 +134,7 @@ impl GroupValues for GroupValuesRows { // tracks to which group each of the input rows belongs groups.clear(); - for (row, &target_hash) in hashes.iter().enumerate() { + for (row, &target_hash) in hashes.values().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 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 ac94c8d4acb7b..5edb121cd0c54 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 @@ -19,6 +19,7 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::{ ArrayRef, AsArray as _, BooleanArray, BooleanBufferBuilder, NullBufferBuilder, + UInt64Array, }; use datafusion_common::Result; use datafusion_expr::EmitTo; @@ -46,7 +47,7 @@ impl GroupValues for GroupValuesBoolean { &mut self, cols: &[ArrayRef], groups: &mut Vec, - hashes: &[u64], + hashes: &UInt64Array, ) -> Result<()> { let array = cols[0].as_boolean(); assert_eq!(array.len(), hashes.len()); 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 71bcb98509699..8a3d99f506a96 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 @@ -19,7 +19,7 @@ use std::mem::size_of; use crate::aggregates::group_values::GroupValues; -use arrow::array::{Array, ArrayRef, OffsetSizeTrait}; +use arrow::array::{Array, ArrayRef, OffsetSizeTrait, UInt64Array}; use datafusion_common::Result; use datafusion_expr::EmitTo; use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; @@ -49,7 +49,7 @@ impl GroupValues for GroupValuesBytes { &mut self, cols: &[ArrayRef], groups: &mut Vec, - hashes: &[u64], + hashes: &UInt64Array, ) -> Result<()> { assert_eq!(cols.len(), 1); 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 e4a2397b06aae..95caa7bbaed72 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 @@ -16,7 +16,7 @@ // under the License. use crate::aggregates::group_values::GroupValues; -use arrow::array::{Array, ArrayRef}; +use arrow::array::{Array, ArrayRef, UInt64Array}; use datafusion_expr::EmitTo; use datafusion_physical_expr::binary_map::OutputType; use datafusion_physical_expr_common::binary_view_map::ArrowBytesViewMap; @@ -47,7 +47,7 @@ impl GroupValues for GroupValuesBytesView { &mut self, cols: &[ArrayRef], groups: &mut Vec, - hashes: &[u64], + hashes: &UInt64Array, ) -> datafusion_common::Result<()> { assert_eq!(cols.len(), 1); 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 b8f1356d95bb7..8b29922e19386 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 @@ -19,7 +19,7 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano}; use arrow::array::{ ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, PrimitiveArray, - cast::AsArray, + UInt64Array, cast::AsArray, }; use arrow::datatypes::{DataType, i256}; use datafusion_common::Result; @@ -136,13 +136,13 @@ where &mut self, cols: &[ArrayRef], groups: &mut Vec, - hashes: &[u64], + hashes: &UInt64Array, ) -> Result<()> { assert_eq!(cols.len(), 1); assert_eq!(cols[0].len(), hashes.len()); groups.clear(); - for (v, &hash) in cols[0].as_primitive::().into_iter().zip(hashes) { + for (v, &hash) in cols[0].as_primitive::().into_iter().zip(hashes.values()) { let group_id = match v { None => *self.null_group.get_or_insert_with(|| { let group_id = self.values.len(); diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index fa0e591f0c6b3..e9f63af498192 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -1455,7 +1455,7 @@ impl GroupedHashAggregateStream { Some(hashes) => hashes, None => self.hasher.compute_hashes(&group_values[0])?, }; - Some(Arc::new(UInt64Array::from(hashes.to_vec())) as ArrayRef) + Some(Arc::new(hashes.clone()) as ArrayRef) } else { None }; diff --git a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs index 6323f2e87dffe..81b88635d2f14 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs @@ -188,8 +188,8 @@ 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))); + let hashes = self.output_hasher.compute_hashes(&cols[..1])?.clone(); + cols.push(Arc::new(hashes)); } Ok(cols) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index abb1059f0de0e..8b50e15cfb5d4 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -162,7 +162,7 @@ use crate::filter_pushdown::{ FilterPushdownPropagation, PushedDownPredicate, }; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; -use crate::repartition::ExpressionHasher; +use crate::repartition::{ExpressionHasher, take_hash_buffer}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, @@ -173,7 +173,7 @@ use datafusion_physical_expr::utils::collect_columns; use parking_lot::Mutex; use std::collections::{HashMap, HashSet}; -use arrow::array::{ArrayRef, UInt8Array, UInt16Array, UInt32Array, UInt64Array}; +use arrow::array::{Array, ArrayRef, UInt8Array, UInt16Array, UInt32Array, UInt64Array}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use arrow_schema::FieldRef; @@ -185,7 +185,6 @@ use datafusion_common::{ }; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryLimit; -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}; @@ -240,9 +239,16 @@ pub fn topk_types_supported(key_type: &DataType, value_type: &DataType) -> bool /// /// 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, + values: UInt64Array, +} + +impl Default for GroupHashTracker { + fn default() -> Self { + Self { + values: UInt64Array::from(Vec::::new()), + } + } } impl GroupHashTracker { @@ -251,34 +257,45 @@ impl GroupHashTracker { previous_group_count: usize, total_group_count: usize, group_indices: &[usize], - hashes: &[u64], + hashes: &UInt64Array, ) { assert_eq!(group_indices.len(), hashes.len()); assert_eq!(self.values.len(), previous_group_count); - self.values.resize(total_group_count, 0); + let mut values = take_hash_buffer(&mut self.values); + values.resize(total_group_count, 0); - for (&group_index, &hash) in group_indices.iter().zip(hashes) { + for (&group_index, &hash) in group_indices.iter().zip(hashes.values()) { if group_index >= previous_group_count { - self.values[group_index] = hash; + values[group_index] = hash; } } + self.values = UInt64Array::from(values); } 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)) + match emit_to { + EmitTo::All => { + let empty = UInt64Array::from(Vec::::new()); + Arc::new(std::mem::replace(&mut self.values, empty)) + } + EmitTo::First(n) => { + let mut values = take_hash_buffer(&mut self.values); + let hashes = split_vec_min_alloc(&mut values, n); + self.values = UInt64Array::from(values); + Arc::new(UInt64Array::from(hashes)) + } + } } fn allocated_size(&self) -> usize { - self.values.allocated_size() + self.values.get_buffer_memory_size() } fn clear_shrink(&mut self, capacity: usize) { - self.values.clear(); - self.values.shrink_to(capacity); + let mut values = take_hash_buffer(&mut self.values); + values.clear(); + values.shrink_to(capacity); + self.values = UInt64Array::from(values); } } @@ -2849,7 +2866,7 @@ mod tests { 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()))); + columns.push(Arc::new(hashes.clone())); Ok(RecordBatch::try_new(schema, columns)?) } @@ -2875,7 +2892,7 @@ mod tests { let hashes = hasher .precomputed_group_by(&group_by, &batch) .expect("batch should contain precomputed group hashes"); - assert_eq!(hashes, &[11, 22]); + assert_eq!(hashes.values(), &[11, 22]); Ok(()) } diff --git a/datafusion/physical-plan/src/repartition/hash.rs b/datafusion/physical-plan/src/repartition/hash.rs index 8a0d1c5d71c7c..09d0f55fa10cd 100644 --- a/datafusion/physical-plan/src/repartition/hash.rs +++ b/datafusion/physical-plan/src/repartition/hash.rs @@ -18,7 +18,7 @@ use crate::aggregates::PhysicalGroupBy; use crate::joins::SeededRandomState; use crate::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory}; -use arrow::array::{ArrayRef, RecordBatch}; +use arrow::array::{Array, ArrayRef, RecordBatch, UInt64Array}; use arrow::datatypes::Schema; use datafusion_common::Result; use datafusion_common::cast::as_uint64_array; @@ -68,7 +68,7 @@ impl HashMetrics { pub(crate) struct ExpressionHasher { hash_exprs: Vec>, - hash_buffer: Vec, + hash_buffer: UInt64Array, random_state: SeededRandomState, metrics: Option, } @@ -77,7 +77,7 @@ impl ExpressionHasher { pub(crate) fn new(hash_exprs: Vec>) -> Self { Self { hash_exprs, - hash_buffer: vec![], + hash_buffer: UInt64Array::from(Vec::::new()), random_state: SeededRandomState::with_seed(0), metrics: None, } @@ -137,21 +137,24 @@ impl ExpressionHasher { format!("{INTERNAL_HASH_COL_PREFIX}_{:016x}", hasher.finish()) } - pub(crate) fn precomputed<'a>(&self, batch: &'a RecordBatch) -> Option<&'a [u64]> { + pub(crate) fn precomputed<'a>( + &self, + batch: &'a RecordBatch, + ) -> Option<&'a UInt64Array> { 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()) + Some(hash_array) } pub(crate) fn precomputed_group_by<'a>( &self, group_by: &PhysicalGroupBy, batch: &'a RecordBatch, - ) -> Option<&'a [u64]> { + ) -> Option<&'a UInt64Array> { if !group_by.is_single() { return None; } @@ -159,16 +162,14 @@ impl ExpressionHasher { self.precomputed(batch) } - pub(crate) fn compute_hashes(&mut self, arrays: &[ArrayRef]) -> Result<&[u64]> { + pub(crate) fn compute_hashes(&mut self, arrays: &[ArrayRef]) -> Result<&UInt64Array> { let num_rows = arrays.first().map(|array| array.len()).unwrap_or(0); - self.hash_buffer.clear(); - self.hash_buffer.resize(num_rows, 0); + let mut hash_buffer = take_hash_buffer(&mut self.hash_buffer); + hash_buffer.clear(); + hash_buffer.resize(num_rows, 0); - create_hashes( - arrays, - self.random_state.random_state(), - &mut self.hash_buffer, - )?; + create_hashes(arrays, self.random_state.random_state(), &mut hash_buffer)?; + self.hash_buffer = UInt64Array::from(hash_buffer); if let Some(metrics) = &self.metrics { metrics.record_computed(num_rows); @@ -180,17 +181,26 @@ impl ExpressionHasher { pub(crate) fn compute_hashes_for_batch( &mut self, batch: &RecordBatch, - ) -> Result<&[u64]> { + ) -> Result<&UInt64Array> { 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() + self.hash_exprs.allocated_size() + self.hash_buffer.get_buffer_memory_size() } pub(crate) fn clear_shrink(&mut self, capacity: usize) { - self.hash_buffer.clear(); - self.hash_buffer.shrink_to(capacity); + let mut hash_buffer = take_hash_buffer(&mut self.hash_buffer); + hash_buffer.clear(); + hash_buffer.shrink_to(capacity); + self.hash_buffer = UInt64Array::from(hash_buffer); } } + +pub(crate) fn take_hash_buffer(hashes: &mut UInt64Array) -> Vec { + let empty = UInt64Array::from(Vec::::new()); + let (_, values, nulls) = std::mem::replace(hashes, empty).into_parts(); + debug_assert!(nulls.is_none()); + values.into() +} diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index fe3300ac551e8..10af8ac086fe1 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -46,7 +46,7 @@ use crate::{ check_if_same_properties, }; -use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions}; +use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array}; use arrow::compute::take_arrays; use arrow::datatypes::{SchemaRef, UInt32Type}; use arrow_schema::SortOptions; @@ -81,7 +81,9 @@ mod distributor_channels; mod hash; use crate::repartition::distributor_channels::SendError; -pub(crate) use crate::repartition::hash::{ExpressionHasher, HashMetrics}; +pub(crate) use crate::repartition::hash::{ + ExpressionHasher, HashMetrics, take_hash_buffer, +}; #[cfg(test)] pub(crate) use crate::repartition::hash::{HASH_ROWS_COMPUTED, HASH_ROWS_REUSED}; use distributor_channels::{ @@ -675,10 +677,10 @@ impl StrengthReducedU64 { } } - fn partition_indices(self, hash_buffer: &[u64], indices: &mut [Vec]) { + fn partition_indices(self, hash_buffer: &UInt64Array, indices: &mut [Vec]) { match self { Self::PowerOfTwo { mask } => { - for (index, hash) in hash_buffer.iter().enumerate() { + for (index, hash) in hash_buffer.values().iter().enumerate() { indices[(*hash & mask) as usize].push(index as u32); } } @@ -686,7 +688,7 @@ impl StrengthReducedU64 { divisor, reciprocal, } => { - for (index, hash) in hash_buffer.iter().enumerate() { + for (index, hash) in hash_buffer.values().iter().enumerate() { let quotient = Self::quotient(*hash, reciprocal); let partition = *hash - quotient * divisor; indices[partition as usize].push(index as u32);