diff --git a/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs b/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs index 323dfd618330..9a43ba01da44 100644 --- a/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs +++ b/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs @@ -104,8 +104,8 @@ async fn test_partial_final() -> Result<()> { Ok(()) } -// Ensure operator respect the soft limit and stops early: `AggregateExec`'s -// `output_rows` metric should be smaller than then total distinct group count. +// Ensure operator respects the soft limit and stops early: `AggregateExec`'s +// `output_rows` metric should be smaller than the total distinct group count. #[tokio::test] async fn limited_distinct_aggregate_stream_respects_soft_limit() -> Result<()> { // Snapshot for an aggregate operator node from `EXPLAIN ANALYZE`. @@ -219,6 +219,125 @@ async fn limited_distinct_aggregate_stream_respects_soft_limit() -> Result<()> { Ok(()) } +// Ensure operator respects the soft limit and stops early: `AggregateExec`'s +// `output_rows` metric should be smaller than the total distinct group count. +#[tokio::test] +async fn single_distinct_aggregate_stream_respects_soft_limit() -> Result<()> { + // Snapshot for an aggregate operator node from `EXPLAIN ANALYZE`. + // + // Example: In an `EXPLAIN ANALYZE` output + // ```txt + // AggregateExec: mode=single, aggr=[], lim=[10], metrics=[output_rows=10, ...] + // ProjectionExec: metrics=[output_rows=10, ...] + // ``` + // + // `output_rows` comes from the `AggregateExec` itself, while `input_rows` + // is the `output_rows` metric of its direct input operator (such as a `ProjectionExec`). + // Tracking both distinguishes early input termination from the downstream `LimitExec` + // merely stopping after it receives enough output rows. + // + // we get: + // ```txt + // AggregateRuntimeMetric { + // mode: Single, + // limit: Some(10), + // input_rows: 10, + // output_rows: 10, + // } + // ``` + #[derive(Debug)] + struct AggregateRuntimeMetric { + mode: AggregateMode, + limit: Option, + input_rows: usize, + output_rows: usize, + } + + fn collect_aggregate_runtime_metrics( + plan: &Arc, + metrics: &mut Vec, + ) { + if let Some(agg) = plan.downcast_ref::() { + let input_rows = agg + .input() + .metrics() + .and_then(|metrics| metrics.aggregate_by_name().output_rows()) + .expect("The input Exec should record output_rows after execution"); + + let output_rows = agg + .metrics() + .and_then(|metrics| metrics.aggregate_by_name().output_rows()) + .expect("AggregateExec should record output_rows after execution"); + + metrics.push(AggregateRuntimeMetric { + mode: *agg.mode(), + limit: agg.limit_options().map(|config| config.limit()), + input_rows, + output_rows, + }); + } + + for child in plan.children() { + collect_aggregate_runtime_metrics(child, metrics); + } + } + + fn aggregate_runtime_metrics( + plan: &Arc, + ) -> Vec { + let mut metrics = vec![]; + collect_aggregate_runtime_metrics(plan, &mut metrics); + metrics + } + + let cfg = SessionConfig::new() + .with_target_partitions(1) + .with_batch_size(10) + .set_bool("datafusion.execution.enable_migration_aggregate", true); + + let ctx = SessionContext::new_with_config(cfg); + + let dataframe = ctx + .sql( + "SELECT DISTINCT value % 100000 AS v \ + FROM generate_series(1000000) \ + LIMIT 10", + ) + .await?; + let plan = dataframe.create_physical_plan().await?; + let formatted_plan = displayable(plan.as_ref()).indent(false).to_string(); + assert!( + formatted_plan.contains("AggregateExec: mode=Single"), + "expected a single aggregate in plan:\n{formatted_plan}" + ); + + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await?; + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 10 + ); + + let metrics = aggregate_runtime_metrics(&plan); + let single = metrics + .iter() + .find(|metric| metric.mode == AggregateMode::Single) + .expect("expected single aggregate metrics"); + + assert_eq!(single.limit, Some(10)); + + assert!( + single.input_rows <= 10, + "single aggregate should stop reading input after reaching the soft limit: {metrics:?}" + ); + + assert!( + single.output_rows <= 10, + "single aggregate should stop before emitting all distinct groups: {metrics:?}" + ); + + Ok(()) +} + #[tokio::test] async fn test_single_local() -> Result<()> { let source = mock_data()?; diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 3ed93e09ce4f..854fbf8484eb 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -868,8 +868,8 @@ pub struct AggregateExec { /// Supported by: /// - [`StreamType::GroupedPriorityQueue`]: retains only the best `limit` /// groups per partition (this stream is selected only when a limit is set) - /// - [`StreamType::PartialHash`], [`StreamType::FinalHash`] and the legacy - /// [`StreamType::GroupedHash`]: stop reading input once `limit` groups + /// - [`StreamType::SingleHash`], [`StreamType::PartialHash`], [`StreamType::FinalHash`] + /// and the legacy [`StreamType::GroupedHash`]: stop reading input once `limit` groups /// have been accumulated /// /// The remaining streams consume all input. diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 6d46b90c2677..cfa708c51129 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -88,6 +88,34 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// into an ordered streaming aggregation, which ensures bounded memory usage and /// evaluates the final result. /// - [`OrderedFinalAggregateStream`] is reused for the streaming aggregation. +/// +/// # Optimization: DISTINCT LIMIT Soft Limit +/// +/// When the input has only one partition or the input is already partitioned, +/// unordered distinct queries such as: +/// +/// ```sql +/// SELECT DISTINCT x FROM t LIMIT 10; +/// ``` +/// +/// are optimized into a single-stage aggregate like: +/// +/// ```txt +/// LimitExec, limit=10 +/// --AggregateExec(Single), group_by=[x], aggr=[], soft_limit=10 +/// ---- Scan(t) +/// ``` +/// +/// After each input batch, the stream checks whether the soft limit has been +/// reached. If so, it emits the accumulated groups and stops reading input. +/// +/// This early termination is skipped after spilling has occurred to keep the +/// spill and replay path simple. In that case, the stream consumes the remaining +/// input and merges all spill runs before producing output. +/// +/// This operator does not guarantee an exact limit because a single batch can +/// cross the threshold. The downstream limit operator enforces the exact result +/// size. pub(crate) struct SingleHashAggregateStream { /// Output schema: group columns followed by final aggregate value columns. schema: SchemaRef, @@ -104,6 +132,10 @@ pub(crate) struct SingleHashAggregateStream { /// Tracks the high-level stream lifecycle. The hash table owns the lower-level /// state for emitting output batches. state: Option, + + /// See the "Optimization: DISTINCT LIMIT Soft Limit" section in + /// [`SingleHashAggregateStream`] for details. + group_values_soft_limit: Option, } /// Spill configuration and accumulated runs for single hash aggregation. @@ -374,6 +406,7 @@ impl SingleHashAggregateStream { hash_table, spill_context, }), + group_values_soft_limit: agg.limit_options().map(|config| config.limit()), }) } @@ -449,6 +482,20 @@ impl SingleHashAggregateStream { return Self::break_with_err(e); } + // Soft group limits are usually small and rarely coincide with + // spilling. Once spilling has occurred, skip this optimization to + // make the internal logic simpler. + let spilled = spill_context + .as_ref() + .is_some_and(|context| context.has_spills()); + + // See the "Optimization: DISTINCT LIMIT Soft Limit" section in + // `SingleHashAggregateStream` for details. + if self.hit_soft_group_limit(&hash_table) && !spilled { + return self + .close_input_and_prepare_output(hash_table, spill_context); + } + // Check memory reservation, and potentially spill. let timer = elapsed_compute.timer(); let resize_result = @@ -490,30 +537,53 @@ impl SingleHashAggregateStream { } Poll::Ready(Some(Err(e))) => Self::break_with_err(e), Poll::Ready(None) => { - self.close_input(); - match spill_context { - Some(spill_context) if spill_context.has_spills() => { - ControlFlow::Continue( - SingleHashAggregateState::PreparingMergeInput { - hash_table, - spill_context, - }, - ) - } - _ => { - let elapsed_compute = - self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = hash_table.start_output(); - timer.done(); - - match result { - Ok(()) => ControlFlow::Continue( - SingleHashAggregateState::ProducingOutput { hash_table }, - ), - Err(e) => Self::break_with_err(e), - } + self.close_input_and_prepare_output(hash_table, spill_context) + } + } + } + + /// See comments in [`Self::group_values_soft_limit`] for details. + fn hit_soft_group_limit( + &self, + hash_table: &AggregateHashTable, + ) -> bool { + self.group_values_soft_limit + .is_some_and(|limit| limit <= hash_table.building_group_count()) + } + + /// Stops consuming input and prepares the next execution phase. + /// Called when the input is exhausted or the distinct soft limit is reached. + /// + /// If data has been spilled, transitions to `PreparingMergeInput` so the + /// spilled and in-memory groups can be merged before output. Otherwise, + /// starts output from the in-memory hash table and transitions to + /// `ProducingOutput`. + fn close_input_and_prepare_output( + &mut self, + mut hash_table: AggregateHashTable, + spill_context: Option>, + ) -> SingleHashAggregateStateTransition { + self.close_input(); + match spill_context { + Some(spill_context) if spill_context.has_spills() => { + ControlFlow::Continue(SingleHashAggregateState::PreparingMergeInput { + hash_table, + spill_context, + }) + } + _ => { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = hash_table.start_output(); + timer.done(); + + match result { + Ok(()) => { + ControlFlow::Continue(SingleHashAggregateState::ProducingOutput { + hash_table, + }) } + Err(e) => Self::break_with_err(e), } } } @@ -728,7 +798,8 @@ impl Stream for SingleHashAggregateStream { /// The table cannot reserve enough memory. Move all current states into /// one fully group-key-sorted spill run. /// -> ProducingOutput - /// Input was exhausted without spilling. Start outputting final values. + /// Input was exhausted without spilling, or the distinct soft limit was + /// reached before spilling. Start outputting final values. /// -> PreparingMergeInput /// Input was exhausted after spilling. Spill the last in-memory run and /// construct the ordered input used to merge all spill files.