perf: materialize completed groups once in OrderedSingleAggregateStream

datafusion/physical-plan/benches/ordered_group_values.rs

@@ -19,7 +19,8 @@
 //! existing streaming hash table on identical batches. Include emission and
 //! cross-batch continuation, not just lookup. Input preparation is untimed.
 //! The physical-plan benchmark includes aggregation but excludes SQL planning,
-//! sorting and I/O; the data already has a proven ordering.
+//! sorting and I/O; the data already has a proven ordering. A partially
+//! ordered case covers ordered prefixes that complete many groups at once.
 
 use std::hint::black_box;
 use std::sync::Arc;
@@ -186,6 +187,8 @@ fn check_case(schema: &SchemaRef, batches: &[Vec<ArrayRef>]) -> (usize, usize) {
 fn aggregate_plan(
     schema: &SchemaRef,
     keys: Vec<Vec<ArrayRef>>,
+    sort_columns: &[&str],
+    input_order_mode: &InputOrderMode,
 ) -> Arc<dyn ExecutionPlan> {
     let mut fields = schema.fields().to_vec();
     fields.push(Arc::new(Field::new("v", DataType::Int64, false)));
@@ -197,10 +200,11 @@ fn aggregate_plan(
             RecordBatch::try_new(Arc::clone(&schema), cols).unwrap()
         })
         .collect::<Vec<_>>();
-    let ordering = LexOrdering::new(vec![
-        PhysicalSortExpr::new_default(col("a", &schema).unwrap()),
-        PhysicalSortExpr::new_default(col("b", &schema).unwrap()),
-    ])
+    let ordering = LexOrdering::new(
+        sort_columns
+            .iter()
+            .map(|name| PhysicalSortExpr::new_default(col(name, &schema).unwrap())),
+    )
     .unwrap();
     let input =
         TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None).unwrap();
@@ -228,7 +232,7 @@ fn aggregate_plan(
         schema,
     )
     .unwrap();
-    assert_eq!(plan.input_order_mode(), &InputOrderMode::Sorted);
+    assert_eq!(plan.input_order_mode(), input_order_mode);
     Arc::new(plan)
 }
 
