You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/07/24 01:09:17 UTC
(datafusion) 01/01: refactor(hash-aggr): Support spilling for ordered aggregation (#23657)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23657-9caaf4b20a9b6a0cd63783f6e92e1d7f5761a9e0
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit d7ec10596b7fcf155ea011827890c42c453126b3
Author: Yongting You <20...@gmail.com>
AuthorDate: Fri Jul 24 09:08:53 2026 +0800
refactor(hash-aggr): Support spilling for ordered aggregation (#23657)
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
Part of https://github.com/apache/datafusion/issues/22710
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly
in the issue, then this section is not needed.
Clearly explaining why changes are proposed helps reviewers understand
your changes and offer better suggestions.
-->
This PR implements larger-than-memory execution for ordered partial and
ordered final aggregation. These two modes are implemented first
because:
- #22710 splits ordered and unordered aggregation into separate
implementations.
- Unordered aggregation internally reuses ordered aggregation for its
larger-than-memory execution logic. So larger-than-memory support for
ordered aggregation become the groundwork for other aggregation modes.
The high-level plan for two-stage ordered hash aggregation is:
```
AggregateExec(mode=final, ordered)
--RepartitionExec(hash)
----AggregateExec(mode=partial, ordered)
```
#### Group keys are fully ordered (input ordered by `a, b`; query uses
`GROUP BY a, b`)
Execution must use bounded memory because only one group is active at
any given time. Therefore, the implementation simply returns an
execution error if a memory reservation fails.
#### Group keys are partially ordered (input ordered by `a`; query uses
`GROUP BY a, b`)
At any given time, we must hold all groups for a particular value of `a`
in memory, so it is possible to run out of memory. This is handled as
follows:
If the partial aggregate runs out of memory, it directly materializes
the accumulated groups as aggregate states and emits them early to the
final stage.
The final aggregate continues aggregating until its memory budget is
reached, then:
1. Sorts the aggregated batches by the group keys and writes them to
disk.
2. Repeats this process until the input is exhausted.
3. Constructs a sort-preserving merge stream over all spill files.
4. Constructs an `OrderedFinalAggregateStream` using the sort-preserving
merge stream as its input. The new stream's group keys are fully
ordered, whereas those of the original stream were only partially
ordered. It then evaluates this ordered stream to produce the final
result.
### Design Tradeoffs
Note that the same implementation is currently shared by two execution
paths, controlled by flags: aggregation with fully ordered group keys
and aggregation with partially ordered group keys. This is not an ideal
pattern. However:
- Fully ordered aggregation can later be split into a separate
implementation with potentially more aggressive optimizations—for
example, it does not require a hash table. The the current
implementation only have to handle partially-ordered cases.
- For now, the behavior remains consistent with the original
implementation, allowing us to complete the migration sooner.
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
The core changes are in `ordered_partial_stream.rs` and
`ordered_final_stream.rs`, for the above algorithm. See code comments
for implementation details.
## Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
6. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
After implementing this change, I checked the coverage using `cargo
llvm-cov` and found that this feature was supported by the original
implementation but was not covered by either regular tests or fuzz
tests. This is a bit of a horror story.
I think the old implementation is heavily multiplexed, so its overall
line coverage is probably good. However, spilling for ordered
aggregation follows a specific execution path that was never exercised.
Therefore, this PR adds `slt` test coverage. I also think we need more
fuzzing cases (filed https://github.com/apache/datafusion/issues/23658)
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
---
.../aggregate_hash_table/common_ordered.rs | 59 +-
.../aggregate_hash_table/ordered_final_table.rs | 8 +-
.../aggregate_hash_table/ordered_partial_table.rs | 8 +-
datafusion/physical-plan/src/aggregates/mod.rs | 25 +-
.../physical-plan/src/aggregates/order/full.rs | 5 +
.../physical-plan/src/aggregates/order/mod.rs | 14 +
.../physical-plan/src/aggregates/order/partial.rs | 6 +
.../src/aggregates/ordered_final_stream.rs | 645 +++++++++++++++++++--
.../src/aggregates/ordered_partial_stream.rs | 144 ++++-
.../test_files/ordered_aggregate_spill.slt | 201 +++++++
10 files changed, 1015 insertions(+), 100 deletions(-)
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs
index c83303c51d..2293e7b1b8 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs
@@ -81,6 +81,10 @@ pub(in crate::aggregates) struct OrderedAggregateTable<OrderedAggrMode> {
/// Output schema: group columns followed by aggregate state or final values.
pub(super) output_schema: SchemaRef,
+ /// Intermediate-state schema used when memory pressure requires the table
+ /// 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,
@@ -129,13 +133,14 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
)]
pub(super) fn new_for_mode(
agg: &AggregateExec,
- partition: usize,
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>>>,
+ group_by_metrics: GroupByMetrics,
) -> Result<Self> {
assert_or_internal_err!(
batch_size > 0,
@@ -168,8 +173,9 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
Ok(Self {
output_schema,
+ state_schema,
batch_size,
- group_by_metrics: GroupByMetrics::new(&agg.metrics, partition),
+ group_by_metrics,
buffer: OrderedAggregateTableBuffer {
group_by: Arc::clone(&agg.group_by),
group_ordering,
@@ -217,9 +223,19 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
self.buffer.group_ordering.input_done();
}
+ /// Returns the ordering state used to decide how memory pressure is handled.
+ pub(in crate::aggregates) fn group_ordering(&self) -> &GroupOrdering {
+ &self.buffer.group_ordering
+ }
+
+ /// Number of groups currently buffered.
+ pub(in crate::aggregates) fn num_groups(&self) -> usize {
+ self.buffer.group_values.len()
+ }
+
/// Check if there is zero groups accumulated so far.
pub(in crate::aggregates) fn is_empty(&self) -> bool {
- self.buffer.group_values.is_empty()
+ self.num_groups() == 0
}
/// All internal buffer's memory size.
@@ -234,6 +250,43 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
+ self.buffer.group_indices.allocated_size()
}
+ pub(in crate::aggregates) fn group_by_metrics(&self) -> GroupByMetrics {
+ self.group_by_metrics.clone()
+ }
+
+ /// Takes every intermediate aggregate state and resets the table so it can
+ /// continue with a new ordered input segment.
+ ///
+ /// 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.
+ pub(in crate::aggregates) fn take_state_batch(
+ &mut self,
+ ) -> Result<Option<RecordBatch>> {
+ if self.buffer.group_values.is_empty() {
+ return Ok(None);
+ }
+
+ let mut output = self.buffer.group_values.emit(EmitTo::All)?;
+ for acc in &mut self.buffer.accumulators {
+ output.extend(acc.state(EmitTo::All)?);
+ }
+
+ let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?;
+ debug_assert!(batch.num_rows() > 0);
+
+ // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the
+ // key/index buffers too so the memory reservation can be released
+ // before the batch is passed downstream or sorted for spilling.
+ self.buffer.group_values.clear_shrink(0);
+ self.buffer.group_indices.clear();
+ self.buffer.group_indices.shrink_to_fit();
+ self.buffer.group_ordering.reset();
+
+ Ok(Some(batch))
+ }
+
/// Returns the [`EmitTo`], clamped to the specified batch size
///
/// Returns `(emit_to, should_remove_groups)`, where `emit_to` is the number
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs
index b7e3fd38ed..fd064ebffe 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs
@@ -19,12 +19,15 @@
//!
//! See comments in [`super::ordered_partial_table`] for details.
+use std::sync::Arc;
+
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use datafusion_common::Result;
use crate::InputOrderMode;
use crate::aggregates::aggregate_hash_table::FinalMarker;
+use crate::aggregates::group_values::GroupByMetrics;
use crate::aggregates::{AggregateExec, AggregateMode};
use super::common_ordered::OrderedAggregateTable;
@@ -41,21 +44,22 @@ use super::common_ordered::OrderedAggregateTable;
impl OrderedAggregateTable<FinalMarker> {
pub(in crate::aggregates) fn new_with_input_order(
agg: &AggregateExec,
- partition: usize,
input_schema: &SchemaRef,
output_schema: SchemaRef,
batch_size: usize,
input_order_mode: &InputOrderMode,
+ group_by_metrics: GroupByMetrics,
) -> Result<Self> {
Self::new_for_mode(
agg,
- partition,
input_schema,
output_schema,
+ Arc::clone(input_schema),
batch_size,
input_order_mode,
&AggregateMode::Final,
vec![None; agg.aggr_expr.len()],
+ group_by_metrics,
)
}
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs
index 033c14056a..a04e4dda8f 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs
@@ -29,12 +29,15 @@
//! The implementation is separated from other aggregate tables because this
//! execution path is likely to be optimized further in the future.
+use std::sync::Arc;
+
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use datafusion_common::Result;
use crate::aggregates::{
AggregateExec, AggregateMode, aggregate_hash_table::PartialMarker,
+ group_values::GroupByMetrics,
};
use super::common_ordered::OrderedAggregateTable;
@@ -56,15 +59,18 @@ impl OrderedAggregateTable<PartialMarker> {
batch_size: usize,
) -> Result<Self> {
let input_schema = agg.input().schema();
+ let state_schema = Arc::clone(&output_schema);
+ let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition);
Self::new_for_mode(
agg,
- partition,
&input_schema,
output_schema,
+ state_schema,
batch_size,
&agg.input_order_mode,
&AggregateMode::Partial,
agg.filter_expr.iter().cloned().collect(),
+ group_by_metrics,
)
}
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index 30e0ad2469..7d93cd7398 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -1235,12 +1235,10 @@ impl AggregateExec {
&& self.limit_options_supported_by_hash_stream()
}
- fn should_use_ordered_partial_aggregate_stream(&self, context: &TaskContext) -> bool {
- // TODO: implement memory-limited path and remove this limitation
- if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) {
- return false;
- }
-
+ fn should_use_ordered_partial_aggregate_stream(
+ &self,
+ _context: &TaskContext,
+ ) -> bool {
self.mode == AggregateMode::Partial
&& self.input_order_mode != InputOrderMode::Linear
&& !self.group_by.is_true_no_grouping()
@@ -1291,12 +1289,7 @@ impl AggregateExec {
&& self.group_by.is_single()
}
- fn should_use_ordered_final_aggregate_stream(&self, context: &TaskContext) -> bool {
- // TODO: implement memory-limited path and remove this limitation
- if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) {
- return false;
- }
-
+ fn should_use_ordered_final_aggregate_stream(&self, _context: &TaskContext) -> bool {
matches!(
self.mode,
AggregateMode::Final | AggregateMode::FinalPartitioned
@@ -3922,10 +3915,10 @@ mod tests {
+----------+-----------+-------------------------+
");
- // Ordered streams don't implement memory limits yet.
+ // Ordered partial aggregation supports finite memory.
let finite_memory_task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?;
let stream = aggregate.execute_typed(0, &finite_memory_task_ctx)?;
- assert!(matches!(stream, StreamType::GroupedHash(_)));
+ assert!(matches!(stream, StreamType::OrderedPartialAggregate(_)));
Ok(())
}
@@ -3999,10 +3992,10 @@ mod tests {
+-----+--------------+
");
- // Ordered streams don't implement memory limits yet.
+ // Ordered final aggregation supports finite memory.
let finite_memory_task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?;
let stream = final_aggregate.execute_typed(0, &finite_memory_task_ctx)?;
- assert!(matches!(stream, StreamType::GroupedHash(_)));
+ assert!(matches!(stream, StreamType::OrderedFinalAggregate(_)));
Ok(())
}
diff --git a/datafusion/physical-plan/src/aggregates/order/full.rs b/datafusion/physical-plan/src/aggregates/order/full.rs
index eb98611f79..ca818d6a2d 100644
--- a/datafusion/physical-plan/src/aggregates/order/full.rs
+++ b/datafusion/physical-plan/src/aggregates/order/full.rs
@@ -115,6 +115,11 @@ impl GroupOrderingFull {
self.state = State::Complete;
}
+ /// Starts tracking a new fully ordered input segment.
+ pub fn reset(&mut self) {
+ self.state = State::Start;
+ }
+
/// Called when new groups are added in a batch. See documentation
/// on [`super::GroupOrdering::new_groups`]
pub fn new_groups(&mut self, total_num_groups: usize) {
diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs
index 97fbd519c8..259411b00b 100644
--- a/datafusion/physical-plan/src/aggregates/order/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/order/mod.rs
@@ -93,6 +93,20 @@ impl GroupOrdering {
}
}
+ /// Resets the ordering state while preserving the configured ordering mode.
+ ///
+ /// Ordered partial aggregation uses this after passing intermediate states
+ /// downstream, and ordered final aggregation uses it after spilling a run.
+ /// In both cases the hash table is empty and can start tracking the next
+ /// input batch from a fresh ordering state.
+ pub fn reset(&mut self) {
+ match self {
+ GroupOrdering::None => {}
+ GroupOrdering::Partial(partial) => partial.reset(),
+ GroupOrdering::Full(full) => full.reset(),
+ }
+ }
+
/// Removes the first `n` groups from the internal state, shifting all
/// existing indexes down by `n`.
pub fn remove_groups(&mut self, n: usize) {
diff --git a/datafusion/physical-plan/src/aggregates/order/partial.rs b/datafusion/physical-plan/src/aggregates/order/partial.rs
index 476551a7ca..1603bb6d07 100644
--- a/datafusion/physical-plan/src/aggregates/order/partial.rs
+++ b/datafusion/physical-plan/src/aggregates/order/partial.rs
@@ -186,6 +186,12 @@ impl GroupOrderingPartial {
};
}
+ /// Starts tracking a new ordered input segment with the same sort-key
+ /// columns.
+ pub fn reset(&mut self) {
+ self.state = State::Start;
+ }
+
fn updated_sort_key(
current_sort: usize,
sort_key: Option<Vec<ScalarValue>>,
diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs
index 89653e05ab..071c1e9011 100644
--- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs
@@ -23,22 +23,44 @@ use std::task::{Context, Poll};
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
-use datafusion_common::Result;
+use datafusion_common::{DataFusionError, Result, internal_err};
use datafusion_execution::TaskContext;
use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
+use datafusion_physical_expr::PhysicalSortExpr;
+use datafusion_physical_expr::expressions::Column;
+use datafusion_physical_expr_common::sort_expr::LexOrdering;
use futures::stream::{Stream, StreamExt};
use super::AggregateExec;
use super::aggregate_hash_table::{FinalMarker, OrderedAggregateTable};
+use super::group_values::GroupByMetrics;
use crate::aggregates::AggregateMode;
use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics};
+use crate::sorts::IncrementalSortIterator;
+use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder};
+use crate::spill::spill_manager::SpillManager;
use crate::stream::EmptyRecordBatchStream;
use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream};
/// Final aggregate stream for `InputOrderMode::Sorted` and
/// `InputOrderMode::PartiallySorted`.
///
-/// See comments at [`super::ordered_partial_stream`] for details.
+/// See comments at [`super::ordered_partial_stream::OrderedPartialAggregateStream`] for details.
+///
+/// # Spilling
+///
+/// This section is only for implementation notes, for background, see [`super::ordered_partial_stream::OrderedPartialAggregateStream`]
+///
+/// For partially sorted input, spilling works as follows:
+///
+/// - Reserve the table footprint plus one `u32` sort index per buffered group. The
+/// extra index array is used in later sorting before spilling.
+/// - On memory pressure, materialize all group states into one batch.
+/// - Use [`IncrementalSortIterator`] to compute the full-batch index, then
+/// materialize and write one sorted `batch_size` slice at a time. The original
+/// batch and full index remain live until the run is written.
+/// - After input ends, merge the sorted runs and replay them through a fully
+/// ordered final aggregate stream.
pub(crate) struct OrderedFinalAggregateStream {
schema: SchemaRef,
input: SendableRecordBatchStream,
@@ -47,14 +69,50 @@ pub(crate) struct OrderedFinalAggregateStream {
state: Option<OrderedFinalAggregateState>,
}
+/// Spill configuration and accumulated runs for partially ordered final
+/// aggregation.
+///
+/// Each spill event drains all currently buffered groups, sorts their intermediate
+/// states by the full group key, and writes them to one spill file. All files are
+/// merged and replayed after the original input ends.
+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>>,
+ },
+ Spilling {
+ table: OrderedAggregateTable<FinalMarker>,
+ spill_context: Box<OrderedFinalSpillContext>,
},
- DrainingFinal {
+ ProducingOutput {
table: OrderedAggregateTable<FinalMarker>,
},
+ PreparingMergeInput {
+ table: OrderedAggregateTable<FinalMarker>,
+ spill_context: Box<OrderedFinalSpillContext>,
+ },
+ MergingSpills {
+ stream: SendableRecordBatchStream,
+ },
Done,
}
@@ -64,6 +122,134 @@ 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()
+ }
+
+ /// Sorts and spills the aggregated groups. Memory reservation should be updated
+ /// by the caller.
+ ///
+ /// Individual spill files are ordered by the `group by` keys.
+ ///
+ /// See [`OrderedFinalAggregateStream`] for spilling details.
+ fn spill_table(
+ &mut self,
+ table: &mut OrderedAggregateTable<FinalMarker>,
+ ) -> Result<()> {
+ let Some(batch) = table.take_state_batch()? else {
+ return Ok(());
+ };
+
+ let sorted_iter =
+ IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size);
+ let spill_file = self
+ .spill_manager
+ .spill_record_batch_iter_and_return_max_batch_memory(
+ sorted_iter,
+ "OrderedFinalAggregateSpill",
+ )?;
+
+ let Some((file, max_record_batch_memory)) = spill_file else {
+ return internal_err!("Ordered final aggregation produced an empty spill");
+ };
+
+ self.spills.push(SortedSpillFile {
+ file,
+ max_record_batch_memory,
+ });
+
+ Ok(())
+ }
+
+ /// Merges every sorted run and finalizes it through the fully ordered path.
+ fn into_replay_stream(
+ self,
+ baseline_metrics: &BaselineMetrics,
+ group_by_metrics: GroupByMetrics,
+ reservation: MemoryReservation,
+ ) -> Result<SendableRecordBatchStream> {
+ let Self {
+ agg,
+ context,
+ partition,
+ batch_size,
+ spill_expr,
+ spill_manager,
+ spills,
+ } = self;
+
+ let spill_schema = Arc::clone(spill_manager.schema());
+ let merged = StreamingMergeBuilder::new()
+ .with_schema(spill_schema)
+ .with_spill_manager(spill_manager)
+ .with_sorted_spill_files(spills)
+ .with_expressions(&spill_expr)
+ .with_metrics(baseline_metrics.intermediate())
+ .with_batch_size(batch_size)
+ .with_reservation(reservation)
+ .build()?;
+ let replay = OrderedFinalAggregateStream::new_with_input_and_metrics(
+ &agg,
+ &context,
+ partition,
+ merged,
+ &InputOrderMode::Sorted,
+ baseline_metrics.clone(),
+ group_by_metrics,
+ None,
+ )?;
+ Ok(Box::pin(replay))
+ }
+}
+
impl OrderedFinalAggregateStream {
pub fn new(
agg: &AggregateExec,
@@ -86,6 +272,35 @@ impl OrderedFinalAggregateStream {
partition: usize,
input: SendableRecordBatchStream,
input_order_mode: &InputOrderMode,
+ ) -> Result<Self> {
+ let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
+ let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition);
+ let spill_metrics = SpillMetrics::new(&agg.metrics, partition);
+ Self::new_with_input_and_metrics(
+ agg,
+ context,
+ partition,
+ input,
+ input_order_mode,
+ baseline_metrics,
+ group_by_metrics,
+ Some(spill_metrics),
+ )
+ }
+
+ #[expect(
+ clippy::too_many_arguments,
+ reason = "keeps replay metric reuse explicit"
+ )]
+ fn new_with_input_and_metrics(
+ agg: &AggregateExec,
+ context: &Arc<TaskContext>,
+ partition: usize,
+ input: SendableRecordBatchStream,
+ input_order_mode: &InputOrderMode,
+ baseline_metrics: BaselineMetrics,
+ group_by_metrics: GroupByMetrics,
+ spill_metrics: Option<SpillMetrics>,
) -> Result<Self> {
debug_assert!(matches!(
agg.mode,
@@ -96,21 +311,37 @@ impl OrderedFinalAggregateStream {
let schema = Arc::clone(&agg.schema);
let input_schema = input.schema();
let batch_size = context.session_config().batch_size();
- let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
- // Preserve the existing aggregate metric surface for this plan node.
- let _spill_metrics = SpillMetrics::new(&agg.metrics, partition);
+ let can_spill = matches!(input_order_mode, InputOrderMode::PartiallySorted(_))
+ && context.runtime_env().disk_manager.tmp_files_enabled();
+ let spill_context = if can_spill {
+ let Some(spill_metrics) = spill_metrics else {
+ return internal_err!("Spillable ordered final stream requires metrics");
+ };
+ Some(Box::new(OrderedFinalSpillContext::new(
+ agg,
+ context,
+ partition,
+ batch_size,
+ input_order_mode,
+ &input_schema,
+ spill_metrics,
+ )?))
+ } else {
+ None
+ };
let table = OrderedAggregateTable::<FinalMarker>::new_with_input_order(
agg,
- partition,
&input_schema,
Arc::clone(&schema),
batch_size,
input_order_mode,
+ group_by_metrics,
)?;
let reservation =
MemoryConsumer::new(format!("OrderedFinalAggregateStream[{partition}]"))
+ .with_can_spill(can_spill)
.register(context.memory_pool());
Ok(Self {
@@ -118,7 +349,10 @@ impl OrderedFinalAggregateStream {
input,
reservation,
baseline_metrics,
- state: Some(OrderedFinalAggregateState::ReadingInput { table }),
+ state: Some(OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context,
+ }),
})
}
@@ -127,6 +361,27 @@ impl OrderedFinalAggregateStream {
self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
}
+ fn break_with_internal_err(message: &str) -> OrderedFinalAggregateStateTransition {
+ ControlFlow::Break((
+ Poll::Ready(Some(internal_err!("{message}"))),
+ OrderedFinalAggregateState::Done,
+ ))
+ }
+
+ /// Reserve memory for the current aggregate table.
+ fn reservation_size_for_table(
+ table: &OrderedAggregateTable<FinalMarker>,
+ spill_context: Option<&OrderedFinalSpillContext>,
+ ) -> usize {
+ let table_size = table.memory_size();
+ if spill_context.is_some() {
+ // See `OrderedFinalAggregateStream` comments for how is it estimated
+ table_size.saturating_add(table.num_groups().saturating_mul(size_of::<u32>()))
+ } else {
+ table_size
+ }
+ }
+
/// Consumes one ordered partial-state input batch, then immediately emits
/// finalized groups if the ordering proves any group is ready.
///
@@ -138,15 +393,23 @@ impl OrderedFinalAggregateStream {
cx: &mut Context<'_>,
original_state: OrderedFinalAggregateState,
) -> OrderedFinalAggregateStateTransition {
- let OrderedFinalAggregateState::ReadingInput { mut table } = original_state
+ let OrderedFinalAggregateState::ReadingInput {
+ mut table,
+ spill_context,
+ } = original_state
else {
- unreachable!("expected reading input state")
+ return Self::break_with_internal_err(
+ "Ordered final aggregate stream expected ReadingInput state",
+ );
};
match self.input.poll_next_unpin(cx) {
Poll::Pending => ControlFlow::Break((
Poll::Pending,
- OrderedFinalAggregateState::ReadingInput { table },
+ OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context,
+ },
)),
Poll::Ready(Some(Ok(batch))) => {
let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
@@ -157,21 +420,86 @@ impl OrderedFinalAggregateStream {
if let Err(e) = result {
return ControlFlow::Break((
Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput { table },
+ OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context,
+ },
));
}
+ // Check memory reservation, and potentially spill.
let timer = elapsed_compute.timer();
- let result = table.next_output_batch();
+ let resize_result =
+ self.reservation
+ .try_resize(Self::reservation_size_for_table(
+ &table,
+ spill_context.as_deref(),
+ ));
timer.done();
+ match resize_result {
+ Ok(()) => {}
+ Err(e @ DataFusionError::ResourcesExhausted(_)) => {
+ let Some(spill_context) = spill_context else {
+ return ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::Done,
+ ));
+ };
+ if table.is_empty() {
+ return ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::Done,
+ ));
+ }
+ return ControlFlow::Continue(
+ OrderedFinalAggregateState::Spilling {
+ table,
+ spill_context,
+ },
+ );
+ }
+ Err(e) => {
+ return ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::Done,
+ ));
+ }
+ }
+
+ let result = if spill_context
+ .as_ref()
+ .is_some_and(|spill_context| spill_context.has_spills())
+ {
+ // Once one incomplete run is spilled, every remaining state
+ // must participate in replay so no group is finalized twice.
+ Ok(None)
+ } else {
+ let timer = elapsed_compute.timer();
+ let result = table.next_output_batch();
+ timer.done();
+ result
+ };
match result {
// Some finalized groups can be emitted. Yield them, then
// continue aggregating input in the current state.
Ok(Some(batch)) => {
- let next_state =
- OrderedFinalAggregateState::ReadingInput { table };
- self.resize_reservation_for_state(&next_state);
+ if let Err(e) =
+ self.reservation
+ .try_resize(Self::reservation_size_for_table(
+ &table,
+ spill_context.as_deref(),
+ ))
+ {
+ return ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::Done,
+ ));
+ }
+ let next_state = OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context,
+ };
ControlFlow::Break((
Poll::Ready(Some(Ok(
@@ -180,36 +508,193 @@ 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 {
+ return Self::break_with_internal_err(
+ "Ordered final aggregate stream expected Spilling state",
+ );
+ };
+
+ // Sanity check: it's impossible to OOM when the table is empty
+ if table.is_empty() {
+ return ControlFlow::Break((
+ Poll::Ready(Some(internal_err!(
+ "Ordered final aggregation entered Spilling with an empty table"
+ ))),
+ OrderedFinalAggregateState::Done,
+ ));
+ }
+
+ 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 let Err(e) = self.reservation.try_resize(table.memory_size()) {
+ result =
+ Err(e.context("Decreasing allocation after spilling should succeed"));
+ }
+
+ timer.done();
+
+ match result {
+ // Finished spilling the aggregate table, continue aggregating from input
+ Ok(()) => ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context: Some(spill_context),
+ }),
+ 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 {
+ return Self::break_with_internal_err(
+ "Ordered final aggregate stream expected PreparingMergeInput 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 {
+ return Self::break_with_internal_err(
+ "Ordered final aggregate stream expected MergingSpills state",
+ );
+ };
+
+ match stream.poll_next_unpin(cx) {
+ Poll::Pending => ControlFlow::Break((
+ Poll::Pending,
+ OrderedFinalAggregateState::MergingSpills { stream },
+ )),
+ Poll::Ready(Some(Ok(batch))) => ControlFlow::Break((
+ Poll::Ready(Some(Ok(batch))),
+ OrderedFinalAggregateState::MergingSpills { stream },
+ )),
+ Poll::Ready(Some(Err(e))) => ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::Done,
+ )),
+ Poll::Ready(None) => ControlFlow::Continue(OrderedFinalAggregateState::Done),
}
}
@@ -221,12 +706,14 @@ impl OrderedFinalAggregateStream {
/// See comments at `poll_next()` for details.
///
/// Returns the next operator state with control flow decision.
- fn handle_draining_final(
+ fn handle_producing_output(
&mut self,
original_state: OrderedFinalAggregateState,
) -> OrderedFinalAggregateStateTransition {
- let OrderedFinalAggregateState::DrainingFinal { table } = original_state else {
- unreachable!("expected draining final state")
+ let OrderedFinalAggregateState::ProducingOutput { table } = original_state else {
+ return Self::break_with_internal_err(
+ "Ordered final aggregate stream expected ProducingOutput state",
+ );
};
let mut table = table;
@@ -238,11 +725,23 @@ impl OrderedFinalAggregateStream {
match result {
Ok(Some(batch)) => {
let next_state = if table.is_empty() {
+ drop(table);
+ if let Err(e) = self.reservation.try_resize(0) {
+ return ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::Done,
+ ));
+ }
OrderedFinalAggregateState::Done
} else {
- OrderedFinalAggregateState::DrainingFinal { table }
+ if let Err(e) = self.reservation.try_resize(table.memory_size()) {
+ return ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::ProducingOutput { table },
+ ));
+ }
+ OrderedFinalAggregateState::ProducingOutput { table }
};
- self.resize_reservation_for_state(&next_state);
ControlFlow::Break((
Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))),
@@ -251,24 +750,18 @@ impl OrderedFinalAggregateStream {
}
Err(e) => ControlFlow::Break((
Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::DrainingFinal { table },
+ OrderedFinalAggregateState::ProducingOutput { table },
)),
Ok(None) => {
+ drop(table);
let next_state = OrderedFinalAggregateState::Done;
- self.resize_reservation_for_state(&next_state);
+ if let Err(e) = self.reservation.try_resize(0) {
+ return ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state));
+ }
ControlFlow::Continue(next_state)
}
}
}
-
- fn resize_reservation_for_state(&mut self, state: &OrderedFinalAggregateState) {
- let new_size = match state {
- OrderedFinalAggregateState::ReadingInput { table }
- | OrderedFinalAggregateState::DrainingFinal { table } => table.memory_size(),
- OrderedFinalAggregateState::Done => 0,
- };
- let _ = self.reservation.try_resize(new_size);
- }
}
impl Stream for OrderedFinalAggregateStream {
@@ -288,15 +781,37 @@ impl Stream for OrderedFinalAggregateStream {
///
/// ReadingInput
/// -> ReadingInput
- /// Merge one input batch. If the ordering proves some groups are
- /// complete, yield one final aggregate batch immediately, then continue
- /// reading input. Otherwise continue directly with the next input batch.
- /// -> DrainingFinal
- /// Input was exhausted. Mark the table input as done so every remaining
- /// group is safe to emit.
+ /// Merge one input batch. If it fits in memory, optionally yield groups
+ /// proven complete by the input ordering, then read the next batch.
+ /// -> Spilling
+ /// The table cannot reserve enough memory. Move all current states into
+ /// one fully group-key-sorted spill run.
+ /// -> ProducingOutput
+ /// Input was exhausted without spilling. Mark every remaining group as
+ /// complete and produce its final result.
+ /// -> PreparingMergeInput
+ /// Input was exhausted after spilling. Spill the last in-memory run and
+ /// construct the ordered input used to merge all spill files.
+ ///
+ /// Spilling
+ /// -> ReadingInput
+ /// One sorted run was written; resume reading the original input.
+ ///
+ /// PreparingMergeInput
+ /// Spill the final in-memory run and build the input ordered replay stream.
+ /// -> MergingSpills
+ /// The final run was spilled and the ordered replay stream was built.
///
- /// DrainingFinal
- /// -> DrainingFinal
+ /// MergingSpills
+ /// Aggregate the merged spill runs and emit final results.
+ /// -> MergingSpills
+ /// Forward one result batch from the fully ordered replay stream that
+ /// consumes the sort-preserving merge.
+ /// -> Done
+ /// The merged spill input was fully aggregated.
+ ///
+ /// ProducingOutput
+ /// -> ProducingOutput
/// One remaining final aggregate batch was yielded; repeat to continue
/// draining the table.
/// -> Done
@@ -319,8 +834,17 @@ impl Stream for OrderedFinalAggregateStream {
state @ OrderedFinalAggregateState::ReadingInput { .. } => {
self.handle_reading_input(cx, state)
}
- state @ OrderedFinalAggregateState::DrainingFinal { .. } => {
- self.handle_draining_final(state)
+ state @ OrderedFinalAggregateState::Spilling { .. } => {
+ self.handle_spilling(state)
+ }
+ state @ OrderedFinalAggregateState::PreparingMergeInput { .. } => {
+ self.handle_preparing_merge_input(state)
+ }
+ state @ OrderedFinalAggregateState::MergingSpills { .. } => {
+ self.handle_merging_spills(cx, state)
+ }
+ state @ OrderedFinalAggregateState::ProducingOutput { .. } => {
+ self.handle_producing_output(state)
}
state @ OrderedFinalAggregateState::Done => {
let _ = self.reservation.try_resize(0);
@@ -334,6 +858,15 @@ impl Stream for OrderedFinalAggregateStream {
self.state = Some(next_state);
continue;
}
+ ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => {
+ // Errors are terminal: discard all operator state and release
+ // its upstream input and memory reservation before returning.
+ drop(next_state);
+ self.close_input();
+ self.reservation.free();
+ self.state = Some(OrderedFinalAggregateState::Done);
+ return Poll::Ready(Some(Err(e)));
+ }
ControlFlow::Break((poll, next_state)) => {
self.state = Some(next_state);
return poll;
diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs
index b4b7fa073a..73d15a8278 100644
--- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs
@@ -23,7 +23,7 @@ use std::task::{Context, Poll};
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
-use datafusion_common::Result;
+use datafusion_common::{DataFusionError, Result};
use datafusion_execution::TaskContext;
use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
use futures::stream::{Stream, StreamExt};
@@ -31,6 +31,7 @@ use futures::stream::{Stream, StreamExt};
use super::AggregateExec;
use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker};
use crate::aggregates::AggregateMode;
+use crate::aggregates::order::GroupOrdering;
use crate::metrics::{BaselineMetrics, MetricBuilder, RecordOutput, SpillMetrics};
use crate::stream::EmptyRecordBatchStream;
use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics};
@@ -69,6 +70,41 @@ 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
@@ -76,7 +112,6 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metric
///
/// See issue for details: <https://github.com/apache/datafusion/issues/22710>
///
-/// More applicable optimizations are left to future work.
pub(crate) struct OrderedPartialAggregateStream {
schema: SchemaRef,
input: SendableRecordBatchStream,
@@ -131,6 +166,10 @@ impl OrderedPartialAggregateStream {
)?;
let reservation =
MemoryConsumer::new(format!("OrderedPartialAggregateStream[{partition}]"))
+ .with_can_spill(matches!(
+ table.group_ordering(),
+ GroupOrdering::Partial(_)
+ ))
.register(context.memory_pool());
Ok(Self {
@@ -185,6 +224,28 @@ impl OrderedPartialAggregateStream {
));
}
+ // Check memory reservation. See function comments for details.
+ match self.resize_or_take_state_batch(&mut table) {
+ Ok(Some(batch)) => {
+ self.reduction_factor.add_part(batch.num_rows());
+ return ControlFlow::Break((
+ Poll::Ready(Some(Ok(
+ batch.record_output(&self.baseline_metrics)
+ ))),
+ OrderedPartialAggregateState::ReadingInput { table },
+ ));
+ }
+ Ok(None) => {}
+ Err(e) => {
+ self.close_input();
+ self.reservation.free();
+ return ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedPartialAggregateState::Done,
+ ));
+ }
+ }
+
let timer = elapsed_compute.timer();
let result = table.next_output_batch();
timer.done();
@@ -195,9 +256,16 @@ impl OrderedPartialAggregateStream {
// current state)
Ok(Some(batch)) => {
self.reduction_factor.add_part(batch.num_rows());
+ if let Err(e) = self.reservation.try_resize(table.memory_size()) {
+ self.close_input();
+ self.reservation.free();
+ return ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedPartialAggregateState::Done,
+ ));
+ }
let next_state =
OrderedPartialAggregateState::ReadingInput { table };
- self.resize_reservation_for_state(&next_state);
ControlFlow::Break((
Poll::Ready(Some(Ok(
@@ -206,21 +274,10 @@ impl OrderedPartialAggregateStream {
next_state,
))
}
- Ok(None) => {
- // Ordered variant don't support memory-limited execution,
- // it have to error when OOM
- if let Err(e) = self.reservation.try_resize(table.memory_size()) {
- return ControlFlow::Break((
- Poll::Ready(Some(Err(e))),
- OrderedPartialAggregateState::ReadingInput { table },
- ));
- }
-
- // Can't do early emit, continue aggregating.
- ControlFlow::Continue(
- OrderedPartialAggregateState::ReadingInput { table },
- )
- }
+ // Can't do early emit, continue aggregating.
+ Ok(None) => ControlFlow::Continue(
+ OrderedPartialAggregateState::ReadingInput { table },
+ ),
Err(e) => ControlFlow::Break((
Poll::Ready(Some(Err(e))),
OrderedPartialAggregateState::ReadingInput { table },
@@ -242,6 +299,42 @@ impl OrderedPartialAggregateStream {
}
}
+ /// Update the memory reservation, and:
+ /// - If memory reservation succeed, returns `Ok(None)`
+ /// - If memory reservation failed,
+ /// - If input is partially ordered, materialize all the output, and
+ /// directly send them to the final aggregation stage.
+ /// Returns `Ok(Some(batch))`
+ /// - If input is fully ordered, directly return error. It's not
+ /// expected to use more than constant memory.
+ /// Returns `Err(..)`
+ ///
+ /// # Implementation Note
+ /// Incrementally output it after the blocked state management is ready, keep
+ /// it simple for now.
+ ///
+ /// Issue: <https://github.com/apache/datafusion/issues/7065>
+ fn resize_or_take_state_batch(
+ &mut self,
+ table: &mut OrderedAggregateTable<PartialMarker>,
+ ) -> Result<Option<RecordBatch>> {
+ let oom = match self.reservation.try_resize(table.memory_size()) {
+ Ok(()) => return Ok(None),
+ Err(e @ DataFusionError::ResourcesExhausted(_)) => e,
+ Err(e) => return Err(e),
+ };
+
+ if matches!(table.group_ordering(), GroupOrdering::Full(_)) {
+ return Err(oom);
+ }
+
+ let Some(batch) = table.take_state_batch()? else {
+ return Err(oom);
+ };
+ self.reservation.try_resize(table.memory_size())?;
+ Ok(Some(batch))
+ }
+
/// Emits one batch after input is exhausted.
///
/// `table.input_done()` has already made every remaining group safe to emit,
@@ -272,7 +365,9 @@ impl OrderedPartialAggregateStream {
} else {
OrderedPartialAggregateState::DrainingFinal { table }
};
- self.resize_reservation_for_state(&next_state);
+ if let Err(e) = self.resize_reservation_for_state(&next_state) {
+ return ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state));
+ }
ControlFlow::Break((
Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))),
@@ -285,13 +380,18 @@ impl OrderedPartialAggregateStream {
)),
Ok(None) => {
let next_state = OrderedPartialAggregateState::Done;
- self.resize_reservation_for_state(&next_state);
+ if let Err(e) = self.resize_reservation_for_state(&next_state) {
+ return ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state));
+ }
ControlFlow::Continue(next_state)
}
}
}
- fn resize_reservation_for_state(&mut self, state: &OrderedPartialAggregateState) {
+ fn resize_reservation_for_state(
+ &mut self,
+ state: &OrderedPartialAggregateState,
+ ) -> Result<()> {
let new_size = match state {
OrderedPartialAggregateState::ReadingInput { table }
| OrderedPartialAggregateState::DrainingFinal { table } => {
@@ -299,7 +399,7 @@ impl OrderedPartialAggregateStream {
}
OrderedPartialAggregateState::Done => 0,
};
- let _ = self.reservation.try_resize(new_size);
+ self.reservation.try_resize(new_size)
}
}
diff --git a/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt b/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt
new file mode 100644
index 0000000000..a4d492ab82
--- /dev/null
+++ b/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
+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 3: The same query spills six times with a 500 KB limit.
+statement ok
+SET datafusion.runtime.memory_limit = '500K'
+
+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
+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
+
+# Exercise the same spill path with a variable-width string aggregate state in
+# the spilled payload. Keep one partial input partition so memory pressure is on
+# the ordered aggregate rather than a repartition merge.
+statement ok
+SET datafusion.runtime.memory_limit = '600K'
+
+# Ensures final aggregate has spill_count > 0
+query TT
+EXPLAIN ANALYZE
+SELECT round(v1, -4), v1 % 5000,
+ sum(v1 * 2), min(CAST(v1 % 2 AS VARCHAR))
+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)), min(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>aggr=[sum(t1.v1 * Int64(2)), min(t1.v1 % Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[<slt:ignore>spill_count=0,<slt:ignore>]
+<slt:ignore>
+
+statement ok
+RESET datafusion.runtime.memory_limit
+
+statement ok
+RESET datafusion.optimizer.enable_round_robin_repartition
+
+statement ok
+RESET datafusion.execution.enable_migration_aggregate
+
+statement ok
+RESET datafusion.optimizer.prefer_existing_sort
+
+statement ok
+RESET datafusion.optimizer.repartition_aggregations
+
+statement ok
+RESET datafusion.execution.batch_size
+
+statement ok
+SET datafusion.execution.target_partitions = 4
+
+statement ok
+RESET datafusion.catalog.create_default_catalog_and_schema
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org