2010YOUY01 commented on code in PR #24259: URL: https://github.com/apache/datafusion/pull/24259#discussion_r3757233396
########## 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: This step is mostly similar to the spilling paths of other aggregation modes. We should extract the common logic later. -- 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]
