alamb commented on code in PR #23657:
URL: https://github.com/apache/datafusion/pull/23657#discussion_r3626983420
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -180,36 +497,183 @@ impl OrderedFinalAggregateStream {
next_state,
))
}
+ // Can't do early emit, continue aggregating.
Ok(None) => {
- // Ordered variant doesn't support memory-limited
- // execution, so it errors when memory reservation
fails.
- if let Err(e) =
self.reservation.try_resize(table.memory_size()) {
- return ControlFlow::Break((
- Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput {
table },
- ));
- }
-
- // Can't do early emit, continue aggregating.
ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput {
table,
+ spill_context,
})
}
Err(e) => ControlFlow::Break((
Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput { table },
+ OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context,
+ },
)),
}
}
Poll::Ready(Some(Err(e))) => ControlFlow::Break((
Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput { table },
+ OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context,
+ },
)),
Poll::Ready(None) => {
self.close_input();
- table.input_done();
-
ControlFlow::Continue(OrderedFinalAggregateState::DrainingFinal { table })
+ match spill_context {
+ Some(spill_context) if spill_context.has_spills() => {
+ ControlFlow::Continue(
+ OrderedFinalAggregateState::PreparingMergeInput {
+ table,
+ spill_context,
+ },
+ )
+ }
+ _ => {
+ table.input_done();
+ ControlFlow::Continue(
+ OrderedFinalAggregateState::ProducingOutput {
table },
+ )
+ }
+ }
+ }
+ }
+ }
+
+ /// Sorts and spills one complete in-memory state run, then resumes input.
+ ///
+ /// See comments at `poll_next()` for details.
+ ///
+ /// Returns the next operator state with control flow decision.
+ fn handle_spilling(
+ &mut self,
+ original_state: OrderedFinalAggregateState,
+ ) -> OrderedFinalAggregateStateTransition {
+ let OrderedFinalAggregateState::Spilling {
+ mut table,
+ mut spill_context,
+ } = original_state
+ else {
+ unreachable!("expected spilling state")
+ };
+
+ let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+ let timer = elapsed_compute.timer();
+ let mut result = spill_context.spill_table(&mut table);
+
+ // Spilling shrinks the aggregate table and releases its accumulated
+ // memory. Update the reservation accordingly.
+ if self.reservation.try_resize(table.memory_size()).is_err() {
+ // Fold spill error and reservation error into one
+ result = internal_err!("Reservation after spilling should succeed")
+ }
+
+ timer.done();
+
+ match result {
+ // Finished spilling the aggregate table, continue aggregating
from input
+ Ok(true) =>
ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context: Some(spill_context),
+ }),
+ Ok(false) => ControlFlow::Break((
+ Poll::Ready(Some(internal_err!(
+ "Ordered final aggregation entered Spilling with an empty
table"
+ ))),
+ OrderedFinalAggregateState::Done,
+ )),
+ Err(e) => ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::Done,
+ )),
+ }
+ }
+
+ /// 1. Spills the last in-memory run.
+ /// 2. Constructs a globally ordered input stream by applying a
sort-preserving
+ /// merge to all spills.
+ /// 3. Constructs a replay stream: an ordered aggregate stream over the
fully
+ /// ordered input constructed from the spills.
+ ///
+ /// See comments at `poll_next()` for details.
+ ///
+ /// Returns the next operator state with control flow decision.
+ fn handle_preparing_merge_input(
+ &mut self,
+ original_state: OrderedFinalAggregateState,
+ ) -> OrderedFinalAggregateStateTransition {
+ let OrderedFinalAggregateState::PreparingMergeInput {
+ mut table,
+ mut spill_context,
+ } = original_state
+ else {
+ unreachable!("expected preparing merge input state")
+ };
+
+ let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+ let timer = elapsed_compute.timer();
+ let replay = match spill_context.spill_table(&mut table) {
+ Ok(_) => {
+ let group_by_metrics = table.group_by_metrics();
+ drop(table);
+ match self.reservation.try_resize(0) {
+ Ok(()) => (*spill_context).into_replay_stream(
+ &self.baseline_metrics,
+ group_by_metrics,
+ self.reservation.new_empty(),
+ ),
+ Err(e) => Err(e),
+ }
+ }
+ Err(e) => Err(e),
+ };
+ timer.done();
+
+ match replay {
+ Ok(stream) => {
+
ControlFlow::Continue(OrderedFinalAggregateState::MergingSpills {
+ stream,
+ })
}
+ Err(e) => ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::Done,
+ )),
+ }
+ }
+
+ /// Forwards output from the fully ordered stream that consumes the merged
+ /// spill runs.
+ ///
+ /// See comments at `poll_next()` for details.
+ ///
+ /// Returns the next operator state with control flow decision.
+ fn handle_merging_spills(
+ &mut self,
+ cx: &mut Context<'_>,
+ original_state: OrderedFinalAggregateState,
+ ) -> OrderedFinalAggregateStateTransition {
+ let OrderedFinalAggregateState::MergingSpills { mut stream } =
original_state
+ else {
+ unreachable!("expected merging spills state")
Review Comment:
ditto here about internal error vs panic
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -64,6 +119,135 @@ type OrderedFinalAggregateStateTransition = ControlFlow<
OrderedFinalAggregateState,
>;
+impl OrderedFinalSpillContext {
+ fn new(
+ agg: &AggregateExec,
+ context: &Arc<TaskContext>,
+ partition: usize,
+ batch_size: usize,
+ input_order_mode: &InputOrderMode,
+ spill_schema: &SchemaRef,
+ spill_metrics: SpillMetrics,
+ ) -> Result<Self> {
+ let group_schema = agg.group_by.group_schema(spill_schema)?;
+ let output_ordering = agg.cache.output_ordering();
+ let InputOrderMode::PartiallySorted(order_indices) = input_order_mode
else {
+ return internal_err!("Ordered final spill requires partially
ordered input");
+ };
+ let spill_indices = order_indices.iter().copied().chain(
+ (0..group_schema.fields().len()).filter(|idx|
!order_indices.contains(idx)),
+ );
+ let spill_sort_exprs = spill_indices.map(|idx| {
+ let field = group_schema.field(idx);
+ let output_expr = Column::new(field.name(), idx);
+ let sort_options = output_ordering
+ .and_then(|ordering| ordering.get_sort_options(&output_expr))
+ .unwrap_or_default();
+ PhysicalSortExpr::new(Arc::new(output_expr), sort_options)
+ });
+ let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else {
+ return internal_err!("Ordered final spill expression is empty");
+ };
+
+ let spill_manager = SpillManager::new(
+ context.runtime_env(),
+ spill_metrics,
+ Arc::clone(spill_schema),
+ )
+ .with_compression_type(context.session_config().spill_compression());
+
+ Ok(Self {
+ agg: agg.clone(),
+ context: Arc::clone(context),
+ partition,
+ batch_size,
+ spill_expr,
+ spill_manager,
+ spills: vec![],
+ })
+ }
+
+ fn has_spills(&self) -> bool {
+ !self.spills.is_empty()
+ }
+
+ /// Sort and spill the aggregated groups. Memory reservation should be
cleared
Review Comment:
it might also help to explain somewhere what sort key the spill files are
sorted by (is it the partial sort prefix, or the full group key)? I think it is
the full group key.
##########
datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs:
##########
@@ -69,14 +70,50 @@ use crate::{InputOrderMode, RecordBatchStream,
SendableRecordBatchStream, metric
/// `k = 100`, it is safe to emit all groups with keys less than 100 because
the
/// input is ordered.
///
+/// # Memory Pressure and Spilling
+///
+/// ## Fully ordered case
+///
+/// If the input is ordered by every group key, for example:
+///
+/// - Input order: `a, b`
+/// - `GROUP BY`: `a, b`
+///
+/// Completed groups can be emitted as soon as the next group is observed.
Thus,
+/// only the current group remains active after completed groups are emitted,
and
+/// memory usage does not grow with the total number of groups.
+///
+/// If a memory reservation nevertheless fails, the stream returns the error
+/// directly, indicating an unexpected behavior.
+///
+/// ## Partially ordered case
+///
+/// If the input is ordered by only a subset of the group keys, for example:
+///
+/// - Input order: `a`
+/// - `GROUP BY`: `a, b`
+///
+/// If one `a` value contains many distinct `b` values, the table may
accumulate
+/// enough groups to exceed the memory limit.
+///
+/// - `OrderedPartialAggregateStream`: On reservation failure, it emits all
current
+/// intermediate states downstream and resets the table. The final stage can
+/// merge repeated `(a, b)` state rows, so no disk spill is required.
+/// - `OrderedFinalAggregateStream`: It cannot emit incomplete final results.
On
+/// reservation failure, it sorts the current intermediate states by the
complete
+/// group key and spills them as one run. After the input ends, it spills any
+/// remaining states, performs a sort-preserving merge of all runs, and
feeds the
+/// merged input into a fully ordered final aggregate stream.
+///
/// ## Implementation Note
///
/// This is intentionally kept simple and closely maps to
/// `GroupedHashAggregateStream` to finish the refactor sooner.
///
+/// More applicable optimizations are left to future work.
Review Comment:
this is somwhat mysterious 😆 maybe we can link to a ticket to track such
potential future optimizations (or just remove this comment)
##########
datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt:
##########
@@ -0,0 +1,201 @@
+# 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.
+
+# End-to-end tests for ordered aggregation under finite memory.
+
+# Result set more than 100 lines will be hashed
+hash-threshold 100
+
+statement ok
+SET datafusion.execution.target_partitions = 2
+
+statement ok
+SET datafusion.execution.batch_size = 128
+
+statement ok
+SET datafusion.optimizer.repartition_aggregations = true
+
+statement ok
+SET datafusion.optimizer.prefer_existing_sort = true
+
+statement ok
+SET datafusion.execution.enable_migration_aggregate = true
Review Comment:
This is already the default, right?
##########
datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs:
##########
@@ -69,14 +70,50 @@ use crate::{InputOrderMode, RecordBatchStream,
SendableRecordBatchStream, metric
/// `k = 100`, it is safe to emit all groups with keys less than 100 because
the
/// input is ordered.
///
+/// # Memory Pressure and Spilling
+///
+/// ## Fully ordered case
+///
+/// If the input is ordered by every group key, for example:
+///
+/// - Input order: `a, b`
+/// - `GROUP BY`: `a, b`
+///
+/// Completed groups can be emitted as soon as the next group is observed.
Thus,
+/// only the current group remains active after completed groups are emitted,
and
+/// memory usage does not grow with the total number of groups.
+///
+/// If a memory reservation nevertheless fails, the stream returns the error
+/// directly, indicating an unexpected behavior.
+///
+/// ## Partially ordered case
+///
+/// If the input is ordered by only a subset of the group keys, for example:
+///
+/// - Input order: `a`
+/// - `GROUP BY`: `a, b`
+///
+/// If one `a` value contains many distinct `b` values, the table may
accumulate
+/// enough groups to exceed the memory limit.
+///
+/// - `OrderedPartialAggregateStream`: On reservation failure, it emits all
current
Review Comment:
this is a great description
##########
datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt:
##########
@@ -0,0 +1,201 @@
+# Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
I actually ran just this test under coverage
```shell
cargo llvm-cov --html --test sqllogictests -- ordered_aggregate_spill.slt
```
And it has quite impressive coverage 👍
<img width="1536" height="271" alt="Image"
src="https://github.com/user-attachments/assets/5b0dcee7-4995-4fef-bbd7-1b9be5b906a4";
/>
##########
datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt:
##########
@@ -0,0 +1,201 @@
+# 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.
+
+# End-to-end tests for ordered aggregation under finite memory.
+
+# Result set more than 100 lines will be hashed
+hash-threshold 100
+
+statement ok
+SET datafusion.execution.target_partitions = 2
+
+statement ok
+SET datafusion.execution.batch_size = 128
+
+statement ok
+SET datafusion.optimizer.repartition_aggregations = true
+
+statement ok
+SET datafusion.optimizer.prefer_existing_sort = true
+
+statement ok
+SET datafusion.execution.enable_migration_aggregate = true
+
+statement ok
+SET datafusion.runtime.memory_limit = '1M'
+
+#
==================================================================================
+# Input is fully ordered by group keys (input order by (a,b), query is 'group
by a,b')
+#
==================================================================================
+
+# Fully ordered input uses the ordered partial and final streams without spill.
+query TT
+EXPLAIN ANALYZE
+SELECT v1, sum(v1 * 2)
+FROM generate_series(20000) AS t1(v1)
+GROUP BY v1
+----
+Plan with Metrics
+01)AggregateExec: mode=FinalPartitioned,<slt:ignore>ordering_mode=Sorted,
metrics=[<slt:ignore>spill_count=0,<slt:ignore>]
+02)--RepartitionExec:<slt:ignore>preserve_order=true<slt:ignore>
+03)----AggregateExec: mode=Partial,<slt:ignore>ordering_mode=Sorted,
metrics=[<slt:ignore>spill_count=0,<slt:ignore>]
+<slt:ignore>
+
+query II rowsort
+SELECT v1, sum(v1 * 2)
+FROM generate_series(20000) AS t1(v1)
+GROUP BY v1
+----
+40002 values hashing to 34c2b23730596cbd2489ed4a627c17d7
+
+# The same fully ordered query cannot spill and reports OOM under tighter
memory.
+statement ok
+SET datafusion.runtime.memory_limit = '1K'
+
+query error Resources exhausted
+SELECT v1, sum(v1 * 2)
+FROM generate_series(20000) AS t1(v1)
+GROUP BY v1
+
+#
==================================================================================
+# Input is partially ordered by group keys (input order by (a), query is
'group by a,b')
+#
+# Try different memory limits, ensure result is the same, but spill count
differ
+
+# HACK: check `spilled_bytes=x KB` to ensure it has spilled. If it has not
spilled,
+# the it shows `spilled_bytes = 0B`. Should better check spill_count, but it's
not
+# stable due to ordered hash repartition, and `sqllogictest` don't support
regex.
+#
==================================================================================
+
+statement ok
+SET datafusion.runtime.memory_limit = '2M'
+
+statement ok
+SET datafusion.optimizer.enable_round_robin_repartition = false
+
+# Round 1: The input is partially ordered and does not spill with a 2 MB limit.
+query TT
+EXPLAIN ANALYZE
+SELECT round(v1, -4), v1 % 5000, sum(v1 * 2)
+FROM generate_series(20000) AS t1(v1)
+GROUP BY round(v1, -4), v1 % 5000
+----
+Plan with Metrics
+01)AggregateExec: mode=FinalPartitioned,<slt:ignore>aggr=[sum(t1.v1 *
Int64(2))], ordering_mode=PartiallySorted([0]),
metrics=[<slt:ignore>spill_count=0,<slt:ignore>]
+02)--RepartitionExec:<slt:ignore>input_partitions=1,
maintains_sort_order=true<slt:ignore>
+03)----AggregateExec:
mode=Partial,<slt:ignore>ordering_mode=PartiallySorted([0]),
metrics=[<slt:ignore>spill_count=0,<slt:ignore>]
+<slt:ignore>
+
+# All rounds should have the same result hash
+query III rowsort
+SELECT round(v1, -4), v1 % 5000, sum(v1 * 2)
+FROM generate_series(20000) AS t1(v1)
+GROUP BY round(v1, -4), v1 % 5000
+----
+45000 values hashing to e6ece4b4b86e6152a1c785e90ab5ba12
+
+# Round 2: The same query spills five times with a 600 KB limit.
+statement ok
+SET datafusion.runtime.memory_limit = '600K'
+
+query TT
+EXPLAIN ANALYZE
+SELECT round(v1, -4), v1 % 5000, sum(v1 * 2)
+FROM generate_series(20000) AS t1(v1)
+GROUP BY round(v1, -4), v1 % 5000
+----
+Plan with Metrics
+01)AggregateExec: mode=FinalPartitioned,<slt:ignore>aggr=[sum(t1.v1 *
Int64(2))], ordering_mode=PartiallySorted([0]),
metrics=[<slt:ignore>spilled_bytes=<slt:ignore> KB,<slt:ignore>]
+02)--RepartitionExec:<slt:ignore>input_partitions=1,
maintains_sort_order=true<slt:ignore>
+03)----AggregateExec:
mode=Partial,<slt:ignore>ordering_mode=PartiallySorted([0]),
metrics=[<slt:ignore>spill_count=0,<slt:ignore>]
+<slt:ignore>
+
+# All rounds should have the same result hash
Review Comment:
it would be nice if there was some way to programatically ensure the hashses
are the same, but I think the comments you have inline here are good enough
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -47,13 +69,46 @@ pub(crate) struct OrderedFinalAggregateStream {
state: Option<OrderedFinalAggregateState>,
}
+/// Spill configuration and accumulated runs for partially ordered final
+/// aggregation. Each file is one fully group-key-sorted intermediate-state
run;
+/// all runs are replayed together after the original input is exhausted.
+struct OrderedFinalSpillContext {
+ /// Aggregate configuration
+ agg: AggregateExec,
+ /// Task context
+ context: Arc<TaskContext>,
+ /// Original partition index
+ partition: usize,
+ /// Target batch size from configuration
+ batch_size: usize,
+ /// Full group-key ordering, such ordering with be kept in: a) individual
spill
+ /// files, b) order after final merging and streaming aggregate
+ spill_expr: LexOrdering,
+ /// Spill I/O and metrics manager.
+ spill_manager: SpillManager,
+ /// Fully sorted spill runs waiting to be merged.
+ spills: Vec<SortedSpillFile>,
+}
+
/// See comments at `poll_next()` for details.
enum OrderedFinalAggregateState {
ReadingInput {
table: OrderedAggregateTable<FinalMarker>,
+ spill_context: Option<Box<OrderedFinalSpillContext>>,
},
- DrainingFinal {
+ Spilling {
table: OrderedAggregateTable<FinalMarker>,
+ spill_context: Box<OrderedFinalSpillContext>,
+ },
+ ProducingOutput {
+ table: OrderedAggregateTable<FinalMarker>,
+ },
+ PreparingMergeInput {
Review Comment:
(I see now it is documted below at `impl Stream for
OrderedFinalAggregateStream {` -- maybe we could leave a link / hit in the
comments to look there for the state transitions
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
