alamb commented on code in PR #24259:
URL: https://github.com/apache/datafusion/pull/24259#discussion_r3761852773
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -1215,6 +1219,12 @@ impl AggregateExec {
)?));
}
+ if self.should_use_ordered_single_aggregate_stream(context) {
Review Comment:
FWIW while I like the explicit name of these methods I do wonder if the
difference between them might be clearer if there was one giant match statement
(rather than a bunch of if ... chains)
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -672,12 +675,15 @@ enum StreamType {
OrderedPartialAggregate(OrderedPartialAggregateStream),
/// Final stage of aggregation for ordered input.
OrderedFinalAggregate(OrderedFinalAggregateStream),
+ /// Single stage of aggregation for ordered input.
+ OrderedSingleAggregate(OrderedSingleAggregateStream),
/// Hash aggregation reused for multiple stages
///
/// Note this is being incrementally migrated to dedicated streams like
/// [`StreamType::PartialHash`], [`StreamType::FinalHash`],
/// [`StreamType::OrderedPartialAggregate`], and
- /// [`StreamType::OrderedFinalAggregate`]
+ /// [`StreamType::OrderedFinalAggregate`], and
Review Comment:
Doc nit: this list now has a double "and" ("… and `OrderedFinalAggregate`,
and `OrderedSingleAggregate`").
```suggestion
/// [`StreamType::OrderedPartialAggregate`],
/// [`StreamType::OrderedFinalAggregate`], and
```
##########
datafusion/physical-plan/src/aggregates/ordered_single_stream.rs:
##########
@@ -0,0 +1,886 @@
+// 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.
+
+//! Single-stage aggregate stream for ordered raw input.
+
+use std::ops::ControlFlow;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use datafusion_common::{DataFusionError, Result, internal_datafusion_err,
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::aggregate_hash_table::{OrderedAggregateTable, SingleMarker};
+use super::group_values::GroupByMetrics;
+use super::ordered_final_stream::OrderedFinalAggregateStream;
+use super::{AggregateExec, create_schema};
+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};
+
+/// Single aggregate stream for `InputOrderMode::Sorted` and
+/// `InputOrderMode::PartiallySorted`.
+///
+/// # Example
+///
+/// SELECT k, AVG(v) FROM t GROUP BY k;
+///
+/// If the input is ordered by `k`, and there are existing key partitioning on
group
+/// by keys, the single mode aggregation with ordering optimization can be
used:
+///
+/// ## Plan
+/// AggregateExec(stage=single, ordered)
+/// -- DataSourceExec(t)
+///
+/// ## Single Stage Behavior
+/// Input: raw rows
+/// Output: final results for all groups (for example, `AVG(x)`)
+///
+/// # Order-based Optimization
+///
+/// For the aggregation work, the hash aggregation implementation is reused.
+///
+/// After each input batch, check whether any groups can be emitted eagerly to
+/// improve memory efficiency. For example, if the last group key seen is
+/// `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.
+///
+/// On reservation failure, the stream 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.
+pub(crate) struct OrderedSingleAggregateStream {
+ schema: SchemaRef,
+ input: SendableRecordBatchStream,
+ reservation: MemoryReservation,
+ baseline_metrics: BaselineMetrics,
+ state: Option<OrderedSingleAggregateState>,
+}
+
+/// Spill configuration and accumulated runs for partially ordered single
+/// 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 OrderedSingleSpillContext {
+ /// Aggregate configuration used to construct the final replay stream.
+ final_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 OrderedSingleAggregateState {
+ ReadingInput {
+ table: OrderedAggregateTable<SingleMarker>,
+ /// None if either
+ /// - Disk Manager doesn't enable temporary file creation
+ /// - The group keys are fully ordered, it's expected to use bounded
memory
+ spill_context: Option<Box<OrderedSingleSpillContext>>,
+ },
+ Spilling {
+ table: OrderedAggregateTable<SingleMarker>,
+ spill_context: Box<OrderedSingleSpillContext>,
+ },
+ ProducingOutput {
+ table: OrderedAggregateTable<SingleMarker>,
+ },
+ PreparingMergeInput {
+ table: OrderedAggregateTable<SingleMarker>,
+ spill_context: Box<OrderedSingleSpillContext>,
+ },
+ MergingSpills {
+ stream: SendableRecordBatchStream,
+ },
+ Done,
+ /// Sentinel state to use when returning error from any other states,
because:
+ /// - It explicitly releases state-owned resources immediately
+ /// - More defensive against accidentally resuming execution after error
+ Error,
+}
+
+type OrderedSingleAggregatePoll = Poll<Option<Result<RecordBatch>>>;
+type OrderedSingleAggregateStateTransition = ControlFlow<
+ (OrderedSingleAggregatePoll, OrderedSingleAggregateState),
+ OrderedSingleAggregateState,
+>;
+
+impl OrderedSingleSpillContext {
+ 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(&agg.input().schema())?;
+ let output_ordering = agg.cache.output_ordering();
+ let InputOrderMode::PartiallySorted(order_indices) = input_order_mode
else {
+ return internal_err!(
+ "Ordered single 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 single 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());
+
+ // Spilled rows contain group keys and intermediate states. Replay must
+ // merge those states and evaluate the final aggregate values.
+ let mut final_agg = agg.clone();
+ final_agg.mode = match agg.mode {
+ AggregateMode::Single => AggregateMode::Final,
+ AggregateMode::SinglePartitioned =>
AggregateMode::FinalPartitioned,
+ mode => {
+ return internal_err!(
+ "Ordered single aggregate spill cannot replay aggregate
mode {mode:?}"
+ );
+ }
+ };
+ final_agg.group_by = Arc::new(agg.group_by.as_final());
+ final_agg.input_order_mode = InputOrderMode::Sorted;
+
+ Ok(Self {
+ final_agg,
+ 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 [`OrderedSingleAggregateStream`] for spilling details.
+ fn spill_table(
+ &mut self,
+ table: &mut OrderedAggregateTable<SingleMarker>,
+ ) -> 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,
+ "OrderedSingleAggregateSpill",
+ )?;
+
+ let Some((file, max_record_batch_memory)) = spill_file else {
+ return internal_err!("Ordered single 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(
Review Comment:
maybe worth a ticket to track
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -1215,6 +1219,12 @@ impl AggregateExec {
)?));
}
+ if self.should_use_ordered_single_aggregate_stream(context) {
Review Comment:
For example the only difference between `should_use_single_hash_stream` and
`should_use_ordered_single_aggregate_stream` seems to be the condition
```rust
&& self.input_order_mode != InputOrderMode::Linear
```
🤔
--
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]