@@ -241,7 +245,8 @@ fn aggregation(c: &mut Criterion) {
     for strings in [false, true] {
         for run_length in [1, 8, 128, 8192] {
             let (schema, keys) = inputs(run_length, 8192, strings);
-            let plan = aggregate_plan(&schema, keys);
+            let plan =
+                aggregate_plan(&schema, keys, &["a", "b"], &InputOrderMode::Sorted);
             let name =
                 format!("{}_run{run_length}", if strings { "string" } else { "int" });
             group.bench_function(name, |b| {
@@ -261,5 +266,57 @@ fn aggregation(c: &mut Criterion) {
     group.finish();
 }
 
-criterion_group!(benches, grouping, aggregation);
+/// Input ordered by `a` only. Each `a` value has 32768 distinct `b` values, so
+/// every `a` boundary completes four default batches of groups at once.
+fn partially_ordered_aggregation(c: &mut Criterion) {
+    const GROUPS_PER_KEY: usize = 32768;
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("a", DataType::Int32, false),
+        Field::new("b", DataType::Int32, false),
+    ]));
+    let keys = (0..ROWS)
+        .step_by(8192)
+        .map(|start| {
+            let rows = start..(start + 8192).min(ROWS);
+            let a: ArrayRef = Arc::new(Int32Array::from_iter_values(
+                rows.clone().map(|row| (row / GROUPS_PER_KEY) as i32),
+            ));
+            let b: ArrayRef = Arc::new(Int32Array::from_iter_values(
+                rows.map(|row| (row % GROUPS_PER_KEY) as i32),
+            ));
+            vec![a, b]
+        })
+        .collect();
+    let plan = aggregate_plan(
+        &schema,
+        keys,
+        &["a"],
+        &InputOrderMode::PartiallySorted(vec![0]),
+    );
+    let runtime = Runtime::new().unwrap();
+    let mut group = c.benchmark_group("partially_ordered_aggregate_exec");
+    group.sample_size(10);
+    group.warm_up_time(Duration::from_millis(250));
+    group.measurement_time(Duration::from_secs(1));
+    group.bench_function("groups_per_key32768", |b| {
+        b.iter(|| {
+            black_box(
+                runtime
+                    .block_on(collect(
+                        Arc::clone(&plan),
+                        Arc::new(TaskContext::default()),
+                    ))
+                    .unwrap(),
+            )
+        });
+    });
+    group.finish();
+}
+
+criterion_group!(
+    benches,
+    grouping,
+    aggregation,
+    partially_ordered_aggregation
+);
 criterion_main!(benches);

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

@@ -24,7 +24,6 @@ use std::sync::Arc;
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
-use datafusion_common::assert_or_internal_err;
 use datafusion_execution::memory_pool::proxy::VecAllocExt;
 use datafusion_expr::{AggregateMetrics, EmitTo};
 
@@ -128,9 +127,6 @@ pub(in crate::aggregates) struct OrderedAggregateTable<OrderedAggrMode> {
     /// to pass through or spill its current state.
     pub(super) state_schema: SchemaRef,
 
-    /// Maximum rows per emitted output batch, from config `batch_size`.
-    pub(super) batch_size: usize,
-
     /// Grouping and accumulator-specific timing metrics.
     pub(super) group_by_metrics: GroupByMetrics,
 
@@ -188,17 +184,11 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         input_schema: &SchemaRef,
         output_schema: SchemaRef,
         state_schema: SchemaRef,
-        batch_size: usize,
         input_order_mode: &InputOrderMode,
         aggregate_mode: &AggregateMode,
         filters: Vec<Option<Arc<dyn PhysicalExpr>>>,
         metrics: OrderedAggregateTableMetrics,
     ) -> Result<Self> {
-        assert_or_internal_err!(
-            batch_size > 0,
-            "OrderedAggregateTable requires config batch_size >= 1"
-        );
-
         let group_ordering = GroupOrdering::try_new(input_order_mode)?;
         let group_schema = agg.group_by().group_schema(input_schema)?;
         let group_values = new_group_values(group_schema, &group_ordering)?;
@@ -229,7 +219,6 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         Ok(Self {
             output_schema,
             state_schema,
-            batch_size,
             group_by_metrics: metrics.group_by,
             aggregate_argument_metrics: metrics.aggregate_arguments,
             aggregate_accumulator_metrics: metrics.accumulator,
@@ -418,38 +407,8 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         Ok(())
     }
 
-    /// Emits groups allowed by `GroupOrdering`, leaving only the current
-    /// unfinished ordered-key range buffered.
-    ///
-    /// 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,
-        materialize_accumulator_fn: MaterializeAccumulatorFn,
-        accumulator_phase: AccumulatorPhase,
-    ) -> Result<Option<RecordBatch>> {
-        if self.buffer.group_values.is_empty() {
-            return Ok(None);
-        }
-
-        let Some(emit_to) = self.buffer.group_ordering.emit_to() else {
-            return Ok(None);
-        };
-        let emit_to = match emit_to {
-            EmitTo::First(n) => EmitTo::First(n.min(self.batch_size)),
-            EmitTo::All if self.num_groups() > self.batch_size => {
-                EmitTo::First(self.batch_size)
-            }
-            EmitTo::All => EmitTo::All,
-        };
-        self.materialize_groups(emit_to, materialize_accumulator_fn, accumulator_phase)
-            .map(Some)
-    }
-
     /// Removes the selected groups once and materializes their output columns.
-    /// The caller chooses the completed prefix and any output-size limit.
+    /// The caller chooses the completed prefix.
     pub(super) fn materialize_groups(
         &mut self,
         emit_to: EmitTo,
@@ -459,11 +418,9 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics);
         let output = self.group_by_metrics.time_emitting(|| {
             let mut output = self.buffer.group_values.emit(emit_to)?;
-            // EOF can also emit a prefix when a caller limits its batch size,
-            // but the completed ordering state no longer tracks group indexes.
-            if let EmitTo::First(n) = emit_to
-                && matches!(self.buffer.group_ordering.emit_to(), Some(EmitTo::First(_)))
-            {
+            // `EmitTo::All` is only used after `input_done`, when the ordering
+            // state no longer tracks group indexes.
+            if let EmitTo::First(n) = emit_to {
                 self.buffer.group_ordering.remove_groups(n);
             }
 

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

@@ -46,7 +46,6 @@ impl OrderedAggregateTable<FinalMarker> {
         agg: &AggregateExec,
         input_schema: &SchemaRef,
         output_schema: SchemaRef,
-        batch_size: usize,
         input_order_mode: &InputOrderMode,
         metrics: OrderedAggregateTableMetrics,
     ) -> Result<Self> {
@@ -55,7 +54,6 @@ impl OrderedAggregateTable<FinalMarker> {
             input_schema,
             output_schema,
             Arc::clone(input_schema),
-            batch_size,
             input_order_mode,
             &AggregateMode::Final,
             vec![None; agg.aggr_expr().len()],

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

@@ -57,7 +57,6 @@ impl OrderedAggregateTable<PartialMarker> {
         agg: &AggregateExec,
         partition: usize,
         output_schema: SchemaRef,
-        batch_size: usize,
     ) -> Result<Self> {
         let input_schema = agg.input().schema();
         let state_schema = Arc::clone(&output_schema);
@@ -67,7 +66,6 @@ impl OrderedAggregateTable<PartialMarker> {
             &input_schema,
             output_schema,
             state_schema,
-            batch_size,
             &agg.input_order_mode,
             &AggregateMode::Partial,
             agg.filter_expr().to_vec(),

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

@@ -44,7 +44,6 @@ impl OrderedAggregateTable<SingleMarker> {
         partition: usize,
         output_schema: SchemaRef,
         state_schema: SchemaRef,
-        batch_size: usize,
     ) -> Result<Self> {
         debug_assert!(matches!(
             agg.mode,
@@ -58,7 +57,6 @@ impl OrderedAggregateTable<SingleMarker> {
             &input_schema,
             output_schema,
             state_schema,
-            batch_size,
             &agg.input_order_mode,
             &agg.mode,
             agg.filter_expr().to_vec(),
@@ -80,14 +78,24 @@ impl OrderedAggregateTable<SingleMarker> {
         )
     }
 
-    /// Emits the next batch of final aggregate values for groups proven complete
-    /// by the input ordering.
-    pub(in crate::aggregates) fn next_output_batch(
+    /// Materializes final results for all groups proven complete by the input
+    /// ordering, leaving the active ordered-key range in the table.
+    ///
+    /// Returns None if there are no completed groups.
+    pub(in crate::aggregates) fn take_completed_result_batch(
         &mut self,
     ) -> Result<Option<RecordBatch>> {
-        self.next_output_batch_inner(
+        if self.is_empty() {
+            return Ok(None);
+        }
+        let Some(emit_to) = self.group_ordering().emit_to() else {
+            return Ok(None);
+        };
+        self.materialize_groups(
+            emit_to,
             HashAggregateAccumulator::evaluate_to_columns,
             AccumulatorPhase::Evaluate,
         )
+        .map(Some)
     }
 }

datafusion/physical-plan/src/aggregates/ordered_final_stream.rs

@@ -189,7 +189,6 @@ impl OrderedFinalAggregateStream {
             agg,
             &input_schema,
             Arc::clone(&schema),
-            batch_size,
             input_order_mode,
             metrics,
         )?;

datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs

@@ -161,7 +161,6 @@ impl OrderedPartialAggregateStream {
             agg,
             partition,
             Arc::clone(&schema),
-            batch_size,
         )?;
         let reservation =
             MemoryConsumer::new(format!("OrderedPartialAggregateStream[{partition}]"))