diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 173b69dcbdec..38b5dce7e02a 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1136,6 +1136,30 @@ 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 + + /// (experimental) Flush unordered partial aggregation after a complete input + /// batch when the group count reaches batch_size. Overrides both other flush + /// thresholds and reserves supported group-key and accumulator storage for + /// twice batch_size groups, initially and after each threshold flush. + /// Requires a single grouping set and the same eligibility as the other + /// flush thresholds. String payload buffers use their normal growth policy. + pub partial_aggregation_flush_batch: bool, default = 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. diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs index 1c004cb70b93..e4a4f034bd90 100644 --- a/datafusion/expr-common/src/groups_accumulator.rs +++ b/datafusion/expr-common/src/groups_accumulator.rs @@ -381,6 +381,10 @@ pub trait GroupsAccumulator: Send + std::any::Any { /// /// May be expensive; check the implementation before calling on hot paths. fn size(&self) -> usize; + + /// Reserves row-addressed state for `capacity` groups where supported. + /// This hint must not change the logical number of groups. + fn reserve_groups(&mut self, _capacity: usize) {} } #[cfg(test)] diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs index 0453124528a0..618d403637ff 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs @@ -216,6 +216,11 @@ where Ok(vec![Arc::new(state_values)]) } + fn reserve_groups(&mut self, capacity: usize) { + self.values + .reserve_exact(capacity.saturating_sub(self.values.len())); + } + fn size(&self) -> usize { self.values.capacity() * size_of::() + self.null_state.size() } diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs index 85b63a9a1797..a81d68b75ba3 100644 --- a/datafusion/functions-aggregate/src/average.rs +++ b/datafusion/functions-aggregate/src/average.rs @@ -1247,6 +1247,13 @@ where Ok(vec![Arc::new(counts) as ArrayRef, Arc::new(sums)]) } + fn reserve_groups(&mut self, capacity: usize) { + self.counts + .reserve_exact(capacity.saturating_sub(self.counts.len())); + self.sums + .reserve_exact(capacity.saturating_sub(self.sums.len())); + } + fn size(&self) -> usize { // Heap buffers self.counts.capacity() * size_of::() diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index 7c3bfdceefdf..0247ccfb89fb 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -823,6 +823,11 @@ impl GroupsAccumulator for CountGroupsAccumulator { Ok(vec![state_array]) } + fn reserve_groups(&mut self, capacity: usize) { + self.counts + .reserve_exact(capacity.saturating_sub(self.counts.len())); + } + fn size(&self) -> usize { self.counts.heap_size(&mut DFHeapSizeCtx::default()) } diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 725a1e3c5b27..aa0e68e1f87c 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -30,8 +30,12 @@ use datafusion_common::hash_utils::create_hashes; use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::{Result, exec_err}; use std::fmt::Debug; +use std::mem; use std::sync::Arc; +const BYTE_VIEW_STARTING_BLOCK_SIZE: usize = 8 * 1024; +const BYTE_VIEW_MAX_BLOCK_SIZE: usize = 2 * 1024 * 1024; + /// HashSet optimized for storing string or binary values that can produce that /// the final set as a `GenericBinaryViewArray` with minimal copies. #[derive(Debug)] @@ -125,9 +129,6 @@ impl ArrowBytesViewSet { /// This map is used by the special `COUNT DISTINCT` aggregate function to /// store the distinct values, and by the `GROUP BY` operator to store /// group values when they are a single string array. -/// Max size of the in-progress buffer before flushing to completed buffers -const BYTE_VIEW_MAX_BLOCK_SIZE: usize = 2 * 1024 * 1024; - pub struct ArrowBytesViewMap where V: Debug + PartialEq + Eq + Clone + Copy + Default, @@ -147,6 +148,9 @@ where in_progress: Vec, /// Completed buffers containing string data completed: Vec, + /// Allocation target, doubled for each new payload block up to 2 MiB. + /// Zero keeps lazy payload growth for the small per-group distinct maps. + block_size: usize, /// random state used to generate hashes random_state: RandomState, @@ -191,18 +195,51 @@ where views: Vec::new(), in_progress: Vec::new(), completed: Vec::new(), + block_size: if map_capacity == 0 { + 0 + } else { + BYTE_VIEW_STARTING_BLOCK_SIZE + }, random_state: RandomState::default(), hashes_buffer: vec![], null: None, } } + /// Reserves group slots without preallocating variable-length string payloads. + pub fn reserve_groups(&mut self, capacity: usize) { + if self.block_size == 0 && capacity != 0 { + self.block_size = BYTE_VIEW_STARTING_BLOCK_SIZE; + } + self.map + .reserve(capacity.saturating_sub(self.map.len()), |entry| entry.hash); + self.views + .reserve_exact(capacity.saturating_sub(self.views.len())); + self.hashes_buffer + .reserve_exact(capacity.saturating_sub(self.hashes_buffer.len())); + } + + /// Emits keys, retains the hash table and scratch buffer, and reserves fixed + /// view capacity. Payload blocks restart their exponential allocation growth. + pub fn take_with_capacity(&mut self, capacity: usize) -> ArrayRef { + let mut outgoing = Self::new(self.output_type); + mem::swap(self, &mut outgoing); + mem::swap(&mut self.map, &mut outgoing.map); + self.map.clear(); + mem::swap(&mut self.hashes_buffer, &mut outgoing.hashes_buffer); + self.hashes_buffer.clear(); + self.initial_map_capacity = outgoing.initial_map_capacity; + mem::swap(&mut self.random_state, &mut outgoing.random_state); + self.reserve_groups(capacity); + outgoing.into_state() + } + /// Return the contents of this map and replace it with a new empty map with /// the same output type pub fn take(&mut self) -> Self { let mut new_self = Self::with_capacity(self.output_type, self.initial_map_capacity); - std::mem::swap(self, &mut new_self); + mem::swap(self, &mut new_self); new_self } @@ -445,7 +482,7 @@ where pub fn into_state(mut self) -> ArrayRef { // Flush any remaining in-progress buffer if !self.in_progress.is_empty() { - let flushed = std::mem::take(&mut self.in_progress); + let flushed = mem::take(&mut self.in_progress); self.completed.push(Buffer::from_vec(flushed)); } @@ -537,13 +574,23 @@ where let view = if len <= 12 { make_view(value, 0, 0) } else { - // Ensure buffer is big enough - if self.in_progress.len() + len > BYTE_VIEW_MAX_BLOCK_SIZE { - let flushed = std::mem::replace( - &mut self.in_progress, - Vec::with_capacity(BYTE_VIEW_MAX_BLOCK_SIZE), - ); - self.completed.push(Buffer::from_vec(flushed)); + let block_limit = if self.block_size == 0 { + BYTE_VIEW_MAX_BLOCK_SIZE + } else { + self.in_progress.capacity() + }; + if self.in_progress.len() + len > block_limit { + if !self.in_progress.is_empty() { + let flushed = mem::take(&mut self.in_progress); + self.completed.push(Buffer::from_vec(flushed)); + } + let capacity = if self.block_size == 0 { + BYTE_VIEW_MAX_BLOCK_SIZE + } else { + self.block_size = (self.block_size * 2).min(BYTE_VIEW_MAX_BLOCK_SIZE); + len.max(self.block_size) + }; + self.in_progress = Vec::with_capacity(capacity); } let buffer_index = self.completed.len() as u32; @@ -947,6 +994,10 @@ mod tests { map.views.reserve_exact(1); map.completed.shrink_to_fit(); map.completed.reserve_exact(1); + let completed = &mut map.completed[0]; + let mut padded = Vec::with_capacity(completed.len() + 16); + padded.extend_from_slice(completed); + *completed = Buffer::from_vec(padded); // The map owns these allocations; `values` and its Arrow buffers remain external. assert!(map.views.capacity() > map.views.len()); 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 020196e7da02..049e6fe8aa5e 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,27 @@ 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_with_capacity(0) + } + + /// Reserves fixed row capacity without changing the number of groups. + pub(in crate::aggregates) fn reserve_groups(&mut self, capacity: usize) { + let state = self.state.building_mut(); + state.group_values.reserve_groups(capacity); + state + .batch_group_indices + .reserve_exact(capacity.saturating_sub(state.batch_group_indices.len())); + for acc in &mut state.accumulators { + acc.accumulator.reserve_groups(capacity); + } + } + + /// Emits all states and optionally reserves a fixed capacity for the next table. + /// A zero capacity uses the normal memory-releasing path. + pub(in crate::aggregates) fn take_state_batch_with_capacity( + &mut self, + capacity: usize, ) -> Result> { let state_schema = Arc::clone(&self.state_schema); let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); @@ -382,7 +403,11 @@ impl AggregateHashTable { } let output = group_by_metrics.time_emitting(|| { - let mut output = state.group_values.emit(EmitTo::All)?; + let mut output = if capacity > 0 { + state.group_values.emit_with_capacity(capacity)? + } else { + state.group_values.emit(EmitTo::All)? + }; for (idx, acc) in state.accumulators.iter_mut().enumerate() { output.extend(accumulator_metrics.time( idx, @@ -399,9 +424,15 @@ 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); state.batch_group_indices.clear(); - state.batch_group_indices.shrink_to_fit(); + if capacity == 0 { + state.group_values.clear_shrink(0); + state.batch_group_indices.shrink_to_fit(); + } else { + for acc in &mut state.accumulators { + acc.accumulator.reserve_groups(capacity); + } + } 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 19d3b94f8a72..707b74225257 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -111,6 +111,18 @@ pub trait GroupValues: Send { /// Emits the group values fn emit(&mut self, emit_to: EmitTo) -> Result>; + /// Reserves room for at least `capacity` groups in supported storage layouts. + /// Does not change the number of groups or the string payload growth policy. + fn reserve_groups(&mut self, _capacity: usize) {} + + /// Emits all keys and reserves a fixed capacity for the next table. + fn emit_with_capacity(&mut self, capacity: usize) -> Result> { + let output = self.emit(EmitTo::All)?; + self.clear_shrink(capacity); + self.reserve_groups(capacity); + Ok(output) + } + /// Materializes selected group values without changing the stored values or /// their group indices. /// 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 24eb94476325..78edb2130fe1 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 @@ -29,9 +29,10 @@ use datafusion_common::Result; use datafusion_common::utils::split_vec_min_alloc; use datafusion_expr::GroupSelection; use std::marker::PhantomData; -use std::mem::{replace, size_of}; +use std::mem::{self, replace, size_of}; use std::sync::Arc; +const BYTE_VIEW_STARTING_BLOCK_SIZE: usize = 8 * 1024; const BYTE_VIEW_MAX_BLOCK_SIZE: usize = 2 * 1024 * 1024; /// An implementation of [`GroupColumn`] for binary view and utf8 view types. @@ -55,19 +56,16 @@ pub struct ByteViewGroupValueBuilder { /// The progressing block /// /// New values will be inserted into it until its capacity - /// is not enough(detail can see `max_block_size`). + /// cannot hold the next value. in_progress: Vec, /// The completed blocks completed: Vec, - /// The max size of `in_progress` - /// - /// `in_progress` will be flushed into `completed`, and create new `in_progress` - /// when found its remaining capacity(`max_block_size` - `len(in_progress)`), - /// is no enough to store the appended value. - /// - /// Currently it is fixed at 2MB. + /// Allocation target, doubled for each new payload block. + block_size: usize, + + /// Maximum allocation target. A single larger value gets its own block. max_block_size: usize, /// Nulls @@ -89,6 +87,7 @@ impl ByteViewGroupValueBuilder { views: Vec::new(), in_progress: Vec::new(), completed: Vec::new(), + block_size: BYTE_VIEW_STARTING_BLOCK_SIZE, max_block_size: BYTE_VIEW_MAX_BLOCK_SIZE, nulls: NullBufferBuilder::empty(), _phantom: PhantomData {}, @@ -98,6 +97,7 @@ impl ByteViewGroupValueBuilder { /// Set the max block size fn with_max_block_size(mut self, max_block_size: usize) -> Self { self.max_block_size = max_block_size; + self.block_size = self.block_size.min(max_block_size); self } @@ -236,16 +236,12 @@ impl ByteViewGroupValueBuilder { fn ensure_in_progress_big_enough(&mut self, value_len: usize) { debug_assert!(value_len > 12); - let require_cap = self.in_progress.len() + value_len; - - // If current block isn't big enough, flush it and create a new in progress block - if require_cap > self.max_block_size { - let flushed_block = replace( - &mut self.in_progress, - Vec::with_capacity(self.max_block_size), - ); - let buffer = Buffer::from_vec(flushed_block); - self.completed.push(buffer); + if value_len > self.in_progress.capacity() - self.in_progress.len() { + if !self.in_progress.is_empty() { + self.flush_in_progress(); + } + self.block_size = (self.block_size * 2).min(self.max_block_size); + self.in_progress = Vec::with_capacity(value_len.max(self.block_size)); } } @@ -564,10 +560,7 @@ impl ByteViewGroupValueBuilder { } fn flush_in_progress(&mut self) { - let flushed_block = replace( - &mut self.in_progress, - Vec::with_capacity(self.max_block_size), - ); + let flushed_block = mem::take(&mut self.in_progress); let buffer = Buffer::from_vec(flushed_block); self.completed.push(buffer); } @@ -626,6 +619,14 @@ impl GroupColumn for ByteViewGroupValueBuilder { self.vectorized_append_inner(array, rows) } + fn reserve_groups(&mut self, capacity: usize) { + self.views + .reserve_exact(capacity.saturating_sub(self.views.len())); + if self.nulls.is_empty() { + self.nulls = NullBufferBuilder::new(capacity); + } + } + fn len(&self) -> usize { self.views.len() } 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 68d56f00a9de..3a17e14202ae 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 @@ -95,6 +95,9 @@ pub trait GroupColumn: Send + Sync { /// The vectorized version `append_val` fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()>; + /// Reserves row-addressed storage without changing the logical length. + fn reserve_groups(&mut self, _capacity: usize) {} + /// Returns the number of rows stored in this builder fn len(&self) -> usize; @@ -1083,6 +1086,44 @@ impl GroupValues for GroupValuesColumn { } } + fn emit_with_capacity(&mut self, capacity: usize) -> Result> { + let output = self.emit(EmitTo::All)?; + 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.reserve_groups(capacity); + Ok(output) + } + + fn reserve_groups(&mut self, capacity: usize) { + self.map + .reserve(capacity.saturating_sub(self.map.len()), |entry| entry.0); + self.map_size = self.map.capacity() * size_of::<(u64, usize)>(); + self.hashes_buffer + .reserve_exact(capacity.saturating_sub(self.hashes_buffer.len())); + for column in &mut self.group_values { + column.reserve_groups(capacity); + } + let buffers = &mut self.vectorized_operation_buffers; + buffers + .append_row_indices + .reserve_exact(capacity.saturating_sub(buffers.append_row_indices.len())); + buffers + .equal_to_row_indices + .reserve_exact(capacity.saturating_sub(buffers.equal_to_row_indices.len())); + buffers + .equal_to_group_indices + .reserve_exact(capacity.saturating_sub(buffers.equal_to_group_indices.len())); + buffers + .remaining_row_indices + .reserve_exact(capacity.saturating_sub(buffers.remaining_row_indices.len())); + buffers + .equal_to_results + .reserve(capacity.saturating_sub(buffers.equal_to_results.len())); + } + fn size(&self) -> usize { let group_values_size = self.group_values.allocated_size() + self 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 ebd1e071cf8d..d67dd5b32b80 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 @@ -255,6 +255,14 @@ where Ok(()) } + fn reserve_groups(&mut self, capacity: usize) { + self.group_values + .reserve_exact(capacity.saturating_sub(self.group_values.len())); + if self.nulls.is_empty() { + self.nulls = NullBufferBuilder::new(capacity); + } + } + fn len(&self) -> usize { self.group_values.len() } 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 a2e3573ce889..541c1da61ba6 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 @@ -78,6 +78,18 @@ impl GroupValues for GroupValuesBytesView { Ok(()) } + fn reserve_groups(&mut self, capacity: usize) { + self.map.reserve_groups(capacity); + } + + fn emit_with_capacity( + &mut self, + capacity: usize, + ) -> datafusion_common::Result> { + self.num_groups = 0; + Ok(vec![self.map.take_with_capacity(capacity)]) + } + fn size(&self) -> usize { self.map.size() + size_of::() } 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 eaeb83e0cd71..04719e41fe56 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 @@ -177,6 +177,13 @@ where Ok(()) } + fn reserve_groups(&mut self, capacity: usize) { + self.map + .reserve(capacity.saturating_sub(self.map.len()), |entry| entry.1); + self.values + .reserve_exact(capacity.saturating_sub(self.values.len())); + } + fn size(&self) -> usize { size_of::() + self.map.capacity() * size_of::<(usize, u64)>() diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 12babfe59cf9..666002966908 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,24 @@ 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, + /// Fixed group capacity reserved for batch-size flushing; zero disables. + reserve_capacity: 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,12 +250,52 @@ impl PartialHashAggregateStream { let early_emit_count = MetricBuilder::new(&agg.metrics).counter("early_emit_count", partition); - let hash_table = AggregateHashTable::::new( + let execution_options = &context.session_config().options().execution; + let flush_batch = execution_options.partial_aggregation_flush_batch + && agg.group_by().is_single(); + let byte_threshold = if flush_batch { + 0 + } else { + execution_options.partial_aggregation_flush_bytes + }; + let group_threshold = if flush_batch { + batch_size + } else { + 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, + reserve_capacity: if flush_batch { + batch_size.saturating_mul(2) + } else { + 0 + }, + flushed_groups: 0, + flush_count: MetricBuilder::new(&agg.metrics) + .counter("table_flush_count", partition), + }); + + let mut hash_table = AggregateHashTable::::new( agg, partition, Arc::clone(&schema), batch_size, )?; + if let Some(flush) = &table_flush + && flush.reserve_capacity > 0 + { + hash_table.reserve_groups(flush.reserve_capacity); + } let skip_aggregation_probe = if agg.group_by().is_single() { let options = &context.session_config().options().execution; let probe_ratio_threshold = @@ -279,8 +337,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,21 +374,50 @@ 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 reserve_capacity = self + .table_flush + .as_ref() + .filter(|_| { + last_state == HandleInputResult::FlushThresholdReached + }) + .map_or(0, |flush| flush.reserve_capacity); + let materialized_group_states = hash_table + .take_state_batch_with_capacity(reserve_capacity)? + .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, - &mut emitter, - hash_table.memory_size(), - ) - .await?; + if last_state == HandleInputResult::FlushThresholdReached { + self.reduction_factor + .add_part(materialized_group_states.num_rows()); + self.reservation.try_resize(hash_table.memory_size())?; + emitter + .emit( + materialized_group_states + .record_output(&self.baseline_metrics), + ) + .await; + } else { + self.emit_on_memory_pressure( + materialized_group_states, + &mut emitter, + hash_table.memory_size(), + ) + .await?; + } } } } @@ -410,7 +498,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 +515,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 +1130,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 +1224,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 +1290,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 000000000000..78c12767bdfd --- /dev/null +++ b/datafusion/sqllogictest/test_files/aggregate_partial_flush.slt @@ -0,0 +1,304 @@ +# 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; + +# Fixed target batch group count overrides the byte and row thresholds. +statement ok +set datafusion.execution.batch_size = 1024; + +statement ok +set datafusion.execution.partial_aggregation_flush_batch = true; + +statement ok +set datafusion.execution.partial_aggregation_flush_bytes = 1; + +statement ok +set datafusion.execution.partial_aggregation_flush_rows = 1; + +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 + +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 + +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 + +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 + +statement ok +reset datafusion.execution.partial_aggregation_flush_batch; + +statement ok +reset datafusion.execution.partial_aggregation_flush_bytes; + +statement ok +reset datafusion.execution.partial_aggregation_flush_rows; + +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 d1898a691d13..55f475f0984c 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -272,6 +272,9 @@ 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_batch false +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 +439,9 @@ 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_batch false (experimental) Flush unordered partial aggregation after a complete input batch when the group count reaches batch_size. Overrides both other flush thresholds and reserves supported group-key and accumulator storage for twice batch_size groups, initially and after each threshold flush. Requires a single grouping set and the same eligibility as the other flush thresholds. String payload buffers use their normal growth policy. +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 16b96a224434..3a25633abcf2 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -143,6 +143,9 @@ 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.partial_aggregation_flush_batch | false | (experimental) Flush unordered partial aggregation after a complete input batch when the group count reaches batch_size. Overrides both other flush thresholds and reserves supported group-key and accumulator storage for twice batch_size groups, initially and after each threshold flush. Requires a single grouping set and the same eligibility as the other flush thresholds. String payload buffers use their normal growth policy. | | 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. |