refactor(hash-aggr): grouping sets routing

datafusion/core/tests/physical_optimizer/filter_pushdown.rs

@@ -2086,7 +2086,7 @@ fn test_pushdown_grouping_sets_filter_on_common_column() {
         -     DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true
       output:
         Ok:
-          - AggregateExec: mode=Final, gby=[(a@0 as a, b@1 as b), (NULL as a, b@1 as b)], aggr=[cnt], ordering_mode=PartiallySorted([1])
+          - AggregateExec: mode=Final, gby=[(a@0 as a, b@1 as b), (NULL as a, b@1 as b)], aggr=[cnt]
           -   DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=b@1 = bar
     "
     );
@@ -2372,7 +2372,7 @@ fn test_pushdown_through_aggregate_grouping_sets_with_reordered_input() {
         -       DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true
       output:
         Ok:
-          - AggregateExec: mode=Final, gby=[(a@1 as a, b@2 as b), (NULL as a, b@2 as b)], aggr=[cnt], ordering_mode=PartiallySorted([1])
+          - AggregateExec: mode=Final, gby=[(a@1 as a, b@2 as b), (NULL as a, b@2 as b)], aggr=[cnt]
           -   ProjectionExec: expr=[c@2 as c, a@0 as a, b@1 as b]
           -     DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=b@1 = bar
     "

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};
@@ -32,6 +33,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,
 };
 
 /// Marker for raw rows -> partial state aggregation.
@@ -213,6 +215,75 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
         Ok(())
     }
 
+    /// Creates the required empty grouping-set rows for a raw-input aggregate.
+    ///
+    /// 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 partial and single aggregation call this method;
+    /// final and partial-reduce aggregation consume states produced upstream.
+    pub(super) fn init_empty_grouping_sets_for_raw_input(&mut self) -> Result<()> {
+        let state = self.state.building_mut();
+        if !state.group_by.has_grouping_set() || !state.group_values.is_empty() {
+            return Ok(());
+        }
+
+        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 acc in state.accumulators.iter_mut() {
+                let null_args = acc.null_arguments(&self.input_schema)?;
+                let values = EvaluatedAccumulatorArgs {
+                    arguments: null_args,
+                    filter: Some(Arc::new(false_filter.clone())),
+                };
+                acc.update_batch(&values, &[0], total_groups)?;
+            }
+        }
+
+        Ok(())
+    }
+
     /// Materializes the full output once, then returns it downstream incrementally
     /// by slicing it into `batch_size` chunks.
     ///

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

@@ -38,20 +38,24 @@ use crate::aggregates::{
     evaluate_group_by,
 };
 
-use super::common::{AggregateAccumulator, EvaluatedAggregateBatch};
+use super::common::{
+    AggregateAccumulator, AggregateBatchFn, EvaluatedAggregateBatch,
+    MaterializeAccumulatorFn,
+};
 
-/// Aggregate table shared by the ordered partial and final paths.
+/// Aggregate table shared by the ordered single, partial and final paths.
 ///
 /// # Ordering optimization
 ///
 /// The table consumes input batches while `GroupOrdering` tracks which groups
 /// are proven complete. Completed groups can be emitted before the input stream
 /// ends, which keeps memory bounded by the active ordered key range.
 ///
-/// # Partial and final variant difference
+/// # Single, partial and final variant difference
 ///
 /// The partial and final aggregate tables implement the two stages of grouped
-/// aggregation. See
+/// aggregation, while the single aggregate table implements both stages in one
+/// table. See
 /// [`OrderedPartialAggregateStream`](crate::aggregates::ordered_partial_stream::OrderedPartialAggregateStream)
 /// for the high-level plan shape.
 ///
@@ -67,6 +71,11 @@ use super::common::{AggregateAccumulator, EvaluatedAggregateBatch};
 /// - Table stores: `k, sum(v), count(v)`
 /// - Output schema: `k, avg(v)`
 ///
+/// Single table ([`AggregateMode::Single`], with optional filter from query):
+/// - Input rows: `k, v`
+/// - Table stores: `k, sum(v), count(v)`
+/// - Output schema: `k, avg(v)`
+///
 /// # Marker Type
 ///
 /// `OrderedAggrMode` selects the aggregate semantics. For example,
@@ -75,7 +84,7 @@ use super::common::{AggregateAccumulator, EvaluatedAggregateBatch};
 /// `OrderedAggregateTable::<FinalMarker>::new_with_input_order(...)`
 /// consumes partial states and emits final values.
 ///
