refactor(hash-aggr): migrate ordered single aggregation

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)
     }
 }

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

@@ -0,0 +1,90 @@
+// 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.
+
+//! Aggregate table for single aggregation when raw input is ordered.
+//!
+//! See comments in [`super::ordered_partial_table`] for details.
+
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use datafusion_common::Result;
+
+use crate::aggregates::aggregate_hash_table::SingleMarker;
+use crate::aggregates::group_values::GroupByMetrics;
+use crate::aggregates::{AggregateExec, AggregateMode};
+
+use super::common::HashAggregateAccumulator;
+use super::common_ordered::OrderedAggregateTable;
+
+/// Implementation specific to single aggregation, where the table stores final
+/// aggregate values and the input rows are raw rows.
+///
+/// Example: `AVG(x) GROUP BY k`
+///
+/// - Aggregate table stores: `k, avg(x)`
+/// - Input rows: `k, x`
+///
+/// See comments at [`OrderedAggregateTable`] for details.
+impl OrderedAggregateTable<SingleMarker> {
+    pub(in crate::aggregates) fn new(
+        agg: &AggregateExec,
+        partition: usize,
+        output_schema: SchemaRef,
+        state_schema: SchemaRef,
+        batch_size: usize,
+    ) -> Result<Self> {
+        debug_assert!(matches!(
+            agg.mode,
+            AggregateMode::Single | AggregateMode::SinglePartitioned
+        ));
+
+        let input_schema = agg.input().schema();
+        let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition);
+        Self::new_for_mode(
+            agg,
+            &input_schema,
+            output_schema,
+            state_schema,
+            batch_size,
+            &agg.input_order_mode,
+            &agg.mode,
+            agg.filter_expr.iter().cloned().collect(),
+            group_by_metrics,
+        )
+    }
+
+    /// Aggregates one raw input batch and updates ordering information for any
+    /// newly observed groups.
+    pub(in crate::aggregates) fn aggregate_batch(
+        &mut self,
+        batch: &RecordBatch,
+    ) -> Result<()> {
+        let evaluated_batch = self.evaluate_batch(batch)?;
+        self.aggregate_evaluated_batch(
+            &evaluated_batch,
+            HashAggregateAccumulator::update_batch,
+        )
+    }
+
+    /// Emits the next batch of final aggregate values for groups proven complete
+    /// by the input ordering.
+    pub(in crate::aggregates) fn next_output_batch(
+        &mut self,
+    ) -> Result<Option<RecordBatch>> {
+        self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns)
+    }
+}