refactor(hash-aggr): route grouping sets to the migrated hash streams

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<AggrMode> AggregateHashTable<AggrMode> {
         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<ArrayRef> = 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

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<PartialMarker> {
         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<ArrayRef> = 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<PartialSkipMarker> {

datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs

@@ -78,6 +78,7 @@ impl AggregateHashTable<SingleMarker> {
     }
 
     pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> {
+        self.init_empty_grouping_sets()?;
         self.start_outputting();
         Ok(())
     }

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

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               |
             +---+-----+---------------+-----------------+
             "
             );

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