-/// Shared methods live on `impl<T>`; partial/final behavior lives on
+/// Shared methods live on `impl<T>`; single/partial/final behavior lives on
 /// marker-specific impls.
 pub(in crate::aggregates) struct OrderedAggregateTable<OrderedAggrMode> {
     /// Output schema: group columns followed by aggregate state or final values.
@@ -129,7 +138,7 @@ pub(super) struct OrderedAggregateTableBuffer {
 impl<AggrMode> OrderedAggregateTable<AggrMode> {
     #[expect(
         clippy::too_many_arguments,
-        reason = "keeps ordered partial and final table construction explicit"
+        reason = "keeps ordered single, partial and final table construction explicit"
     )]
     pub(super) fn new_for_mode(
         agg: &AggregateExec,
@@ -259,8 +268,8 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
     ///
     /// Unlike normal ordered emission, this operation is allowed to take the
     /// active (incomplete) groups. Partial aggregation can pass those states to
-    /// its final stage, while final aggregation sorts and spills them before
-    /// replay.
+    /// its final stage, while single and final aggregation sort and spill them
+    /// before replay.
     pub(in crate::aggregates) fn take_state_batch(
         &mut self,
     ) -> Result<Option<RecordBatch>> {
@@ -304,18 +313,18 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
             EmitTo::All => (EmitTo::First(self.batch_size), false),
         }
     }
-    /// Aggregates one evaluated input batch.
-    ///
-    /// This common utility is used by ordered partial and ordered final aggregation.
-    ///
-    /// # Argument: `is_final`
+
+    /// Aggregates one evaluated input batch after selecting the mode-specific
+    /// accumulator operation.
     ///
-    /// - `true`: merge partial aggregate states for final aggregation.
-    /// - `false`: update aggregate states from raw input for partial aggregation.
+    /// Each aggregation mode chooses a different `aggregate_fn` according to its
+    /// semantics. For example, partial aggregation takes raw inputs and updates
+    /// stored partial states, so it uses
+    /// [`datafusion_expr::GroupsAccumulator::update_batch`].
     pub(super) fn aggregate_evaluated_batch(
         &mut self,
         evaluated_batch: &EvaluatedAggregateBatch,
-        is_final: bool,
+        aggregate_fn: AggregateBatchFn,
     ) -> Result<()> {
         for group_values in &evaluated_batch.grouping_set_args {
             let starting_num_groups = self.buffer.group_values.len();
@@ -338,19 +347,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
                 .iter_mut()
                 .zip(evaluated_batch.accumulator_args.iter())
             {
-                if is_final {
-                    acc.merge_batch(
-                        values,
-                        &self.buffer.group_indices,
-                        total_num_groups,
-                    )?;
-                } else {
-                    acc.update_batch(
-                        values,
-                        &self.buffer.group_indices,
-                        total_num_groups,
-                    )?;
-                }
+                aggregate_fn(acc, values, &self.buffer.group_indices, total_num_groups)?;
             }
             drop(timer);
         }
@@ -361,15 +358,13 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
     /// Emits groups allowed by `GroupOrdering`, leaving only the current
     /// unfinished ordered-key range buffered.
     ///
-    /// This common utility is used by ordered partial and ordered final aggregation.
-    ///
-    /// # Argument: `is_final`
-    ///
-    /// - `true`: output final aggregate values.
-    /// - `false`: output partial accumulator states.
-    pub(super) fn next_output_batch_for_mode(
+    /// Each aggregation mode chooses a different `materialize_accumulator_fn`
+    /// according to its semantics. For example, partial aggregation emits
+    /// partial states to feed the final stage, so it uses
+    /// [`datafusion_expr::GroupsAccumulator::state`].
+    pub(super) fn next_output_batch_inner(
         &mut self,
-        is_final: bool,
+        materialize_accumulator_fn: MaterializeAccumulatorFn,
     ) -> Result<Option<RecordBatch>> {
         if self.buffer.group_values.is_empty() {
             return Ok(None);
@@ -394,11 +389,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         }
 
         for acc in &mut self.buffer.accumulators {
-            if is_final {
-                output.push(acc.evaluate(emit_to)?);
-            } else {
-                output.extend(acc.state(emit_to)?);
-            }
+            output.extend(materialize_accumulator_fn(acc, emit_to)?);
         }
         drop(timer);
 

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

@@ -20,6 +20,7 @@ mod common_ordered;
 mod final_table;
 mod ordered_final_table;
 mod ordered_partial_table;
+mod ordered_single_table;
 mod partial_reduce_table;
 mod partial_table;
 mod single_table;

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

@@ -30,6 +30,7 @@ use crate::aggregates::aggregate_hash_table::FinalMarker;
 use crate::aggregates::group_values::GroupByMetrics;
 use crate::aggregates::{AggregateExec, AggregateMode};
 
+use super::common::HashAggregateAccumulator;
 use super::common_ordered::OrderedAggregateTable;
 
 /// Implementation specific to final aggregation, where the table stores partial
@@ -73,13 +74,16 @@ impl OrderedAggregateTable<FinalMarker> {
         // `PhysicalGroupBy::as_final()` removes grouping sets while planning
         // final aggregation, so final ordered aggregation sees one grouping.
         debug_assert_eq!(evaluated_batch.grouping_set_args.len(), 1);
-        self.aggregate_evaluated_batch(&evaluated_batch, true)
+        self.aggregate_evaluated_batch(
+            &evaluated_batch,
+            HashAggregateAccumulator::merge_batch,
+        )
     }
 
     /// See comments in `ordered_partial_stream::next_output_batch`
     pub(in crate::aggregates) fn next_output_batch(
         &mut self,
     ) -> Result<Option<RecordBatch>> {
-        self.next_output_batch_for_mode(true)
+        self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns)
     }
 }

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

@@ -40,6 +40,7 @@ use crate::aggregates::{
     group_values::GroupByMetrics,
 };
 
+use super::common::HashAggregateAccumulator;
 use super::common_ordered::OrderedAggregateTable;
 
 /// Implementation specific to partial aggregation, where the table stores
@@ -81,7 +82,10 @@ impl OrderedAggregateTable<PartialMarker> {
         batch: &RecordBatch,
     ) -> Result<()> {
         let evaluated_batch = self.evaluate_batch(batch)?;
-        self.aggregate_evaluated_batch(&evaluated_batch, false)
+        self.aggregate_evaluated_batch(
+            &evaluated_batch,
+            HashAggregateAccumulator::update_batch,
+        )
     }
 
     /// Emits the next batch of partial state rows for groups proven complete by
@@ -101,6 +105,6 @@ impl OrderedAggregateTable<PartialMarker> {
     pub(in crate::aggregates) fn next_output_batch(
         &mut self,
     ) -> Result<Option<RecordBatch>> {
-        self.next_output_batch_for_mode(false)
+        self.next_output_batch_inner(HashAggregateAccumulator::state)
     }
 }