From 634a59fff53e19c630657d728c4db3a57204727e Mon Sep 17 00:00:00 2001 From: Rachel Int Date: Wed, 7 Oct 2026 15:23:41 +0800 Subject: [PATCH 1/3] 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. | From 522c7492f0ee2596a8d244c46bda805196f786e2 Mon Sep 17 00:00:00 2001 From: Rachel Int Date: Wed, 7 Oct 2026 17:16:56 +0800 Subject: [PATCH 2/3] perf: reuse StringView storage across partial flushes --- datafusion/common/src/config.rs | 3 + .../src/binary_view_map.rs | 70 ++++++++++++++++++- .../aggregates/aggregate_hash_table/common.rs | 24 ++++++- .../src/aggregates/group_values/mod.rs | 8 +++ .../single_group_by/bytes_view.rs | 6 ++ .../src/aggregates/hash_stream.rs | 7 +- docs/source/user-guide/configs.md | 2 +- 7 files changed, 115 insertions(+), 5 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 905ceefa50d2d..375342d6ba2d7 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1141,6 +1141,9 @@ config_namespace! { /// 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. + /// After a threshold flush, single StringView keys retain hash-table + /// capacity and reserve a contiguous payload buffer using the previous + /// payload length. Memory-pressure emission still releases allocations. /// 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 diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 725a1e3c5b271..616810fabd377 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -145,6 +145,8 @@ where views: Vec, /// In-progress buffer for out-of-line string data in_progress: Vec, + /// Maximum payload block size; extended after a reusable Partial flush. + payload_block_limit: usize, /// Completed buffers containing string data completed: Vec, @@ -190,6 +192,7 @@ where initial_map_capacity: map_capacity, views: Vec::new(), in_progress: Vec::new(), + payload_block_limit: BYTE_VIEW_MAX_BLOCK_SIZE, completed: Vec::new(), random_state: RandomState::default(), hashes_buffer: vec![], @@ -206,6 +209,29 @@ where new_self } + /// Emits all keys while retaining the hash allocation for a Partial flush. + /// The returned array owns the emitted strings. The empty map uses one payload + /// buffer, reserved from the previous payload length, up to the u32 offset + /// limit. Normal output and memory-pressure emission should use [`Self::take`]. + pub fn take_state_reusing_allocation(&mut self) -> ArrayRef { + let payload_bytes = self + .completed + .iter() + .fold(self.in_progress.len(), |total, buffer| { + total.saturating_add(buffer.len()) + }); + let mut outgoing = Self::with_capacity(self.output_type, 0); + std::mem::swap(self, &mut outgoing); + std::mem::swap(&mut self.map, &mut outgoing.map); + self.map.clear(); + self.initial_map_capacity = outgoing.initial_map_capacity; + std::mem::swap(&mut self.random_state, &mut outgoing.random_state); + self.payload_block_limit = u32::MAX as usize; + self.in_progress + .reserve_exact(payload_bytes.min(self.payload_block_limit)); + outgoing.into_state() + } + /// Empties this map and releases every allocation it holds, so /// [`Self::size`] drops to approximately zero. /// @@ -538,7 +564,7 @@ where make_view(value, 0, 0) } else { // Ensure buffer is big enough - if self.in_progress.len() + len > BYTE_VIEW_MAX_BLOCK_SIZE { + if self.in_progress.len() + len > self.payload_block_limit { let flushed = std::mem::replace( &mut self.in_progress, Vec::with_capacity(BYTE_VIEW_MAX_BLOCK_SIZE), @@ -890,6 +916,48 @@ mod tests { assert_eq!(lazy.map.capacity(), 0); } + #[test] + fn partial_flush_reuses_table_and_preserves_emitted_strings() { + let mut map = ArrowBytesViewMap::::new(OutputType::Utf8View); + let long_value = "x".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1); + let other_value = "y".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1); + let expected = vec![ + Some("inline"), + None, + Some(long_value.as_str()), + Some(other_value.as_str()), + ]; + let values: ArrayRef = Arc::new(StringViewArray::from(expected.clone())); + let mut emitted = Vec::new(); + for _ in 0..3 { + let mut next_group = 0; + let mut groups = Vec::new(); + map.insert_if_new( + &values, + |_| { + let group = next_group; + next_group += 1; + group + }, + |group| groups.push(group), + ); + assert_eq!(groups, vec![0, 1, 2, 3]); + let capacity = map.map.capacity(); + emitted.push(map.take_state_reusing_allocation()); + assert!(map.is_empty()); + assert_eq!(map.map.capacity(), capacity); + assert!(map.in_progress.capacity() >= long_value.len() + other_value.len()); + } + for (index, array) in emitted.into_iter().enumerate() { + let array = array.as_string_view(); + assert_eq!(array.iter().collect::>(), expected); + assert_eq!(array.data_buffers().len(), if index == 0 { 2 } else { 1 }); + } + map.clear_and_release(); + assert_eq!(map.map.allocation_size(), 0); + assert_eq!(map.in_progress.capacity(), 0); + } + #[test] fn clear_and_release_frees_the_preallocation_that_take_keeps() { let mut map = ArrowBytesViewMap::<()>::with_capacity( 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 020196e7da022..9265e0982cfae 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -372,6 +372,20 @@ impl AggregateHashTable { /// spilling without finalizing the same group more than once. pub(in crate::aggregates) fn take_state_batch( &mut self, + ) -> Result> { + self.take_state_batch_inner(false) + } + + /// Emits a threshold-triggered Partial batch while allowing key storage reuse. + pub(in crate::aggregates) fn take_state_batch_for_partial_flush( + &mut self, + ) -> Result> { + self.take_state_batch_inner(true) + } + + fn take_state_batch_inner( + &mut self, + is_partial_flush: bool, ) -> Result> { let state_schema = Arc::clone(&self.state_schema); let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); @@ -382,7 +396,11 @@ impl AggregateHashTable { } let output = group_by_metrics.time_emitting(|| { - let mut output = state.group_values.emit(EmitTo::All)?; + let mut output = if is_partial_flush { + state.group_values.emit_for_partial_flush()? + } else { + state.group_values.emit(EmitTo::All)? + }; for (idx, acc) in state.accumulators.iter_mut().enumerate() { output.extend(accumulator_metrics.time( idx, @@ -399,7 +417,9 @@ impl AggregateHashTable { // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the // key/index buffers too so the memory reservation can be released // before the batch is sorted for spilling. - state.group_values.clear_shrink(0); + if !is_partial_flush { + state.group_values.clear_shrink(0); + } state.batch_group_indices.clear(); state.batch_group_indices.shrink_to_fit(); diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index 19d3b94f8a72a..c6d0b8b45e5b1 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -111,6 +111,14 @@ pub trait GroupValues: Send { /// Emits the group values fn emit(&mut self, emit_to: EmitTo) -> Result>; + /// Emits all keys for a Partial threshold flush, preparing for more input. + /// Implementations may retain allocations when the table is filled again. + fn emit_for_partial_flush(&mut self) -> Result> { + let values = self.emit(EmitTo::All)?; + self.clear_shrink(0); + Ok(values) + } + /// Materializes selected group values without changing the stored values or /// their group indices. /// 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 a2e3573ce889f..a8b8f4c883456 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 @@ -126,6 +126,12 @@ impl GroupValues for GroupValuesBytesView { Ok(vec![group_values]) } + fn emit_for_partial_flush(&mut self) -> datafusion_common::Result> { + let values = self.map.take_state_reusing_allocation(); + self.num_groups = 0; + Ok(vec![values]) + } + fn values_preserving( &mut self, selection: GroupSelection<'_>, diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 888cfb8afd566..aba2e7c18c300 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -354,7 +354,12 @@ impl PartialHashAggregateStream { } HandleInputResult::OOM | HandleInputResult::FlushThresholdReached => { let materialized_group_states = - hash_table.take_state_batch()?.ok_or_else(|| { + if last_state == HandleInputResult::FlushThresholdReached { + hash_table.take_state_batch_for_partial_flush()? + } else { + hash_table.take_state_batch()? + } + .ok_or_else(|| { internal_datafusion_err!( "Partial hash aggregate tried to flush an empty table" ) diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index c71f506f2e7e1..5ecd738378102 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -143,7 +143,7 @@ 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_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. After a threshold flush, single StringView keys retain hash-table capacity and reserve a contiguous payload buffer using the previous payload length. Memory-pressure emission still releases allocations. 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. | From 5008df083795701eac24112064ac6711bb75791a Mon Sep 17 00:00:00 2001 From: Rachel Int Date: Wed, 7 Oct 2026 18:22:41 +0800 Subject: [PATCH 3/3] perf: reserve group-key storage across partial flushes --- datafusion/common/src/config.rs | 9 +- .../physical-expr-common/src/binary_map.rs | 19 ++++ .../src/binary_view_map.rs | 4 + .../aggregates/aggregate_hash_table/common.rs | 4 +- .../src/aggregates/group_values/mod.rs | 98 ++++++++++++++++++- .../group_values/multi_group_by/bytes.rs | 9 ++ .../group_values/multi_group_by/bytes_view.rs | 17 ++++ .../group_values/multi_group_by/dictionary.rs | 12 +++ .../multi_group_by/fixed_size_binary.rs | 7 ++ .../group_values/multi_group_by/list.rs | 20 ++++ .../group_values/multi_group_by/mod.rs | 23 +++++ .../group_values/multi_group_by/primitive.rs | 7 ++ .../group_values/multi_group_by/row_backed.rs | 6 ++ .../src/aggregates/group_values/row.rs | 4 + .../group_values/single_group_by/bytes.rs | 6 ++ .../group_values/single_group_by/primitive.rs | 7 ++ docs/source/user-guide/configs.md | 2 +- 17 files changed, 248 insertions(+), 6 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 375342d6ba2d7..82441b03b4065 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1141,9 +1141,12 @@ config_namespace! { /// 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. - /// After a threshold flush, single StringView keys retain hash-table - /// capacity and reserve a contiguous payload buffer using the previous - /// payload length. Memory-pressure emission still releases allocations. + /// Threshold flushes retain group-key hash tables and scratch buffers. + /// Key buffers transferred to output are reserved again using the previous + /// lengths, including primitive, string, and multi-column keys. StringView + /// payloads use contiguous buffers up to the view offset limit. Retained + /// capacity counts toward this threshold and memory reservations. + /// Memory-pressure emission still releases allocations. /// 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 diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index c9f8c5d119ce9..b66b087f08656 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -336,6 +336,25 @@ where new_self } + /// Emits keys while retaining the hash table and reserving replacement key buffers. + /// Output arrays own the previous key buffers, so subsequent input cannot change them. + pub fn take_state_reusing_allocation(&mut self) -> ArrayRef { + let payload_bytes = self.buffer.len(); + let num_offsets = self.offsets.len(); + let mut outgoing = Self::new_inner(self.output_type, 0, 0); + swap(self, &mut outgoing); + swap(&mut self.map, &mut outgoing.map); + self.map.clear(); + swap(&mut self.hashes_buffer, &mut outgoing.hashes_buffer); + self.hashes_buffer.clear(); + swap(&mut self.random_state, &mut outgoing.random_state); + self.initial_map_capacity = outgoing.initial_map_capacity; + self.initial_buffer_capacity = outgoing.initial_buffer_capacity; + self.buffer.reserve_exact(payload_bytes); + self.offsets.reserve_exact(num_offsets.saturating_sub(1)); + outgoing.into_state() + } + /// Empties this map and releases every allocation it holds, so /// [`Self::size`] drops to approximately zero. /// diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 616810fabd377..0a986982e68db 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -220,10 +220,14 @@ where .fold(self.in_progress.len(), |total, buffer| { total.saturating_add(buffer.len()) }); + let num_views = self.views.len(); let mut outgoing = Self::with_capacity(self.output_type, 0); std::mem::swap(self, &mut outgoing); std::mem::swap(&mut self.map, &mut outgoing.map); self.map.clear(); + std::mem::swap(&mut self.hashes_buffer, &mut outgoing.hashes_buffer); + self.hashes_buffer.clear(); + self.views.reserve_exact(num_views); self.initial_map_capacity = outgoing.initial_map_capacity; std::mem::swap(&mut self.random_state, &mut outgoing.random_state); self.payload_block_limit = u32::MAX as usize; 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 9265e0982cfae..9543a839b7014 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -421,7 +421,9 @@ impl AggregateHashTable { state.group_values.clear_shrink(0); } state.batch_group_indices.clear(); - state.batch_group_indices.shrink_to_fit(); + if !is_partial_flush { + state.batch_group_indices.shrink_to_fit(); + } Ok(Some(batch)) } diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index c6d0b8b45e5b1..dc7060414a2b8 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -223,14 +223,110 @@ mod tests { use std::sync::Arc; use arrow::array::{ - ArrayRef, AsArray, BooleanArray, Int32Array, StringArray, StringViewArray, + Array, ArrayRef, AsArray, BooleanArray, FixedSizeBinaryArray, Int32Array, + ListArray, StringArray, StringViewArray, StructArray, UInt32Array, }; + use arrow::compute; use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use datafusion_expr::{EmitTo, GroupSelection}; use super::new_group_values; use crate::aggregates::order::GroupOrdering; + #[test] + fn partial_flush_preserves_output_and_resets_all_key_types() { + let strings: ArrayRef = Arc::new(StringArray::from(vec![ + Some("first long string value"), + None, + Some("second long string value"), + ])); + let mut arrays: Vec = vec![ + Arc::new(Int32Array::from(vec![Some(11), None, Some(22)])), + Arc::new(BooleanArray::from(vec![Some(true), None, Some(false)])), + Arc::clone(&strings), + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), None]), + None, + Some(vec![Some(2), Some(3)]), + ])), + Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + vec![Some(b"abc".as_slice()), None, Some(b"def".as_slice())] + .into_iter(), + 3, + ) + .unwrap(), + ), + ]; + arrays.push(Arc::new(StructArray::from(vec![( + Arc::new(Field::new("nested", DataType::Int32, true)), + Arc::clone(&arrays[0]), + )]))); + for data_type in [ + DataType::LargeUtf8, + DataType::Binary, + DataType::LargeBinary, + DataType::Utf8View, + DataType::BinaryView, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + ] { + arrays.push(compute::cast(&strings, &data_type).unwrap()); + } + let mut cases: Vec> = + arrays.iter().map(|a| vec![Arc::clone(a)]).collect(); + cases.push(arrays); + for columns in cases { + let schema = Arc::new(Schema::new( + columns + .iter() + .enumerate() + .map(|(i, a)| { + Field::new(format!("key{i}"), a.data_type().clone(), true) + }) + .collect::>(), + )); + let mut values = new_group_values(schema, &GroupOrdering::None).unwrap(); + let mut retained_outputs = vec![]; + let mut groups = vec![]; + for indices in [[0, 1, 2, 0], [2, 0, 1, 2], [1, 2, 0, 1]] { + let input = columns + .iter() + .map(|a| { + compute::take( + a.as_ref(), + &UInt32Array::from(indices.to_vec()), + None, + ) + .unwrap() + }) + .collect::>(); + values.intern(&input, &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2, 0]); + let output = values.emit_for_partial_flush().unwrap(); + assert!(values.is_empty()); + let expected = input.iter().map(|a| a.slice(0, 3)).collect::>(); + retained_outputs.push((output, expected)); + } + // Earlier batches stay alive while the same builders are refilled. + for (output, expected) in retained_outputs { + for (actual, expected) in output.iter().zip(expected.iter()) { + // Dictionary output may compact or reorder its dictionary values. + let data_type = match expected.data_type() { + DataType::Dictionary(_, value_type) => value_type.as_ref(), + data_type => data_type, + }; + let actual = compute::cast(actual, data_type).unwrap(); + let expected = compute::cast(expected, data_type).unwrap(); + assert_eq!(actual.as_ref(), expected.as_ref()); + } + } + values.clear_shrink(0); + assert!(values.is_empty()); + values.intern(&columns, &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2]); + } + } + #[test] fn preserving_values_keep_group_indices_valid() { let schema = Arc::new(Schema::new(vec![Field::new( diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs index 7fb4f26c6a71c..0e8cb6df3ed4f 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs @@ -419,6 +419,15 @@ where }) } + fn take_for_partial_flush(&mut self) -> ArrayRef { + let num_values = self.len(); + let payload_bytes = self.buffer.len(); + let output = self.take_n(num_values); + self.offsets.reserve_exact(num_values); + self.buffer.reserve(payload_bytes); + output + } + fn take_n(&mut self, n: usize) -> ArrayRef { debug_assert!(self.len() >= n); let null_buffer = self.nulls.take_n(n); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs index 24eb94476325a..8a5b849a2158b 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs @@ -652,6 +652,23 @@ impl GroupColumn for ByteViewGroupValueBuilder { self.values_preserving_inner(selection) } + fn take_for_partial_flush(&mut self) -> ArrayRef { + let num_views = self.views.len(); + let payload_bytes = self + .completed + .iter() + .fold(self.in_progress.len(), |total, buffer| { + total.saturating_add(buffer.len()) + }); + let mut fresh = Self::new().with_max_block_size(u32::MAX as usize); + fresh.views.reserve_exact(num_views); + fresh + .in_progress + .reserve_exact(payload_bytes.min(u32::MAX as usize)); + fresh.nulls = NullBufferBuilder::new(num_views); + replace(self, fresh).build_inner() + } + fn take_n(&mut self, n: usize) -> ArrayRef { self.take_n_inner(n) } diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs index c58ce2a948e3e..5b619a0bb7c6c 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -549,6 +549,18 @@ impl GroupColumn Ok(Self::into_dict(values, &selected_groups, null_inner_slot)) } + fn take_for_partial_flush(&mut self) -> ArrayRef { + let values = self.inner.take_for_partial_flush(); + let output = Self::into_dict(values, &self.group_to_inner, self.null_inner_slot); + self.group_to_inner.clear(); + self.value_dedup.clear(); + self.null_inner_slot = None; + self.val_to_inner.clear(); + self.val_hashes.clear(); + self.cached_values = None; + output + } + fn take_n(&mut self, n: usize) -> ArrayRef { let old_inner_len = self.inner.len(); let all_inner_values = self.inner.take_n(old_inner_len); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs index 521bc8e2f005a..c3221fefeb612 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs @@ -244,6 +244,13 @@ impl GroupColumn for FixedSizeBinaryGroupValueBuilder { )) } + fn take_for_partial_flush(&mut self) -> ArrayRef { + let payload_bytes = self.buffer.len(); + let output = self.take_n(self.len); + self.buffer.reserve_exact(payload_bytes); + output + } + fn take_n(&mut self, n: usize) -> ArrayRef { debug_assert!(self.len >= n); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/list.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/list.rs index b61e1987d7ee2..fa38360eca483 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/list.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/list.rs @@ -261,6 +261,26 @@ impl GroupColumn for ListGroupValueBuilder { ))) } + fn take_for_partial_flush(&mut self) -> ArrayRef { + let num_offsets = self.offsets.len(); + let offsets = + std::mem::replace(&mut self.offsets, Vec::with_capacity(num_offsets)); + self.offsets.push(O::usize_as(0)); + let nulls = std::mem::replace( + &mut self.outer_nulls, + NullBufferBuilder::new(self.outer_len), + ) + .build(); + self.outer_len = 0; + let child = self.child.take_for_partial_flush(); + Arc::new(GenericListArray::::new( + Arc::clone(&self.field), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + child, + nulls, + )) + } + fn take_n(&mut self, n: usize) -> ArrayRef { // Number of child elements consumed by the first n outer rows. let cut_offset = self.offsets[n]; 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 68d56f00a9de5..e6d0709838195 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 @@ -116,6 +116,12 @@ pub trait GroupColumn: Send + Sync { /// Builds a new array from the first `n` stored rows, shifting the /// remaining rows to the start of the builder fn take_n(&mut self, n: usize) -> ArrayRef; + + /// Emits all rows and prepares reusable storage for the next Partial flush cycle. + /// Buffers transferred to the result must not be modified by subsequent input. + fn take_for_partial_flush(&mut self) -> ArrayRef { + self.take_n(self.len()) + } } /// Determines if the nullability of the existing and new input array can be used @@ -1211,6 +1217,23 @@ impl GroupValues for GroupValuesColumn { Ok(output) } + fn emit_for_partial_flush(&mut self) -> Result> { + let output = self + .group_values + .iter_mut() + .map(|column| column.take_for_partial_flush()) + .collect(); + self.map.clear(); + self.hashes_buffer.clear(); + self.group_index_lists.clear(); + self.emit_group_index_list_buffer.clear(); + self.vectorized_operation_buffers.clear(); + self.vectorized_operation_buffers + .equal_to_results + .truncate(0); + Ok(output) + } + fn values_preserving( &mut self, selection: GroupSelection<'_>, diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs index ebd1e071cf8d5..39162b71e04a9 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs @@ -299,6 +299,13 @@ where )) } + fn take_for_partial_flush(&mut self) -> ArrayRef { + let num_values = self.group_values.len(); + let output = self.take_n(num_values); + self.group_values.reserve_exact(num_values); + output + } + fn take_n(&mut self, n: usize) -> ArrayRef { let first_n = split_vec_min_alloc(&mut self.group_values, n); let first_n_nulls = if NULLABLE { self.nulls.take_n(n) } else { None }; diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs index c45942edac6ff..0b8f3d420d383 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs @@ -304,6 +304,12 @@ impl GroupColumn for RowsGroupColumn { Ok(self.rows_to_array(rows)) } + fn take_for_partial_flush(&mut self) -> ArrayRef { + let output = self.rows_to_array(&self.group_values); + self.group_values.clear(); + output + } + fn take_n(&mut self, n: usize) -> ArrayRef { debug_assert!(n <= self.group_values.num_rows()); diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index a2ff0b6adb2b2..683fa2c2bde2d 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -266,6 +266,10 @@ impl GroupValues for GroupValuesRows { Ok(output) } + fn emit_for_partial_flush(&mut self) -> Result> { + self.emit(EmitTo::All) + } + fn values_preserving( &mut self, selection: GroupSelection<'_>, 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 3c92ce787dc2b..c6c9b91e5d9ef 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 @@ -124,6 +124,12 @@ impl GroupValues for GroupValuesBytes { Ok(vec![group_values]) } + fn emit_for_partial_flush(&mut self) -> Result> { + let values = self.map.take_state_reusing_allocation(); + self.num_groups = 0; + Ok(vec![values]) + } + fn values_preserving( &mut self, selection: GroupSelection<'_>, 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 eaeb83e0cd712..f50c31a965a2f 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 @@ -241,6 +241,13 @@ where Ok(vec![Arc::new(array.with_data_type(self.data_type.clone()))]) } + fn emit_for_partial_flush(&mut self) -> Result> { + let num_values = self.values.len(); + let output = self.emit(EmitTo::All)?; + self.values.reserve_exact(num_values); + Ok(output) + } + fn values_preserving( &mut self, selection: GroupSelection<'_>, diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 5ecd738378102..320001d8cb424 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -143,7 +143,7 @@ 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. After a threshold flush, single StringView keys retain hash-table capacity and reserve a contiguous payload buffer using the previous payload length. Memory-pressure emission still releases allocations. 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_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. Threshold flushes retain group-key hash tables and scratch buffers. Key buffers transferred to output are reserved again using the previous lengths, including primitive, string, and multi-column keys. StringView payloads use contiguous buffers up to the view offset limit. Retained capacity counts toward this threshold and memory reservations. Memory-pressure emission still releases allocations. 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. |