diff --git a/datafusion/core/tests/execution/hash_reuse.rs b/datafusion/core/tests/execution/hash_reuse.rs new file mode 100644 index 0000000000000..343cbc619766c --- /dev/null +++ b/datafusion/core/tests/execution/hash_reuse.rs @@ -0,0 +1,334 @@ +// 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::ExecutionPlan; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode}; +use datafusion::physical_plan::collect; +use datafusion::physical_plan::display::DisplayableExecutionPlan; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use insta::assert_snapshot; +use std::sync::Arc; + +#[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(()) +} + +#[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(()) +} + +#[tokio::test] +async fn hash_output_is_configured_during_physical_planning() -> Result<()> { + let string_test = HashReuseTest::new( + r"SELECT l_returnflag, COUNT(*) AS row_count + FROM lineitem + GROUP BY l_returnflag", + ); + let (_, string_plan) = string_test.create_plan().await?; + let (partial_hashes, repartition_hashes) = hash_output_flags(&string_plan); + assert_eq!(partial_hashes, vec![true]); + assert_eq!(repartition_hashes, vec![true]); + + let primitive_test = HashReuseTest::new( + r"SELECT l_linenumber, COUNT(*) AS row_count + FROM lineitem + GROUP BY l_linenumber", + ); + let (_, primitive_plan) = primitive_test.create_plan().await?; + let (partial_hashes, repartition_hashes) = hash_output_flags(&primitive_plan); + assert_eq!(partial_hashes, vec![false]); + assert_eq!(repartition_hashes, vec![false]); + + Ok(()) +} + +fn hash_output_flags(plan: &Arc) -> (Vec, Vec) { + fn visit( + plan: &Arc, + partial_hashes: &mut Vec, + repartition_hashes: &mut Vec, + ) { + if let Some(aggregate) = plan.downcast_ref::() + && aggregate.mode() == &AggregateMode::Partial + { + partial_hashes.push(aggregate.emits_hashes()); + } + if let Some(repartition) = plan.downcast_ref::() { + repartition_hashes.push(repartition.emits_hashes()); + } + for child in plan.children() { + visit(child, partial_hashes, repartition_hashes); + } + } + + let mut partial_hashes = vec![]; + let mut repartition_hashes = vec![]; + visit(plan, &mut partial_hashes, &mut repartition_hashes); + (partial_hashes, repartition_hashes) +} + +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 create_plan(&self) -> Result<(SessionContext, Arc)> { + 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?; + Ok((ctx, plan)) + } + + async fn run(self) -> Result { + let (ctx, plan) = self.create_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/core/tests/physical_optimizer/partition_statistics.rs b/datafusion/core/tests/physical_optimizer/partition_statistics.rs index 6cabcdb710393..5d6672c9cd3b3 100644 --- a/datafusion/core/tests/physical_optimizer/partition_statistics.rs +++ b/datafusion/core/tests/physical_optimizer/partition_statistics.rs @@ -911,7 +911,6 @@ mod test { ColumnStatistics::new_unknown(), ], }; - assert_eq!(*p0_statistics, expected_p0_statistics); let expected_p1_statistics = Statistics { @@ -930,7 +929,6 @@ mod test { ColumnStatistics::new_unknown(), ], }; - let p1_statistics = StatisticsContext::new().compute( aggregate_exec_partial.as_ref(), &StatisticsArgs::new().with_partition(Some(1)), @@ -1009,7 +1007,6 @@ mod test { ColumnStatistics::new_unknown(), ], }; - assert_eq!( empty_stat, *StatisticsContext::new().compute( 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-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 952aae9846d0f..ddc8e2663f92e 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -52,7 +52,7 @@ use datafusion_physical_expr::{ }; use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::aggregates::{ - AggregateExec, AggregateMode, PhysicalGroupBy, + AggregateExec, AggregateMode, AggregateOutputMode, PhysicalGroupBy, }; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::execution_plan::EmissionType; @@ -728,6 +728,63 @@ fn partial_aggregate_output_satisfies_final_partitioning( .is_satisfied() } +/// Returns whether avoiding repeated hashing is expected to offset the cost of +/// propagating an additional `UInt64` column. +fn aggregate_hash_output_is_worthwhile(aggregate: &AggregateExec) -> Result { + let input_schema = aggregate.input().schema(); + for expr in aggregate.group_expr().input_exprs() { + if !expr.data_type(&input_schema)?.is_primitive() { + return Ok(true); + } + } + Ok(false) +} + +/// Configures the partial aggregate below distribution-preserving operators to +/// emit reusable group hashes. +/// +/// This runs while enforcing the final aggregate's distribution, when the +/// complete partial -> exchange -> final shape is known. Keeping the decision +/// here avoids runtime operators independently inferring whether a private hash +/// column should be part of their output. +fn configure_partial_aggregate_hash_output( + mut ctx: DistributionContext, +) -> Result { + if let Some(aggregate) = ctx.plan.downcast_ref::() { + if aggregate.mode() == &AggregateMode::Partial { + let emit_hashes = aggregate_hash_output_is_worthwhile(aggregate)?; + ctx.plan = Arc::new(aggregate.clone().with_hash_reuse(emit_hashes)?); + } + return Ok(ctx); + } + + if ctx.children.len() != 1 { + return Ok(ctx); + } + + if let Some(repartition) = ctx.plan.downcast_ref::() { + let child = configure_partial_aggregate_hash_output(ctx.children.swap_remove(0))?; + ctx.plan = Arc::new( + repartition + .clone() + .with_new_input(Arc::clone(&child.plan))? + .with_hash_reuse(true), + ); + ctx.children.push(child); + return Ok(ctx); + } + + if ctx.plan.is::() || ctx.plan.is::() + { + let child = configure_partial_aggregate_hash_output(ctx.children.swap_remove(0))?; + ctx.plan = + Arc::clone(&ctx.plan).with_new_children(vec![Arc::clone(&child.plan)])?; + ctx.children.push(child); + } + + Ok(ctx) +} + /// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator /// on top of the given plan node to satisfy a single partition requirement /// while preserving ordering constraints. @@ -1407,7 +1464,7 @@ pub fn ensure_distribution( target_partitions, )?; - let children = children + let mut children = children .into_iter() .map( |DistributionChildState { @@ -1483,6 +1540,16 @@ pub fn ensure_distribution( ) .collect::>>()?; + if plan + .downcast_ref::() + .is_some_and(|aggregate| { + matches!(aggregate.mode().output_mode(), AggregateOutputMode::Final) + }) + { + let child = children.swap_remove(0); + children.push(configure_partial_aggregate_hash_output(child)?); + } + let children_plans = children .iter() .map(|c| Arc::clone(&c.plan)) 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..e90d4c8aa45ff 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, ); }, @@ -432,9 +445,10 @@ fn bench_fixed_size_binary(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 42014f336f3d8..88d95df269bda 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, AggregateOutputMode, GroupHashTracker, PhysicalGroupBy, + aggregate_expressions, evaluate_group_by, }; +use crate::repartition::ExpressionHasher; /// Marker for raw rows -> partial state aggregation. pub(in crate::aggregates) struct PartialMarker; @@ -129,15 +131,21 @@ impl AggregateHashTable { let group_schema = agg.group_by.group_schema(&input_schema)?; let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + let hasher = agg.hashing_config().create_hasher(&agg.metrics, partition); + let should_output_hashes = + agg.emits_hashes() && 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: should_output_hashes.then(GroupHashTracker::default), + hasher, 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,14 @@ 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. + /// 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. + 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 2293e7b1b8e89..d446ff87ad15a 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, AggregateOutputMode, GroupHashTracker, PhysicalGroupBy, + aggregate_expressions, evaluate_group_by, }; +use crate::repartition::ExpressionHasher; use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; @@ -118,6 +119,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 @@ -136,6 +143,7 @@ impl OrderedAggregateTable { input_schema: &SchemaRef, output_schema: SchemaRef, state_schema: SchemaRef, + partition: usize, batch_size: usize, input_order_mode: &InputOrderMode, aggregate_mode: &AggregateMode, @@ -171,6 +179,10 @@ impl OrderedAggregateTable { }) .collect::>()?; + let hasher = agg.hashing_config().create_hasher(&agg.metrics, partition); + let should_output_hashes = agg.emits_hashes() + && aggregate_mode.output_mode() == AggregateOutputMode::Partial; + Ok(Self { output_schema, state_schema, @@ -181,6 +193,8 @@ impl OrderedAggregateTable { group_ordering, group_values, group_indices: vec![], + group_hash_tracker: should_output_hashes.then(GroupHashTracker::default), + hasher, accumulators, }, _mode: PhantomData, @@ -191,10 +205,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); @@ -211,6 +225,10 @@ impl OrderedAggregateTable { Ok(EvaluatedAggregateBatch { grouping_set_args, accumulator_args, + precomputed_group_hashes: self + .buffer + .hasher + .precomputed_group_by(&self.buffer.group_by, batch), }) } @@ -248,6 +266,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() } pub(in crate::aggregates) fn group_by_metrics(&self) -> GroupByMetrics { @@ -314,15 +338,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, @@ -400,6 +438,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/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs index fd064ebffec12..c3b3905d9aeff 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs @@ -46,6 +46,7 @@ impl OrderedAggregateTable { agg: &AggregateExec, input_schema: &SchemaRef, output_schema: SchemaRef, + partition: usize, batch_size: usize, input_order_mode: &InputOrderMode, group_by_metrics: GroupByMetrics, @@ -55,6 +56,7 @@ impl OrderedAggregateTable { input_schema, output_schema, Arc::clone(input_schema), + partition, batch_size, input_order_mode, &AggregateMode::Final, diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs index a04e4dda8fb39..8dfbb9c916b81 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs @@ -66,6 +66,7 @@ impl OrderedAggregateTable { &input_schema, output_schema, state_schema, + partition, batch_size, &agg.input_order_mode, &AggregateMode::Partial, 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 4bcacb49afb04..0b876f432d33b 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, @@ -92,6 +94,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, @@ -159,9 +166,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; } @@ -194,13 +211,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() @@ -208,6 +237,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 1101d535311e4..08464feba0c11 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`]. /// 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..a730d12cfad41 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 { @@ -1264,10 +1248,12 @@ mod tests { use arrow::{compute::concat_batches, util::pretty::pretty_format_batches}; use datafusion_common::utils::proxy::HashTableAllocExt; use datafusion_expr::EmitTo; + use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet; use crate::aggregates::group_values::{ GroupValues, multi_group_by::GroupValuesColumn, }; + use crate::repartition::ExpressionHasher; use super::{ GroupIndexView, group_column_supported_type, make_group_column, supported_schema, @@ -1877,8 +1863,11 @@ mod tests { } fn load_to_group_values(&self, group_values: &mut impl GroupValues) { + let metrics = ExecutionPlanMetricsSet::new(); + let mut hasher = ExpressionHasher::new(vec![], &metrics, 0); 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..55ff00285aba2 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 @@ -18,11 +18,13 @@ use std::mem::size_of; use crate::aggregates::group_values::GroupValues; +use crate::repartition::ExpressionHasher; use arrow::array::{Array, ArrayRef, OffsetSizeTrait}; use datafusion_common::Result; use datafusion_expr::EmitTo; use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; +use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet; /// A [`GroupValues`] storing single column of Utf8/LargeUtf8/Binary/LargeBinary values /// @@ -45,15 +47,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 +116,11 @@ 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 metrics = ExecutionPlanMetricsSet::new(); + let mut hasher = ExpressionHasher::new(vec![], &metrics, 0); + 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..0250d0ed8f858 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,10 +16,12 @@ // under the License. use crate::aggregates::group_values::GroupValues; +use crate::repartition::ExpressionHasher; use arrow::array::{Array, ArrayRef}; use datafusion_expr::EmitTo; use datafusion_physical_expr::binary_map::OutputType; use datafusion_physical_expr_common::binary_view_map::ArrowBytesViewMap; +use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet; use std::mem::size_of; /// A [`GroupValues`] storing single column of Utf8View/BinaryView values @@ -47,6 +49,7 @@ impl GroupValues for GroupValuesBytesView { &mut self, cols: &[ArrayRef], groups: &mut Vec, + hashes: &[u64], ) -> datafusion_common::Result<()> { assert_eq!(cols.len(), 1); @@ -54,8 +57,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 +114,11 @@ 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 metrics = ExecutionPlanMetricsSet::new(); + let mut hasher = ExpressionHasher::new(vec![], &metrics, 0); + 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..320bad969d4ae 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 { @@ -250,10 +251,12 @@ where #[cfg(test)] mod tests { use super::*; + use crate::repartition::ExpressionHasher; use arrow::array::types::Int32Type; use arrow::array::{ArrayRef, Int32Array}; use arrow::datatypes::DataType; use datafusion_expr::EmitTo; + use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet; use std::sync::Arc; /// Mirror of the `EmitTo::take_needed` regression test, applied to the @@ -273,7 +276,11 @@ 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 metrics = ExecutionPlanMetricsSet::new(); + let mut hasher = ExpressionHasher::new(vec![], &metrics, 0); + 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 99c101199459f..6ddfff9b31a4d 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; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::spill_manager::{GetSlicedSize, SpillManager}; use crate::stream::EmptyRecordBatchStream; @@ -333,6 +334,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)`, @@ -444,6 +451,7 @@ impl GroupedHashAggregateStream { &agg_group_by, &aggregate_exprs, AggregateMode::Partial, + agg.hashing_config(), )?); // Need to update the GROUP BY expressions to point to the correct column after schema change @@ -582,6 +590,13 @@ impl GroupedHashAggregateStream { None }; + let hasher = agg.hashing_config().create_hasher(&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 = agg.emits_hashes() + && (agg.mode.output_mode() == AggregateOutputMode::Partial + || oom_mode == OutOfMemoryMode::Spill); + Ok(GroupedHashAggregateStream { schema: agg_schema, input_schema: agg.input().schema(), @@ -595,6 +610,8 @@ impl GroupedHashAggregateStream { oom_mode, group_values, current_group_indices: Default::default(), + group_hash_tracker: should_output_hashes.then(GroupHashTracker::default), + hasher, exec_state, baseline_metrics, group_by_metrics, @@ -872,17 +889,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, @@ -985,7 +1024,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 @@ -1042,6 +1086,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 @@ -1097,9 +1147,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, @@ -1232,6 +1291,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. @@ -1283,6 +1346,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)) @@ -1366,7 +1432,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)?; @@ -1376,6 +1442,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 @@ -1388,6 +1463,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..d24222e99ce5c 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs @@ -22,10 +22,11 @@ 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, + AggregateExec, AggregateOutputMode, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, evaluate_many, }; use crate::metrics::BaselineMetrics; +use crate::repartition::ExpressionHasher; 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,12 @@ 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 = aggr.emits_hashes() + && aggr.mode.output_mode() == AggregateOutputMode::Partial; + // aggr.hashing_config was created based on group_by.input_exprs(), but we need an + // ExpressionHasher created based on group_by.output_exprs(), so create on from scratch. + let output_hasher = + ExpressionHasher::new(group_by.output_exprs(), &aggr.metrics, partition); Ok(GroupedTopKAggregateStream { partition, @@ -122,6 +131,8 @@ impl GroupedTopKAggregateStream { aggregate_arguments, group_by, priority_map, + output_hasher, + output_group_hashes, null_group_seen: false, }) } @@ -175,6 +186,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/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index e7f0f075b33a5..9e698ea5f7103 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, @@ -455,7 +455,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)), }, ); } @@ -545,9 +545,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, .. @@ -571,7 +571,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, .. diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 787cd4f03ff6c..0fdef5ff6bd45 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -162,6 +162,9 @@ use crate::filter_pushdown::{ FilterPushdownPropagation, PushedDownPredicate, }; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +#[cfg(test)] +use crate::repartition::ExpressionHasher; +use crate::repartition::HashingConfig; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, @@ -177,13 +180,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 +200,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 +231,58 @@ 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); +/// 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. @@ -856,6 +909,9 @@ pub struct AggregateExec { required_input_ordering: Option, /// Describes how the input is ordered relative to the group by columns input_order_mode: InputOrderMode, + /// Immutable planning-time configuration for group hashing and private + /// hash-column output. + hashing_config: HashingConfig, cache: Arc, /// During initialization, if the plan supports dynamic filtering (see [`AggrDynFilter`]), /// it is set to `Some(..)` regardless of whether it can be pushed down to a child node. @@ -880,6 +936,7 @@ impl AggregateExec { required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), + hashing_config: self.hashing_config.clone(), cache: Arc::clone(&self.cache), mode: self.mode, group_by: Arc::clone(&self.group_by), @@ -900,6 +957,7 @@ impl AggregateExec { required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), + hashing_config: self.hashing_config.clone(), cache: Arc::clone(&self.cache), mode: self.mode, group_by: Arc::clone(&self.group_by), @@ -926,7 +984,14 @@ impl AggregateExec { input_schema: SchemaRef, ) -> Result { let group_by = group_by.into(); - let schema = create_schema(&input.schema(), &group_by, &aggr_expr, mode)?; + let hashing_config = HashingConfig::new(group_by.input_exprs()); + let schema = create_schema( + &input.schema(), + &group_by, + &aggr_expr, + mode, + &hashing_config, + )?; let schema = Arc::new(schema); AggregateExec::try_new_with_schema( @@ -937,6 +1002,7 @@ impl AggregateExec { input, input_schema, schema, + hashing_config, ) } @@ -948,6 +1014,10 @@ impl AggregateExec { /// a rule may re-write aggregate expressions (e.g. reverse them) during /// initialization, field names may change inadvertently if one re-creates /// the schema in such cases. + #[expect( + clippy::too_many_arguments, + reason = "The hashing configuration must be preserved alongside the explicit schema" + )] fn try_new_with_schema( mode: AggregateMode, group_by: impl Into>, @@ -956,6 +1026,7 @@ impl AggregateExec { input: Arc, input_schema: SchemaRef, schema: SchemaRef, + hashing_config: HashingConfig, ) -> Result { let group_by = group_by.into(); let filter_expr = filter_expr.into(); @@ -1042,6 +1113,7 @@ impl AggregateExec { required_input_ordering, limit_options: None, input_order_mode, + hashing_config, cache: Arc::new(cache), dynamic_filter: None, }; @@ -1056,6 +1128,45 @@ impl AggregateExec { &self.mode } + /// Configures whether this aggregate should emit reusable group hashes. + /// + /// Physical planning calls this after distribution enforcement has + /// determined that a downstream operator can consume the hashes. + pub fn with_hash_reuse(mut self, enabled: bool) -> Result { + self.hashing_config = self.hashing_config.with_hash_output(enabled); + // Schema changes, because now this node will output one new extra internal column + // containing the hashes, so both schema and PlanProperties need to be recomputed. + self.schema = Arc::new(create_schema( + &self.input.schema(), + &self.group_by, + &self.aggr_expr, + self.mode, + &self.hashing_config, + )?); + + let group_expr_mapping = + ProjectionMapping::try_new(self.group_by.expr.clone(), &self.input.schema())?; + self.cache = Arc::new(Self::compute_properties( + &self.input, + Arc::clone(&self.schema), + &group_expr_mapping, + self.group_by.is_true_no_grouping(), + &self.mode, + &self.input_order_mode, + self.aggr_expr.as_ref(), + )?); + Ok(self) + } + + pub(crate) fn hashing_config(&self) -> &HashingConfig { + &self.hashing_config + } + + /// Returns whether this plan emits group hashes in a private output column. + pub fn emits_hashes(&self) -> bool { + self.hashing_config.should_output_hashes() + } + /// Set the limit options for this AggExec pub fn with_limit_options(mut self, limit_options: Option) -> Self { self.limit_options = limit_options; @@ -1975,6 +2086,7 @@ impl ExecutionPlan for AggregateExec { Arc::clone(&children[0]), Arc::clone(&self.input_schema), Arc::clone(&self.schema), + self.hashing_config.clone(), )?; me.limit_options = self.limit_options; me.dynamic_filter.clone_from(&self.dynamic_filter); @@ -2252,6 +2364,7 @@ impl ExecutionPlan for AggregateExec { limit, has_grouping_set: group_by.has_grouping_set(), dynamic_filter, + emit_hashes: self.emits_hashes(), }, )), ), @@ -2517,7 +2630,8 @@ impl AggregateExec { filter_expr, input, Arc::clone(&input_schema), - )?; + )? + .with_hash_reuse(hash_agg.emit_hashes)?; let aggregate = if let Some(limit) = &hash_agg.limit { let options = match limit.descending { Some(descending) => { @@ -2556,8 +2670,13 @@ fn create_schema( group_by: &PhysicalGroupBy, aggr_expr: &[Arc], mode: AggregateMode, + hashing_config: &HashingConfig, ) -> Result { - let mut fields = Vec::with_capacity(group_by.num_output_exprs() + aggr_expr.len()); + let output_hashes = hashing_config.should_output_hashes() + && mode.output_mode() == AggregateOutputMode::Partial; + let mut fields = Vec::with_capacity( + group_by.num_output_exprs() + aggr_expr.len() + output_hashes as usize, + ); fields.extend(group_by.output_fields(input_schema)?); match mode.output_mode() { @@ -2575,6 +2694,11 @@ fn create_schema( } } + if output_hashes { + let output_hashing_config = HashingConfig::new(group_by.output_exprs()); + fields.push(output_hashing_config.hash_field()); + } + Ok(Schema::new_with_metadata( fields, input_schema.metadata().clone(), @@ -3055,7 +3179,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; @@ -3083,6 +3207,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}; @@ -3121,6 +3246,89 @@ mod tests { Ok(schema) } + fn assert_and_strip_group_hashes( + batches: &[RecordBatch], + hash_exprs: &[Arc], + ) -> Result> { + let metrics = ExecutionPlanMetricsSet::new(); + let mut hasher = ExpressionHasher::new(hash_exprs.to_vec(), &metrics, 0); + let hash_name = hasher.internal_hash_col_name(); + 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"); + let arrays = evaluate_expressions_to_arrays(hash_exprs, batch)?; + let expected = hasher.compute_hashes(&arrays)?; + assert_eq!(actual, expected); + + 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 metrics = ExecutionPlanMetricsSet::new(); + let mut hasher = ExpressionHasher::new(hash_exprs.clone(), &metrics, 0); + 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() + .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 metrics = ExecutionPlanMetricsSet::new(); + let hasher = ExpressionHasher::new(hash_exprs, &metrics, 0); + 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. @@ -3244,6 +3452,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, @@ -3292,6 +3506,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, @@ -3442,6 +3658,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! { @@ -3907,6 +4125,137 @@ 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::Utf8, false), + Field::new("value", DataType::Int32, false), + ])); + let partitions = vec![ + vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![ + 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(StringArray::from(vec!["b", "c"])), + 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), + )? + .with_hash_reuse(true)?, + ); + 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), + )? + .with_hash_reuse(true), + ); + 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), + )? + .with_hash_reuse(true)?, + ); + let repartition = Arc::new( + RepartitionExec::try_new( + Arc::clone(&partial) as Arc, + Partitioning::Hash(hash_exprs, 2), + )? + .with_hash_reuse(true), + ); + 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) | ++-----+--------------+ +| a | 2 | +| b | 2 | +| c | 1 | ++-----+--------------+ +"); + + Ok(()) + } + #[tokio::test] async fn partial_grouped_aggregate_materializes_before_slicing() -> Result<()> { let schema = Arc::new(Schema::new(vec![ @@ -3924,6 +4273,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)?]) @@ -3964,6 +4314,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] | @@ -4028,6 +4379,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 | @@ -4173,8 +4526,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])), @@ -4280,6 +4634,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)) @@ -4305,6 +4660,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] | @@ -4351,8 +4707,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])), @@ -5162,6 +5519,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)?]) @@ -5216,6 +5574,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" @@ -5246,6 +5605,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)?]) @@ -5308,6 +5668,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" @@ -5337,6 +5698,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)?]) @@ -5389,6 +5751,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" @@ -5424,6 +5787,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)?]) @@ -5484,6 +5848,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" @@ -5625,11 +5990,13 @@ mod tests { ], true, ); + let hashing_config = HashingConfig::new(grouping_set.input_exprs()); let aggr_schema = create_schema( &input_schema, &grouping_set, &aggr_expr, AggregateMode::Final, + &hashing_config, )?; let expected_schema = Schema::new(vec![ Field::new("a", DataType::Float32, false), diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 071c1e9011f41..b42a38889ecd6 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -335,6 +335,7 @@ impl OrderedFinalAggregateStream { agg, &input_schema, Arc::clone(&schema), + partition, batch_size, input_order_mode, group_by_metrics, diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 34493a5f51742..cc7ee9a946602 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -866,6 +866,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(); diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 00df227cb87db..79e0093e0a395 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}; @@ -202,12 +204,23 @@ impl ExecutionPlan for RecursiveQueryExec { let static_stream = self.static_term.execute(partition, Arc::clone(&context))?; let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + let distinct_deduplicator = self + .is_distinct + .then(|| { + DistinctDeduplicator::new( + static_stream.schema(), + &context, + &self.metrics, + partition, + ) + }) + .transpose()?; Ok(Box::pin(RecursiveQueryStream::new( context, Arc::clone(&self.work_table), Arc::clone(&self.recursive_term), static_stream, - self.is_distinct, + distinct_deduplicator, baseline_metrics, )?)) } @@ -289,15 +302,12 @@ impl RecursiveQueryStream { work_table: Arc, recursive_term: Arc, static_stream: SendableRecordBatchStream, - is_distinct: bool, + distinct_deduplicator: Option, baseline_metrics: BaselineMetrics, ) -> Result { let schema = static_stream.schema(); let reservation = MemoryConsumer::new("RecursiveQuery").register(task_context.memory_pool()); - let distinct_deduplicator = is_distinct - .then(|| DistinctDeduplicator::new(Arc::clone(&schema), &task_context)) - .transpose()?; Ok(Self { task_context, work_table, @@ -440,10 +450,24 @@ struct DistinctDeduplicator { group_values: Box, reservation: MemoryReservation, intern_output_buffer: Vec, + hasher: ExpressionHasher, } impl DistinctDeduplicator { - fn new(schema: SchemaRef, task_context: &TaskContext) -> Result { + fn new( + schema: SchemaRef, + task_context: &TaskContext, + metrics: &ExecutionPlanMetricsSet, + partition: usize, + ) -> 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 +475,7 @@ impl DistinctDeduplicator { group_values, reservation, intern_output_buffer: Vec::new(), + hasher: ExpressionHasher::new(hash_exprs, metrics, partition), }) } @@ -470,8 +495,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..96b892eefe0a5 --- /dev/null +++ b/datafusion/physical-plan/src/repartition/hash.rs @@ -0,0 +1,229 @@ +// 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 arrow::datatypes::{DataType, Field, FieldRef, Schema}; +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"; + +/// Immutable planning-time configuration for expression hashing. +/// +/// This configuration determines whether a plan node exposes an internal hash +/// column and constructs the mutable [`ExpressionHasher`] used by each runtime +/// stream. It deliberately contains no hash buffer, random state, or metrics. +#[derive(Debug, Clone)] +pub(crate) struct HashingConfig { + hash_exprs: Vec>, + emit_hashes: bool, +} + +impl HashingConfig { + pub(crate) fn new(hash_exprs: Vec>) -> Self { + Self { + hash_exprs, + emit_hashes: false, + } + } + + /// Records the hash-output decision made during physical planning. + pub(crate) fn with_hash_output(mut self, enabled: bool) -> Self { + self.emit_hashes = enabled; + self + } + + pub(crate) fn should_output_hashes(&self) -> bool { + self.emit_hashes + } + + pub(crate) fn internal_hash_col_name(&self) -> String { + internal_hash_col_name(&self.hash_exprs) + } + + pub(crate) fn hash_field(&self) -> FieldRef { + Field::new(self.internal_hash_col_name(), DataType::UInt64, false).into() + } + + pub(crate) fn input_has_hash_column(&self, input_schema: &Schema) -> bool { + input_schema + .field_with_name(&self.internal_hash_col_name()) + .is_ok() + } + + pub(crate) fn create_hasher( + &self, + metrics: &ExecutionPlanMetricsSet, + partition: usize, + ) -> ExpressionHasher { + ExpressionHasher::new(self.hash_exprs.clone(), metrics, partition) + } +} + +/// 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); + } +} + +#[derive(Debug)] +pub(crate) struct ExpressionHasher { + hash_exprs: Vec>, + hash_buffer: Vec, + random_state: SeededRandomState, + metrics: HashMetrics, +} + +impl Clone for ExpressionHasher { + fn clone(&self) -> Self { + self.new_for_exprs(self.hash_exprs.clone()) + } +} + +impl ExpressionHasher { + pub(crate) fn new( + hash_exprs: Vec>, + metrics: &ExecutionPlanMetricsSet, + partition: usize, + ) -> Self { + Self { + hash_exprs, + hash_buffer: vec![], + random_state: SeededRandomState::with_seed(0), + metrics: HashMetrics::new(metrics, partition), + } + } + + /// Creates a hasher for another expression list while preserving metrics. + pub(crate) fn new_for_exprs(&self, hash_exprs: Vec>) -> Self { + Self { + hash_exprs, + hash_buffer: vec![], + random_state: SeededRandomState::with_seed(0), + metrics: self.metrics.clone(), + } + } + + /// 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 { + internal_hash_col_name(&self.hash_exprs) + } + + 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()?; + self.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, + )?; + + self.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); + } +} + +fn internal_hash_col_name(hash_exprs: &[Arc]) -> String { + let mut hasher = DefaultHasher::new(); + hash_exprs.hash(&mut hasher); + format!("{INTERNAL_HASH_COL_PREFIX}_{:016x}", hasher.finish()) +} diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 3473aad9b3fc0..293c6f690e702 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, HashingConfig}; +#[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,16 +410,18 @@ 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 partition_metrics = + RepartitionMetrics::new(i, output_partitions, metrics); - let timer = metrics.fetch_time.timer(); + let timer = partition_metrics.fetch_time.timer(); let stream = input.execute(i, Arc::clone(ctx))?; timer.done(); - streams_and_metrics.push((stream, metrics)); + streams_and_metrics.push((stream, partition_metrics)); } *self = RepartitionExecState::InputStreamsInitialized(streams_and_metrics); Ok(()) @@ -427,6 +433,7 @@ impl RepartitionExecState { input: &Arc, metrics: &ExecutionPlanMetricsSet, partitioning: &Partitioning, + hashing_config: &HashingConfig, preserve_order: bool, name: &str, context: &Arc, @@ -437,7 +444,7 @@ impl RepartitionExecState { self.ensure_input_streams_initialized( input, metrics, - partitioning.partition_count(), + partitioning, context, )?; let RepartitionExecState::InputStreamsInitialized(value) = self else { @@ -548,7 +555,7 @@ impl RepartitionExecState { // launch one async task per *input* partition let mut spawned_tasks = Vec::with_capacity(num_input_partitions); - for (i, (stream, metrics)) in + for (i, (stream, partition_metrics)) in std::mem::take(streams_and_metrics).into_iter().enumerate() { let txs: HashMap<_, _> = channels @@ -579,7 +586,8 @@ impl RepartitionExecState { stream, txs, partitioning.clone(), - metrics, + hashing_config.clone(), + partition_metrics, // preserve_order depends on partition index to start from 0 if preserve_order { 0 } else { i }, num_input_partitions, @@ -610,9 +618,8 @@ pub struct BatchPartitioner { enum BatchPartitionerState { Hash { - exprs: Vec>, partition_reducer: StrengthReducedU64, - hash_buffer: Vec, + hasher: ExpressionHasher, indices: Vec>, }, RoundRobin { @@ -726,6 +733,23 @@ impl BatchPartitioner { exprs: Vec>, num_partitions: usize, timer: metrics::Time, + ) -> Result { + let metrics = ExecutionPlanMetricsSet::new(); + Self::new_from_hashing_config( + &HashingConfig::new(exprs), + num_partitions, + timer, + &metrics, + 0, + ) + } + + fn new_from_hashing_config( + hashing_config: &HashingConfig, + num_partitions: usize, + timer: metrics::Time, + metrics: &ExecutionPlanMetricsSet, + partition: usize, ) -> Result { if num_partitions == 0 { return internal_err!("Hash repartition requires at least one partition"); @@ -733,9 +757,8 @@ impl BatchPartitioner { Ok(Self { state: BatchPartitionerState::Hash { - exprs, partition_reducer: StrengthReducedU64::new(num_partitions as u64), - hash_buffer: vec![], + hasher: hashing_config.create_hasher(metrics, partition), indices: vec![vec![]; num_partitions], }, timer, @@ -835,6 +858,39 @@ impl BatchPartitioner { } } + fn try_new_with_hashing_config( + partitioning: Partitioning, + timer: metrics::Time, + input_partition: usize, + num_input_partitions: usize, + hashing_config: &HashingConfig, + metrics: &ExecutionPlanMetricsSet, + ) -> Result { + match partitioning { + Partitioning::Hash(_, num_partitions) => Self::new_from_hashing_config( + hashing_config, + num_partitions, + timer, + metrics, + input_partition, + ), + Partitioning::RoundRobinBatch(num_partitions) => { + Ok(Self::new_round_robin_partitioner( + num_partitions, + timer, + input_partition, + num_input_partitions, + )) + } + Partitioning::Range(range_repartitioning) => { + Ok(Self::new_range_partitioner(&range_repartitioning, timer)) + } + other => { + not_impl_err!("Unsupported repartitioning scheme {other:?}") + } + } + } + /// Partition the provided [`RecordBatch`] into one or more partitioned [`RecordBatch`] /// based on the [`Partitioning`] specified on construction /// @@ -883,28 +939,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 @@ -1204,12 +1252,17 @@ pub struct RepartitionExec { /// Boolean flag to decide whether to preserve ordering. If true means /// `SortPreservingRepartitionExec`, false means `RepartitionExec`. preserve_order: bool, + /// Immutable planning-time configuration for partition hashing and private + /// hash-column output. + hashing_config: HashingConfig, /// Cache holding plan properties like equivalences, output partitioning etc. cache: Arc, } #[derive(Debug, Clone)] struct RepartitionMetrics { + /// Registry shared by all metrics for this operator. + execution_metrics: ExecutionPlanMetricsSet, /// Time in nanos to execute child operator and fetch batches fetch_time: metrics::Time, /// Repartitioning elapsed time in nanos @@ -1246,6 +1299,7 @@ impl RepartitionMetrics { .collect(); Self { + execution_metrics: metrics.clone(), fetch_time, repartition_time, send_time, @@ -1342,14 +1396,9 @@ impl ExecutionPlan for RepartitionExec { mut children: Vec>, ) -> Result> { check_if_same_properties!(self, children); - let mut repartition = RepartitionExec::try_new( - children.swap_remove(0), - self.partitioning().clone(), - )?; - if self.preserve_order { - repartition = repartition.with_preserve_order(); - } - Ok(Arc::new(repartition)) + Ok(Arc::new( + (*self).clone().with_new_input(children.swap_remove(0))?, + )) } fn with_new_children_and_same_properties( @@ -1387,6 +1436,7 @@ impl ExecutionPlan for RepartitionExec { let input = Arc::clone(&self.input); let partitioning = self.partitioning().clone(); + let hashing_config = self.hashing_config.clone(); let metrics = self.metrics.clone(); let preserve_order = self.sort_exprs().is_some(); let name = self.name().to_owned(); @@ -1407,7 +1457,7 @@ impl ExecutionPlan for RepartitionExec { state.ensure_input_streams_initialized( &input, &metrics, - partitioning.partition_count(), + &partitioning, &context, )?; } @@ -1423,6 +1473,7 @@ impl ExecutionPlan for RepartitionExec { &input, &metrics, &partitioning, + &hashing_config, preserve_order, &name, &context, @@ -1672,6 +1723,7 @@ impl ExecutionPlan for RepartitionExec { self.input.try_pushdown_sort(order)?.try_map(|new_input| { let mut new_repartition = RepartitionExec::try_new(new_input, self.partitioning().clone())?; + new_repartition = new_repartition.with_hash_reuse(self.emits_hashes()); if self.preserve_order { new_repartition = new_repartition.with_preserve_order(); } @@ -1700,6 +1752,7 @@ impl ExecutionPlan for RepartitionExec { state: Arc::clone(&self.state), metrics: self.metrics.clone(), preserve_order: self.preserve_order, + hashing_config: self.hashing_config.clone(), cache: new_properties.into(), }))) } @@ -1772,6 +1825,7 @@ impl ExecutionPlan for RepartitionExec { partition_method: Some(partition_method), }), preserve_order: self.preserve_order(), + emit_hashes: self.emits_hashes(), }, )), ), @@ -1878,7 +1932,8 @@ impl RepartitionExec { } }; - let mut repart_exec = RepartitionExec::try_new(input, partitioning)?; + let mut repart_exec = RepartitionExec::try_new(input, partitioning)? + .with_hash_reuse(repart.emit_hashes); if repart.preserve_order { repart_exec = repart_exec.with_preserve_order(); } @@ -1895,16 +1950,55 @@ impl RepartitionExec { partitioning: Partitioning, ) -> Result { let preserve_order = false; + let hash_exprs = match &partitioning { + Partitioning::Hash(exprs, _) => exprs.clone(), + _ => vec![], + }; let cache = Self::compute_properties(&input, partitioning, preserve_order); Ok(RepartitionExec { input, state: Default::default(), metrics: ExecutionPlanMetricsSet::new(), preserve_order, + hashing_config: HashingConfig::new(hash_exprs), cache: Arc::new(cache), }) } + /// Configures whether this repartition should propagate reusable hashes. + /// + /// Enabling this only takes effect when the input schema contains the + /// internal hash column for this repartition's expressions. Repartition + /// therefore never manufactures a private output column that planning did + /// not arrange upstream. + pub fn with_hash_reuse(mut self, enabled: bool) -> Self { + let hash_column_available = self + .hashing_config + .input_has_hash_column(&self.input.schema()); + self.hashing_config = self + .hashing_config + .with_hash_output(enabled && hash_column_available); + self + } + + /// Replaces this repartition's input while preserving its planned + /// partitioning, ordering, and hash-output configuration. + pub fn with_new_input(self, input: Arc) -> Result { + let emit_hashes = self.emits_hashes(); + let preserve_order = self.preserve_order; + let mut repartition = Self::try_new(input, self.partitioning().clone())? + .with_hash_reuse(emit_hashes); + if preserve_order { + repartition = repartition.with_preserve_order(); + } + Ok(repartition) + } + + /// Returns whether this plan propagates hashes in a private output column. + pub fn emits_hashes(&self) -> bool { + self.hashing_config.should_output_hashes() + } + fn maintains_input_order_helper( input: &Arc, preserve_order: bool, @@ -1983,15 +2077,18 @@ impl RepartitionExec { mut stream: SendableRecordBatchStream, mut output_channels: HashMap, partitioning: Partitioning, + hashing_config: HashingConfig, metrics: RepartitionMetrics, input_partition: usize, num_input_partitions: usize, ) -> Result<()> { - let mut partitioner = BatchPartitioner::try_new( + let mut partitioner = BatchPartitioner::try_new_with_hashing_config( partitioning, metrics.repartition_time.clone(), input_partition, num_input_partitions, + &hashing_config, + &metrics.execution_metrics, )?; // While there are still outputs to send to, keep pulling inputs diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 16b1b1532f518..3fc481801c04e 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1446,6 +1446,8 @@ message AggregateExecNode { bool has_grouping_set = 12; // Optional dynamic filter expression for pushing down to the child. PhysicalExprNode dynamic_filter = 13; + // Whether the partial aggregate should emit reusable group hashes. + bool emit_hashes = 14; } message GlobalLimitExecNode { @@ -1520,6 +1522,8 @@ message RepartitionExecNode{ // New partitioning variants are stored in `partitioning`. Partitioning partitioning = 5; bool preserve_order = 6; + // Whether hash repartitioning should propagate reusable group hashes. + bool emit_hashes = 7; } message Partitioning { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index c5d7c003013a1..ded75296f3ef0 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -157,6 +157,9 @@ impl serde::Serialize for AggregateExecNode { if self.dynamic_filter.is_some() { len += 1; } + if self.emit_hashes { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.AggregateExecNode", len)?; if !self.group_expr.is_empty() { struct_ser.serialize_field("groupExpr", &self.group_expr)?; @@ -199,6 +202,9 @@ impl serde::Serialize for AggregateExecNode { if let Some(v) = self.dynamic_filter.as_ref() { struct_ser.serialize_field("dynamicFilter", v)?; } + if self.emit_hashes { + struct_ser.serialize_field("emitHashes", &self.emit_hashes)?; + } struct_ser.end() } } @@ -231,6 +237,8 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode { "hasGroupingSet", "dynamic_filter", "dynamicFilter", + "emit_hashes", + "emitHashes", ]; #[allow(clippy::enum_variant_names)] @@ -248,6 +256,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode { Limit, HasGroupingSet, DynamicFilter, + EmitHashes, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -282,6 +291,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode { "limit" => Ok(GeneratedField::Limit), "hasGroupingSet" | "has_grouping_set" => Ok(GeneratedField::HasGroupingSet), "dynamicFilter" | "dynamic_filter" => Ok(GeneratedField::DynamicFilter), + "emitHashes" | "emit_hashes" => Ok(GeneratedField::EmitHashes), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -314,6 +324,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode { let mut limit__ = None; let mut has_grouping_set__ = None; let mut dynamic_filter__ = None; + let mut emit_hashes__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::GroupExpr => { @@ -394,6 +405,12 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode { } dynamic_filter__ = map_.next_value()?; } + GeneratedField::EmitHashes => { + if emit_hashes__.is_some() { + return Err(serde::de::Error::duplicate_field("emitHashes")); + } + emit_hashes__ = Some(map_.next_value()?); + } } } Ok(AggregateExecNode { @@ -410,6 +427,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode { limit: limit__, has_grouping_set: has_grouping_set__.unwrap_or_default(), dynamic_filter: dynamic_filter__, + emit_hashes: emit_hashes__.unwrap_or_default(), }) } } @@ -23674,6 +23692,9 @@ impl serde::Serialize for RepartitionExecNode { if self.preserve_order { len += 1; } + if self.emit_hashes { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.RepartitionExecNode", len)?; if let Some(v) = self.input.as_ref() { struct_ser.serialize_field("input", v)?; @@ -23684,6 +23705,9 @@ impl serde::Serialize for RepartitionExecNode { if self.preserve_order { struct_ser.serialize_field("preserveOrder", &self.preserve_order)?; } + if self.emit_hashes { + struct_ser.serialize_field("emitHashes", &self.emit_hashes)?; + } struct_ser.end() } } @@ -23698,6 +23722,8 @@ impl<'de> serde::Deserialize<'de> for RepartitionExecNode { "partitioning", "preserve_order", "preserveOrder", + "emit_hashes", + "emitHashes", ]; #[allow(clippy::enum_variant_names)] @@ -23705,6 +23731,7 @@ impl<'de> serde::Deserialize<'de> for RepartitionExecNode { Input, Partitioning, PreserveOrder, + EmitHashes, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -23729,6 +23756,7 @@ impl<'de> serde::Deserialize<'de> for RepartitionExecNode { "input" => Ok(GeneratedField::Input), "partitioning" => Ok(GeneratedField::Partitioning), "preserveOrder" | "preserve_order" => Ok(GeneratedField::PreserveOrder), + "emitHashes" | "emit_hashes" => Ok(GeneratedField::EmitHashes), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -23751,6 +23779,7 @@ impl<'de> serde::Deserialize<'de> for RepartitionExecNode { let mut input__ = None; let mut partitioning__ = None; let mut preserve_order__ = None; + let mut emit_hashes__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Input => { @@ -23771,12 +23800,19 @@ impl<'de> serde::Deserialize<'de> for RepartitionExecNode { } preserve_order__ = Some(map_.next_value()?); } + GeneratedField::EmitHashes => { + if emit_hashes__.is_some() { + return Err(serde::de::Error::duplicate_field("emitHashes")); + } + emit_hashes__ = Some(map_.next_value()?); + } } } Ok(RepartitionExecNode { input: input__, partitioning: partitioning__, preserve_order: preserve_order__.unwrap_or_default(), + emit_hashes: emit_hashes__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 2300f7192fb97..fe24a832aa74f 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2171,6 +2171,9 @@ pub struct AggregateExecNode { /// Optional dynamic filter expression for pushing down to the child. #[prost(message, optional, tag = "13")] pub dynamic_filter: ::core::option::Option, + /// Whether the partial aggregate should emit reusable group hashes. + #[prost(bool, tag = "14")] + pub emit_hashes: bool, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct GlobalLimitExecNode { @@ -2276,6 +2279,9 @@ pub struct RepartitionExecNode { pub partitioning: ::core::option::Option, #[prost(bool, tag = "6")] pub preserve_order: bool, + /// Whether hash repartitioning should propagate reusable group hashes. + #[prost(bool, tag = "7")] + pub emit_hashes: bool, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct Partitioning { diff --git a/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt b/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt index f1e787441d5e1..3dcf395eea162 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