From 634a59fff53e19c630657d728c4db3a57204727e Mon Sep 17 00:00:00 2001 From: Rachel Int Date: Wed, 7 Oct 2026 15:23:41 +0800 Subject: [PATCH] perf: add independent partial flush row and byte thresholds --- datafusion/common/src/config.rs | 16 ++ .../src/aggregates/hash_stream.rs | 119 ++++++-- .../test_files/aggregate_partial_flush.slt | 259 ++++++++++++++++++ .../test_files/information_schema.slt | 4 + docs/source/user-guide/configs.md | 2 + 5 files changed, 384 insertions(+), 16 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/aggregate_partial_flush.slt diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 173b69dcbdec5..905ceefa50d2d 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1136,6 +1136,22 @@ config_namespace! { /// aggregation ratio check and trying to switch to skipping aggregation mode pub skip_partial_aggregation_probe_rows_threshold: usize, default = 100_000 + /// (experimental) Allocated-byte threshold for flushing unordered partial + /// hash aggregation tables. Checked after each input batch; the table can + /// exceed this size by one batch. Flushed groups count toward the existing + /// skip-partial probe. Repeated keys do not disable flushing. Aggregations + /// with a soft group limit or nested aggregate state are excluded. + /// Set to 0 to disable this threshold. If both flush thresholds are + /// enabled, reaching either triggers a flush. Final aggregation is unchanged. + pub partial_aggregation_flush_bytes: usize, default = 0 + + /// (experimental) Number of distinct group rows above which an unordered + /// partial hash aggregation table is flushed. This counts groups held in + /// the table, not input rows. Uses the same eligibility and skip-partial + /// accounting as partial_aggregation_flush_bytes. Set to 0 to disable this + /// threshold. If both thresholds are enabled, reaching either triggers a flush. + pub partial_aggregation_flush_rows: usize, default = 0 + /// Should DataFusion use row number estimates at the input to decide /// whether increasing parallelism is beneficial or not. By default, /// only exact row numbers (not estimates) are used for this decision. diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 12babfe59cf9f..888cfb8afd566 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -159,6 +159,9 @@ pub(crate) struct PartialHashAggregateStream { /// Number of times accumulated states were emitted due to memory pressure. early_emit_count: metrics::Count, + /// Optional thresholds for emitting partial states before memory pressure. + table_flush: Option, + /// Tracks whether partial aggregation should switch to direct state conversion. skip_aggregation_probe: Option, @@ -207,9 +210,22 @@ enum HandleInputResult { ReachedLimit, #[expect(clippy::upper_case_acronyms)] OOM, + FlushThresholdReached, SwitchToSkipAggregation, } +/// Emits partial states at size thresholds while preserving skip-probe counts. +struct PartialTableFlush { + /// Allocated bytes above which the current table is emitted. + byte_threshold: usize, + /// Distinct group rows above which the current table is emitted; zero disables. + group_threshold: usize, + /// Previously emitted groups, included in the skip-partial reduction ratio. + flushed_groups: usize, + /// Number of emissions triggered by either size threshold. + flush_count: metrics::Count, +} + impl PartialHashAggregateStream { pub fn new( agg: &AggregateExec, @@ -232,6 +248,26 @@ impl PartialHashAggregateStream { let early_emit_count = MetricBuilder::new(&agg.metrics).counter("early_emit_count", partition); + let execution_options = &context.session_config().options().execution; + let byte_threshold = execution_options.partial_aggregation_flush_bytes; + let group_threshold = execution_options.partial_aggregation_flush_rows; + let group_values_soft_limit = agg.limit_options().map(|config| config.limit()); + let has_nested_state = schema + .fields() + .iter() + .skip(agg.group_by().num_group_exprs()) + .any(|field| field.data_type().is_nested()); + let table_flush = ((byte_threshold > 0 || group_threshold > 0) + && group_values_soft_limit.is_none() + && !has_nested_state) + .then(|| PartialTableFlush { + byte_threshold, + group_threshold, + flushed_groups: 0, + flush_count: MetricBuilder::new(&agg.metrics) + .counter("table_flush_count", partition), + }); + let hash_table = AggregateHashTable::::new( agg, partition, @@ -279,8 +315,9 @@ impl PartialHashAggregateStream { reservation, reduction_factor, early_emit_count, + table_flush, skip_aggregation_probe, - group_values_soft_limit: agg.limit_options().map(|config| config.limit()), + group_values_soft_limit, hash_table: Some(hash_table), }) } @@ -315,14 +352,23 @@ impl PartialHashAggregateStream { | HandleInputResult::SwitchToSkipAggregation => { break; } - HandleInputResult::OOM => { - let materialized_group_states = hash_table.take_state_batch()?.ok_or_else(|| { - internal_datafusion_err!( - "Partial hash aggregate ran out of memory with no aggregated groups" - ) - })?; - - self.early_emit_count.add(1); + HandleInputResult::OOM | HandleInputResult::FlushThresholdReached => { + let materialized_group_states = + hash_table.take_state_batch()?.ok_or_else(|| { + internal_datafusion_err!( + "Partial hash aggregate tried to flush an empty table" + ) + })?; + + if let Some(table_flush) = self.table_flush.as_mut() + && last_state == HandleInputResult::FlushThresholdReached + { + table_flush.flushed_groups += + materialized_group_states.num_rows(); + table_flush.flush_count.add(1); + } else { + self.early_emit_count.add(1); + } timer.done(); self.emit_on_memory_pressure( materialized_group_states, @@ -410,7 +456,14 @@ impl PartialHashAggregateStream { // ---------------------------------------------- // Step 3: Skip partial aggregation optimization // ---------------------------------------------- - self.update_skip_aggregation_probe(input_rows, hash_table.building_group_count()); + let flushed_groups = self + .table_flush + .as_ref() + .map_or(0, |table_flush| table_flush.flushed_groups); + self.update_skip_aggregation_probe( + input_rows, + flushed_groups + hash_table.building_group_count(), + ); // True branch: a decision has been made to skip partial aggregation. if self.should_skip_aggregation() { @@ -420,12 +473,26 @@ impl PartialHashAggregateStream { // ------------------------------------------------- // Step 4: Larger-than-memory execution (early emit) // ------------------------------------------------- - let resize_result = self.reservation.try_resize(hash_table.memory_size()); + let allocated_bytes = hash_table.memory_size(); + let resize_result = self.reservation.try_resize(allocated_bytes); match resize_result { - Ok(()) => Ok(HandleInputResult::ProcessNext), - Err(DataFusionError::ResourcesExhausted(_)) => Ok(HandleInputResult::OOM), - Err(e) => Err(e), + Ok(()) => {} + Err(DataFusionError::ResourcesExhausted(_)) => { + return Ok(HandleInputResult::OOM); + } + Err(e) => return Err(e), } + + // Step 5: Emit partial states once either configured threshold is reached. + if self.table_flush.as_ref().is_some_and(|table_flush| { + (table_flush.byte_threshold > 0 + && allocated_bytes >= table_flush.byte_threshold) + || (table_flush.group_threshold > 0 + && hash_table.building_group_count() >= table_flush.group_threshold) + }) { + return Ok(HandleInputResult::FlushThresholdReached); + } + Ok(HandleInputResult::ProcessNext) } /// Emit a materialized partial-state batch in `batch_size` slices. @@ -1021,9 +1088,16 @@ mod tests { Ok(()) } + #[rstest::rstest] + #[case(0, 0)] + #[case(1, 0)] + #[case(0, 1)] + #[case(1, 1)] #[tokio::test] - async fn test_partial_hash_stream_skip_aggregation_probe_not_locked_until_skip() - -> Result<()> { + async fn test_partial_hash_stream_skip_aggregation_probe_not_locked_until_skip( + #[case] flush_bytes: usize, + #[case] flush_rows: usize, + ) -> Result<()> { // Test that the probe is not locked until we actually decide to skip. // This allows us to continue evaluating the skip condition across multiple batches. // @@ -1108,6 +1182,14 @@ mod tests { // Configure skip aggregation settings let mut session_config = task_ctx.session_config().clone(); + session_config + .options_mut() + .execution + .partial_aggregation_flush_bytes = flush_bytes; + session_config + .options_mut() + .execution + .partial_aggregation_flush_rows = flush_rows; session_config = session_config.set( "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", &datafusion_common::ScalarValue::UInt64(Some(probe_rows_threshold)), @@ -1166,6 +1248,11 @@ mod tests { "Expected batch 3's rows ({batch3_rows}) to be skipped", ); + let flushes = metrics + .sum_by_name("table_flush_count") + .map_or(0, |value| value.as_usize()); + assert_eq!(flushes, usize::from(flush_bytes > 0 || flush_rows > 0)); + Ok(()) } diff --git a/datafusion/sqllogictest/test_files/aggregate_partial_flush.slt b/datafusion/sqllogictest/test_files/aggregate_partial_flush.slt new file mode 100644 index 0000000000000..f12e5109d9a0a --- /dev/null +++ b/datafusion/sqllogictest/test_files/aggregate_partial_flush.slt @@ -0,0 +1,259 @@ +# 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. + +# Compare ordinary aggregation with Partial flushing by bytes or group rows. +# Repeated integer, long string, and mixed keys must merge to the same results. +# Aggregates with nested state keep their existing execution path. + +# The test context runs with 4 target partitions, so the aggregation is two-stage. + +statement ok +set datafusion.execution.batch_size = 1024; + +statement ok +create table t as select value as v from generate_series(1, 200000); + +# --- partial_aggregation_flush_bytes = 0 + +statement ok +set datafusion.execution.partial_aggregation_flush_bytes = 0; + +# integer key, several accumulators +query IIIII +select count(*), sum(c), sum(s), min(mn), max(mx) from (select v % 10007 as k, count(*) c, sum(v) s, min(v) mn, max(v) mx from t group by k); +---- +10007 200000 20000100000 1 200000 + +# avg keeps sum and count as state +query IR +select count(*), round(sum(a), 3) from (select v % 10007 as k, avg(v) a from t group by k); +---- +10007 1000705003.5 + +# string key +query II +select count(*), sum(c) from (select concat('key-', cast(v % 9973 as varchar)) as k, count(*) c from t group by k); +---- +9973 200000 + +# string view key +query II +select count(*), sum(c) from (select arrow_cast(concat('a-long-enough-string-view-key-', cast(v % 9973 as varchar)), 'Utf8View') as k, count(*) c from t group by k); +---- +9973 200000 + +# several key columns of mixed types +query III +select count(*), sum(c), sum(s) from (select v % 101 as k1, concat('k', cast(v % 103 as varchar)) as k2, v % 7 = 0 as k3, count(*) c, sum(v) s from t group by k1, k2, k3); +---- +20806 200000 20000100000 + +# nullable key +query II +select count(*), sum(c) from (select case when v % 11 = 0 then null else v % 5003 end as k, count(*) c from t group by k); +---- +5004 200000 + +# distinct has no accumulators +query I +select count(*) from (select distinct v % 20011 as k, v % 3 as j from t); +---- +60033 + +# accumulators with list state +query II +select count(*), sum(d) from (select v % 5003 as k, count(distinct v % 17) d from t group by k); +---- +5003 85051 + +# median +query IR +select count(*), sum(m) from (select v % 5003 as k, median(v) m from t group by k); +---- +5003 500302501.5 + +# array_agg +query II +select count(*), sum(array_length(a)) from (select v % 5003 as k, array_agg(v) a from t group by k); +---- +5003 200000 + +# every row its own group, largest groups first +query III +select k, c, s from (select v as k, count(*) c, sum(v) s from t group by v) order by k desc limit 3; +---- +200000 1 200000 +199999 1 199999 +199998 1 199998 + +# --- partial_aggregation_flush_bytes = 32768 + +statement ok +set datafusion.execution.partial_aggregation_flush_bytes = 32768; + +# integer key, several accumulators +query IIIII +select count(*), sum(c), sum(s), min(mn), max(mx) from (select v % 10007 as k, count(*) c, sum(v) s, min(v) mn, max(v) mx from t group by k); +---- +10007 200000 20000100000 1 200000 + +# avg keeps sum and count as state +query IR +select count(*), round(sum(a), 3) from (select v % 10007 as k, avg(v) a from t group by k); +---- +10007 1000705003.5 + +# string key +query II +select count(*), sum(c) from (select concat('key-', cast(v % 9973 as varchar)) as k, count(*) c from t group by k); +---- +9973 200000 + +# string view key +query II +select count(*), sum(c) from (select arrow_cast(concat('a-long-enough-string-view-key-', cast(v % 9973 as varchar)), 'Utf8View') as k, count(*) c from t group by k); +---- +9973 200000 + +# several key columns of mixed types +query III +select count(*), sum(c), sum(s) from (select v % 101 as k1, concat('k', cast(v % 103 as varchar)) as k2, v % 7 = 0 as k3, count(*) c, sum(v) s from t group by k1, k2, k3); +---- +20806 200000 20000100000 + +# nullable key +query II +select count(*), sum(c) from (select case when v % 11 = 0 then null else v % 5003 end as k, count(*) c from t group by k); +---- +5004 200000 + +# distinct has no accumulators +query I +select count(*) from (select distinct v % 20011 as k, v % 3 as j from t); +---- +60033 + +# accumulators with list state +query II +select count(*), sum(d) from (select v % 5003 as k, count(distinct v % 17) d from t group by k); +---- +5003 85051 + +# median +query IR +select count(*), sum(m) from (select v % 5003 as k, median(v) m from t group by k); +---- +5003 500302501.5 + +# array_agg +query II +select count(*), sum(array_length(a)) from (select v % 5003 as k, array_agg(v) a from t group by k); +---- +5003 200000 + +# every row its own group, largest groups first +query III +select k, c, s from (select v as k, count(*) c, sum(v) s from t group by v) order by k desc limit 3; +---- +200000 1 200000 +199999 1 199999 +199998 1 199998 + +statement ok +set datafusion.execution.partial_aggregation_flush_bytes = 0; + +# --- partial_aggregation_flush_rows = 100 + +statement ok +set datafusion.execution.partial_aggregation_flush_rows = 100; + +# integer key, several accumulators +query IIIII +select count(*), sum(c), sum(s), min(mn), max(mx) from (select v % 10007 as k, count(*) c, sum(v) s, min(v) mn, max(v) mx from t group by k); +---- +10007 200000 20000100000 1 200000 + +# avg keeps sum and count as state +query IR +select count(*), round(sum(a), 3) from (select v % 10007 as k, avg(v) a from t group by k); +---- +10007 1000705003.5 + +# string key +query II +select count(*), sum(c) from (select concat('key-', cast(v % 9973 as varchar)) as k, count(*) c from t group by k); +---- +9973 200000 + +# string view key +query II +select count(*), sum(c) from (select arrow_cast(concat('a-long-enough-string-view-key-', cast(v % 9973 as varchar)), 'Utf8View') as k, count(*) c from t group by k); +---- +9973 200000 + +# several key columns of mixed types +query III +select count(*), sum(c), sum(s) from (select v % 101 as k1, concat('k', cast(v % 103 as varchar)) as k2, v % 7 = 0 as k3, count(*) c, sum(v) s from t group by k1, k2, k3); +---- +20806 200000 20000100000 + +# nullable key +query II +select count(*), sum(c) from (select case when v % 11 = 0 then null else v % 5003 end as k, count(*) c from t group by k); +---- +5004 200000 + +# distinct has no accumulators +query I +select count(*) from (select distinct v % 20011 as k, v % 3 as j from t); +---- +60033 + +# accumulators with list state +query II +select count(*), sum(d) from (select v % 5003 as k, count(distinct v % 17) d from t group by k); +---- +5003 85051 + +# median +query IR +select count(*), sum(m) from (select v % 5003 as k, median(v) m from t group by k); +---- +5003 500302501.5 + +# array_agg +query II +select count(*), sum(array_length(a)) from (select v % 5003 as k, array_agg(v) a from t group by k); +---- +5003 200000 + +# every row its own group, largest groups first +query III +select k, c, s from (select v as k, count(*) c, sum(v) s from t group by v) order by k desc limit 3; +---- +200000 1 200000 +199999 1 199999 +199998 1 199998 + +statement ok +reset datafusion.execution.partial_aggregation_flush_rows; + +statement ok +reset datafusion.execution.partial_aggregation_flush_bytes; + +statement ok +reset datafusion.execution.batch_size; diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index d1898a691d139..cd8331d09b8b6 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -272,6 +272,8 @@ datafusion.execution.parquet.statistics_enabled page datafusion.execution.parquet.statistics_truncate_length 64 datafusion.execution.parquet.write_batch_size 1024 datafusion.execution.parquet.writer_version 1.0 +datafusion.execution.partial_aggregation_flush_bytes 0 +datafusion.execution.partial_aggregation_flush_rows 0 datafusion.execution.perfect_hash_join_min_key_density 0.15 datafusion.execution.perfect_hash_join_small_build_threshold 1024 datafusion.execution.planning_concurrency 13 @@ -436,6 +438,8 @@ datafusion.execution.parquet.statistics_enabled page (writing) Sets if statistic datafusion.execution.parquet.statistics_truncate_length 64 (writing) Sets statistics truncate length. If NULL, uses default parquet writer setting datafusion.execution.parquet.write_batch_size 1024 (writing) Sets write_batch_size in rows datafusion.execution.parquet.writer_version 1.0 (writing) Sets parquet writer version valid values are "1.0" and "2.0" +datafusion.execution.partial_aggregation_flush_bytes 0 (experimental) Allocated-byte threshold for flushing unordered partial hash aggregation tables. Checked after each input batch; the table can exceed this size by one batch. Flushed groups count toward the existing skip-partial probe. Repeated keys do not disable flushing. Aggregations with a soft group limit or nested aggregate state are excluded. Set to 0 to disable this threshold. If both flush thresholds are enabled, reaching either triggers a flush. Final aggregation is unchanged. +datafusion.execution.partial_aggregation_flush_rows 0 (experimental) Number of distinct group rows above which an unordered partial hash aggregation table is flushed. This counts groups held in the table, not input rows. Uses the same eligibility and skip-partial accounting as partial_aggregation_flush_bytes. Set to 0 to disable this threshold. If both thresholds are enabled, reaching either triggers a flush. datafusion.execution.perfect_hash_join_min_key_density 0.15 The minimum required density of join keys on the build side to consider a perfect hash join (see `HashJoinExec` for more details). Density is calculated as: `(number of rows) / (max_key - min_key + 1)`. A perfect hash join may be used if the actual key density > this value. Currently only supports cases where build_side.num_rows() < u32::MAX. Support for build_side.num_rows() >= u32::MAX will be added in the future. datafusion.execution.perfect_hash_join_small_build_threshold 1024 A perfect hash join (see `HashJoinExec` for more details) will be considered if the range of keys (max - min) on the build side is < this threshold. This provides a fast path for joins with very small key ranges, bypassing the density check. Currently only supports cases where build_side.num_rows() < u32::MAX. Support for build_side.num_rows() >= u32::MAX will be added in the future. datafusion.execution.planning_concurrency 13 Fan-out during initial physical planning. This is mostly use to plan `UNION` children in parallel. Defaults to the number of CPU cores on the system diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 16b96a2244344..c71f506f2e7e1 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -143,6 +143,8 @@ The following configuration settings are available: | datafusion.execution.enable_file_stream_work_stealing | true | When `true` (the default), DataFusion's built-in file scans dynamically rebalance files across partitions at query execution time: a partition that goes idle reads files (or byte-range morsels) originally assigned to a sibling partition, which keeps all partitions busy in a single process. Executors that depend on the plan-time partition assignment — such as Ballista and datafusion-distributed, which run each partition as an isolated task and never poll the siblings — should set this to `false` so each partition reads only its own file group and no runtime reassignment occurs. | | datafusion.execution.skip_partial_aggregation_probe_ratio_threshold | 0.8 | Aggregation ratio (number of distinct groups / number of input rows) threshold for skipping partial aggregation. If the value is greater then partial aggregation will skip aggregation for further input | | datafusion.execution.skip_partial_aggregation_probe_rows_threshold | 100000 | Number of input rows partial aggregation partition should process, before aggregation ratio check and trying to switch to skipping aggregation mode | +| datafusion.execution.partial_aggregation_flush_bytes | 0 | (experimental) Allocated-byte threshold for flushing unordered partial hash aggregation tables. Checked after each input batch; the table can exceed this size by one batch. Flushed groups count toward the existing skip-partial probe. Repeated keys do not disable flushing. Aggregations with a soft group limit or nested aggregate state are excluded. Set to 0 to disable this threshold. If both flush thresholds are enabled, reaching either triggers a flush. Final aggregation is unchanged. | +| datafusion.execution.partial_aggregation_flush_rows | 0 | (experimental) Number of distinct group rows above which an unordered partial hash aggregation table is flushed. This counts groups held in the table, not input rows. Uses the same eligibility and skip-partial accounting as partial_aggregation_flush_bytes. Set to 0 to disable this threshold. If both thresholds are enabled, reaching either triggers a flush. | | datafusion.execution.use_row_number_estimates_to_optimize_partitioning | false | Should DataFusion use row number estimates at the input to decide whether increasing parallelism is beneficial or not. By default, only exact row numbers (not estimates) are used for this decision. Setting this flag to `true` will likely produce better plans. if the source of statistics is accurate. We plan to make this the default in the future. | | datafusion.execution.enforce_batch_size_in_joins | false | Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. Note: this option currently only applies to the symmetric hash join. | | datafusion.execution.objectstore_writer_buffer_size | 10485760 | Size (bytes) of data buffer DataFusion uses when writing output files. This affects the size of the data chunks that are uploaded to remote object stores (e.g. AWS S3). If very large (>= 100 GiB) output files are being written, it may be necessary to increase this size to avoid errors from the remote end point. |