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 6aa246beedb9d..0c481a6b3bc0b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -15,10 +15,11 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::marker::PhantomData; use std::sync::Arc; -use arrow::array::{ArrayRef, AsArray, new_null_array}; +use arrow::array::{ArrayRef, AsArray, BooleanArray, new_null_array}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; @@ -35,6 +36,7 @@ use crate::aggregates::grouped_hash_stream::create_group_accumulator; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, + group_id_array, max_duplicate_ordinal, }; use super::AggregateTableMetrics; @@ -377,6 +379,83 @@ impl AggregateHashTable { state.batch_group_indices = Vec::new(); self.state = AggregateHashTableState::Outputting(state); } + + /// Creates the required empty grouping-set rows when the input is empty. + /// + /// For example, this query must still produce one grand-total group even if + /// `t` has no rows: + /// + /// ```sql + /// SELECT COUNT(v) + /// FROM t + /// GROUP BY GROUPING SETS (()); + /// ``` + /// + /// The synthetic row is filtered out before accumulator update so aggregates + /// see the same state they would see for an empty input, rather than a real + /// null-valued row. + /// + /// Only the raw-input tables (partial and single aggregation) call this + /// method: grouping sets are expanded while consuming raw rows, so the + /// state-input stages (final and partial-reduce aggregation) receive the + /// already expanded keys as plain group columns (see + /// [`PhysicalGroupBy::as_final`]) and never own a grouping set. + pub(super) fn init_empty_grouping_sets(&mut self) -> Result<()> { + let state = self.state.building_mut(); + if !state.group_by.has_grouping_set() || !state.group_values.is_empty() { + return Ok(()); + } + + let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); + let max_ordinal = max_duplicate_ordinal(state.group_by.groups()); + let mut ordinals: HashMap<&[bool], usize> = HashMap::new(); + let group_schema = state.group_by.group_schema(&self.input_schema)?; + let n_expr = state.group_by.expr().len(); + let mut any_interned = false; + + for group in state.group_by.groups() { + let ordinal = { + let entry = ordinals.entry(group.as_slice()).or_insert(0); + let ordinal = *entry; + *entry += 1; + ordinal + }; + + if !group.iter().all(|&is_null| is_null) { + continue; + } + + let mut cols: Vec = group_schema + .fields() + .iter() + .take(n_expr) + .map(|field| new_null_array(field.data_type(), 1)) + .collect(); + cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); + + state + .group_values + .intern(&cols, &mut state.batch_group_indices)?; + any_interned = true; + } + + if any_interned { + let total_groups = state.group_values.len(); + let false_filter = BooleanArray::from(vec![false]); + for (idx, acc) in state.accumulators.iter_mut().enumerate() { + let null_args = acc.null_arguments(&self.input_schema)?; + let values = EvaluatedAccumulatorArgs { + arguments: null_args, + filter: Some(Arc::new(false_filter.clone())), + }; + accumulator_metrics.time(idx, AccumulatorPhase::Update, || { + acc.update_batch(&values, &[0], total_groups) + })?; + } + } + + Ok(()) + } } /// State and argument information for a single Aggregate diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index 54997f0537b87..bbc51ae666ab7 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -15,22 +15,20 @@ // specific language governing permissions and limitations // under the License. -use std::collections::HashMap; use std::marker::PhantomData; use std::sync::Arc; -use arrow::array::{ArrayRef, BooleanArray, new_null_array}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, assert_eq_or_internal_err}; +use crate::aggregates::AggregateExec; use crate::aggregates::group_values::{AccumulatorPhase, new_group_values}; use crate::aggregates::order::GroupOrdering; -use crate::aggregates::{AggregateExec, group_id_array, max_duplicate_ordinal}; use super::common::{ AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, - EvaluatedAccumulatorArgs, HashAggregateAccumulator, PartialMarker, PartialSkipMarker, + HashAggregateAccumulator, PartialMarker, PartialSkipMarker, }; /// Implementation specific to partial aggregation, where the table stores @@ -125,77 +123,6 @@ impl AggregateHashTable { self.start_outputting(); Ok(()) } - - /// Creates the required empty grouping-set rows when the input is empty. - /// - /// For example, this query must still produce one grand-total group even if - /// `t` has no rows: - /// - /// ```sql - /// SELECT COUNT(v) - /// FROM t - /// GROUP BY GROUPING SETS (()); - /// ``` - /// - /// The synthetic row is filtered out before accumulator update so aggregates - /// see the same state they would see for an empty input, rather than a real - /// null-valued row. - fn init_empty_grouping_sets(&mut self) -> Result<()> { - let state = self.state.building_mut(); - if !state.group_by.has_grouping_set() || !state.group_values.is_empty() { - return Ok(()); - } - - let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics); - let max_ordinal = max_duplicate_ordinal(state.group_by.groups()); - let mut ordinals: HashMap<&[bool], usize> = HashMap::new(); - let group_schema = state.group_by.group_schema(&self.input_schema)?; - let n_expr = state.group_by.expr().len(); - let mut any_interned = false; - - for group in state.group_by.groups() { - let ordinal = { - let entry = ordinals.entry(group.as_slice()).or_insert(0); - let ordinal = *entry; - *entry += 1; - ordinal - }; - - if !group.iter().all(|&is_null| is_null) { - continue; - } - - let mut cols: Vec = group_schema - .fields() - .iter() - .take(n_expr) - .map(|field| new_null_array(field.data_type(), 1)) - .collect(); - cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); - - state - .group_values - .intern(&cols, &mut state.batch_group_indices)?; - any_interned = true; - } - - if any_interned { - let total_groups = state.group_values.len(); - let false_filter = BooleanArray::from(vec![false]); - for (idx, acc) in state.accumulators.iter_mut().enumerate() { - let null_args = acc.null_arguments(&self.input_schema)?; - let values = EvaluatedAccumulatorArgs { - arguments: null_args, - filter: Some(Arc::new(false_filter.clone())), - }; - accumulator_metrics.time(idx, AccumulatorPhase::Update, || { - acc.update_batch(&values, &[0], total_groups) - })?; - } - } - - Ok(()) - } } impl AggregateHashTable { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs index 1ff05fc79d224..2d7dc2a63d086 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs @@ -78,6 +78,7 @@ impl AggregateHashTable { } pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.init_empty_grouping_sets()?; self.start_outputting(); Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 340bf5cfc12d6..f67f8119bece3 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -117,6 +117,16 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metric /// remaining input batch is converted directly to partial aggregate state rows /// without inserting the rows into the grouped hash table. /// +/// # Feature: Grouping Sets +/// +/// `GROUPING SETS`, `CUBE` and `ROLLUP` are expanded in the partial stage: every +/// grouping set of an input batch is evaluated (with the grouping expressions +/// that are not part of the set replaced by `NULL`, plus an internal +/// `__grouping_id` column) and interned into the same hash table. The final +/// stage then merges the expanded keys as a plain group by. +/// +/// The partial aggregation skip optimization is disabled for grouping sets. +/// /// # Feature: Memory-limited Execution /// /// ## Partial Aggregation diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 6da7ee1018dc5..2b09df3edbc19 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1278,11 +1278,20 @@ impl AggregateExec { )?)) } + // # Grouping sets + // + // `GROUPING SETS`, `CUBE` and `ROLLUP` are expanded by the raw-input stages + // (partial and single aggregation), which evaluate every grouping set of + // each input batch into the same hash table. They are always planned with + // `InputOrderMode::Linear` (see `try_new_with_schema`), so they only reach + // the unordered hash streams. State-input stages (final and partial-reduce + // aggregation) must receive the expanded keys as a plain group by (see + // `PhysicalGroupBy::as_final`), which their predicates check explicitly. + fn should_use_partial_hash_stream(&self, _context: &TaskContext) -> bool { self.mode == AggregateMode::Partial && self.input_order_mode == InputOrderMode::Linear && !self.group_by.is_true_no_grouping() - && self.group_by.is_single() && self.limit_options_supported_by_hash_stream() } @@ -1293,7 +1302,6 @@ impl AggregateExec { self.mode == AggregateMode::Partial && self.input_order_mode != InputOrderMode::Linear && !self.group_by.is_true_no_grouping() - && self.group_by.is_single() && self.limit_options_supported_by_hash_stream() } @@ -1322,7 +1330,6 @@ impl AggregateExec { ) && self.limit_options.is_none() && self.input_order_mode == InputOrderMode::Linear && !self.group_by.is_true_no_grouping() - && self.group_by.is_single() } fn should_use_ordered_single_aggregate_stream(&self, _context: &TaskContext) -> bool { @@ -1332,7 +1339,6 @@ impl AggregateExec { ) && self.limit_options.is_none() && self.input_order_mode != InputOrderMode::Linear && !self.group_by.is_true_no_grouping() - && self.group_by.is_single() } fn should_use_ordered_final_aggregate_stream(&self, _context: &TaskContext) -> bool { @@ -3519,12 +3525,14 @@ mod tests { | 2 | 1.0 | 0 | 1 | | 3 | | 1 | 1 | | 3 | | 1 | 2 | - | 3 | 2.0 | 0 | 2 | + | 3 | 2.0 | 0 | 1 | + | 3 | 2.0 | 0 | 1 | | 3 | 3.0 | 0 | 1 | | 4 | | 1 | 1 | | 4 | | 1 | 2 | | 4 | 3.0 | 0 | 1 | - | 4 | 4.0 | 0 | 2 | + | 4 | 4.0 | 0 | 1 | + | 4 | 4.0 | 0 | 1 | +---+-----+---------------+-----------------+ " ); diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 7b8b08c15ba30..40412385efbc7 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -73,6 +73,14 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// This stream implements the complete aggregation without a partial/final /// split. It consumes raw input rows and emits final aggregate values. /// +/// # Grouping Sets +/// +/// `GROUPING SETS`, `CUBE` and `ROLLUP` are expanded while consuming raw input: +/// every grouping set of an input batch is evaluated and interned into the same +/// hash table, the same way [`super::hash_stream::PartialHashAggregateStream`] +/// does it. When spilling, the expanded keys are sorted and replayed as a plain +/// group by. +/// /// # Spilling /// /// During aggregation, group keys and states accumulate. If memory usage exceeds