From 8b4b772ef6b3046616c362e5684a6572600608d7 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Wed, 9 Sep 2026 01:26:41 +0800 Subject: [PATCH 1/2] fix(physical-plan): honor distinct soft limits in SingleHashAggregateStream ## Which issue does this PR close? - Closes [#24980](https://github.com/apache/datafusion/issues/24980) ## Rationale for this change `SingleHashAggregateStream` ignores the distinct soft limit pushed into `AggregateExec`. As a result, single-stage `SELECT DISTINCT ... LIMIT n` queries consume all input even after enough distinct groups have been collected. This change stops input consumption once the in-memory hash table contains at least the requested number of distinct groups. Any existing spills are merged before producing output, and the downstream limit operator enforces the exact output row count. ## What changes are included in this PR? * Add a distinct soft-limit check to `SingleHashAggregateStream` after processing each input batch. * Reuse the input-exhausted transition when the soft limit is reached, preserving spill merging and output preparation. * Add unit tests covering early termination with and without spilling. - Update the `AggregateExec::limit_options` documentation to list `SingleHash` as supporting distinct soft limits. ## Are these changes tested? Added unit tests covering: * Reaching the soft limit without spilling. * Reaching the soft limit after spilling, preserving spilled groups and deduplicating overlapping groups. * Rejecting further input consumption after the soft limit is reached in the spill case. Following test commands have been executed and passed - `cargo test --profile=ci --test sqllogictests` - `cargo test -p datafusion` - `cargo test -p datafusion-cli` ## Are there any user-facing changes? Single-stage `DISTINCT` queries with a limit can stop consuming input earlier, reducing unnecessary work. Query semantics are unchanged. --- .../physical-plan/src/aggregates/mod.rs | 197 +++++++++++++++++- .../src/aggregates/single_stream.rs | 95 +++++++-- 2 files changed, 267 insertions(+), 25 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 3ed93e09ce4f4..e417b5a1bcd1f 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. @@ -4424,6 +4424,199 @@ mod tests { Ok(()) } + #[tokio::test] + async fn single_hash_distinct_aggregation_soft_limit_no_spill() -> Result<()> { + // Verify that single hash aggregation stops reading + // input once it has collected enough distinct groups to + // satisfy the soft limit. + + let soft_limit: usize = 2; + + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, false)])); + + let input_batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + // The accumulated groups should be [1] + vec![Arc::new(UInt32Array::from(vec![1, 1]))], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + // The accumulated groups should be [1, 2] + vec![Arc::new(UInt32Array::from(vec![2, 2]))], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + // This batch must not be read. + vec![Arc::new(UInt32Array::from(vec![3, 4, 5]))], + )?, + ]; + + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + + let input = TestMemoryExec::try_new_exec( + std::slice::from_ref(&input_batches), + Arc::clone(&schema), + None, + )?; + + let single_aggregate = Arc::new( + AggregateExec::try_new( + AggregateMode::Single, + group_by.clone(), + vec![], + vec![], + input, + Arc::clone(&schema), + )? + .with_limit_options(Some(LimitOptions::new(soft_limit))), + ); + + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .set_bool("datafusion.execution.enable_migration_aggregate", true), + ), + ); + + let single_stream = single_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(single_stream, StreamType::SingleHash(_))); + let stream: SendableRecordBatchStream = single_stream.into(); + let single_output = collect(stream).await?; + assert_eq!( + single_output + .iter() + .map(RecordBatch::num_rows) + .sum::(), + soft_limit + ); + assert_snapshot!(batches_to_sort_string(&single_output), @r" ++---+ +| a | ++---+ +| 1 | +| 2 | ++---+ +"); + + Ok(()) + } + + #[tokio::test] + async fn single_hash_distinct_aggregation_soft_limit_with_spill() -> Result<()> { + // Verify that reaching the soft limit after spilling stops input consumption + // and merges the spilled and in-memory groups before producing output. + + use crate::test::exec::MockExec; + + // Force the first batch to spill while leaving room for spill replay. + let pool_size = 8 * 1024; // 8 KiB + let batch_size = 2; + let memory_pool = Arc::new(FairSpillPool::new(pool_size)); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config( + SessionConfig::new().with_batch_size(batch_size).set_bool( + "datafusion.execution.enable_migration_aggregate", + true, + ), + ) + .with_runtime(Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(memory_pool) + .build()?, + )), + ); + + let soft_limit = 1024; + + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, false)])); + + let input_batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + // Fewer groups than the soft limit, but enough to force a spill. + vec![Arc::new(UInt32Array::from_iter_values(0..512))], + ) + .map_err(DataFusionError::from), + RecordBatch::try_new( + Arc::clone(&schema), + // Reach the soft limit after spilling, with overlapping groups. + vec![Arc::new(UInt32Array::from_iter_values(256..1280))], + ) + .map_err(DataFusionError::from), + // Reading past the soft limit must fail the test, even if those + // additional rows would not change the aggregate's output. + Err(DataFusionError::Execution( + "input must not be polled after reaching the soft limit".to_string(), + )), + ]; + + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + + let input = Arc::new( + MockExec::new(input_batches, Arc::clone(&schema)) + .with_use_task(false) + // Avoid inspecting the intentional input error (in `vec!`) while + // computing statistics before execution. + .with_unknown_statistics(), + ); + + let single_aggregate = Arc::new( + AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![], + vec![], + input, + Arc::clone(&schema), + )? + .with_limit_options(Some(LimitOptions::new(soft_limit))), + ); + + let stream = single_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::SingleHash(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + + let mut values: Vec = output + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect(); + values.sort_unstable(); + + //The merged output must contain each value in 0..1280 + // exactly once. + assert_eq!(values, (0..1280).collect::>()); + + let metrics = single_aggregate.metrics().expect("metrics must exist"); + + assert!( + metrics.spill_count().unwrap_or(0) > 0, + "test must actually spill" + ); + + assert!( + metrics.spilled_rows().unwrap_or(0) > 0, + "test must spill group rows", + ); + + Ok(()) + } + #[tokio::test] async fn limited_distinct_aggregate_uses_migrated_hash_streams() -> Result<()> { let schema = diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 6d46b90c2677c..5b20505e10c96 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -104,6 +104,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, + + /// When set, there are no aggregate expressions: AggregateExec routes + /// limited non-DISTINCT aggregates to a different stream. + group_values_soft_limit: Option, } /// Spill configuration and accumulated runs for single hash aggregation. @@ -374,6 +378,7 @@ impl SingleHashAggregateStream { hash_table, spill_context, }), + group_values_soft_limit: agg.limit_options().map(|config| config.limit()), }) } @@ -449,6 +454,27 @@ impl SingleHashAggregateStream { return Self::break_with_err(e); } + // Soft limit optimization: + // + // Stop reading input once the in-memory table contains enough distinct + // groups to satisfy the soft limit. + // + // When a limit is present, AggregateExec routes only unordered, + // unfiltered DISTINCT aggregates to this stream. + // + // With no aggregate expressions, additional input can only match existing + // groups or add new ones; it cannot change any existing group's output. + // Since there is no ordering requirement and we already have enough + // distinct groups, we can finish reading as if the input were exhausted. + // + // Reuse the input-exhausted transition to merge any existing spills + // before producing output. The downstream limit operator enforces + // the exact output row count. + if self.hit_soft_group_limit(&hash_table) { + 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 +516,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), } } } From e92aee75ea5bf774d0dbd6be648849a124c7797e Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 12 Sep 2026 15:56:53 +0800 Subject: [PATCH 2/2] fix(physical-plan): skip soft-limit early exit after spilling - Add an end-to-end SQL test that verifies `SingleHashAggregateStream` stops consuming input when the distinct soft limit is reached before spilling. - Add e2e test to `limited_distinct_aggregation.rs`, in order to test whether distinct aggregate will truely stop reading input after it hit soft limit. --- .../limited_distinct_aggregation.rs | 123 ++++++++++- .../physical-plan/src/aggregates/mod.rs | 193 ------------------ .../src/aggregates/single_stream.rs | 62 ++++-- 3 files changed, 163 insertions(+), 215 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs b/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs index 323dfd6183306..9a43ba01da44e 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 e417b5a1bcd1f..854fbf8484eb9 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -4424,199 +4424,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn single_hash_distinct_aggregation_soft_limit_no_spill() -> Result<()> { - // Verify that single hash aggregation stops reading - // input once it has collected enough distinct groups to - // satisfy the soft limit. - - let soft_limit: usize = 2; - - let schema = - Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, false)])); - - let input_batches = vec![ - RecordBatch::try_new( - Arc::clone(&schema), - // The accumulated groups should be [1] - vec![Arc::new(UInt32Array::from(vec![1, 1]))], - )?, - RecordBatch::try_new( - Arc::clone(&schema), - // The accumulated groups should be [1, 2] - vec![Arc::new(UInt32Array::from(vec![2, 2]))], - )?, - RecordBatch::try_new( - Arc::clone(&schema), - // This batch must not be read. - vec![Arc::new(UInt32Array::from(vec![3, 4, 5]))], - )?, - ]; - - let group_by = - PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); - - let input = TestMemoryExec::try_new_exec( - std::slice::from_ref(&input_batches), - Arc::clone(&schema), - None, - )?; - - let single_aggregate = Arc::new( - AggregateExec::try_new( - AggregateMode::Single, - group_by.clone(), - vec![], - vec![], - input, - Arc::clone(&schema), - )? - .with_limit_options(Some(LimitOptions::new(soft_limit))), - ); - - let task_ctx = Arc::new( - TaskContext::default().with_session_config( - SessionConfig::new() - .set_bool("datafusion.execution.enable_migration_aggregate", true), - ), - ); - - let single_stream = single_aggregate.execute_typed(0, &task_ctx)?; - assert!(matches!(single_stream, StreamType::SingleHash(_))); - let stream: SendableRecordBatchStream = single_stream.into(); - let single_output = collect(stream).await?; - assert_eq!( - single_output - .iter() - .map(RecordBatch::num_rows) - .sum::(), - soft_limit - ); - assert_snapshot!(batches_to_sort_string(&single_output), @r" -+---+ -| a | -+---+ -| 1 | -| 2 | -+---+ -"); - - Ok(()) - } - - #[tokio::test] - async fn single_hash_distinct_aggregation_soft_limit_with_spill() -> Result<()> { - // Verify that reaching the soft limit after spilling stops input consumption - // and merges the spilled and in-memory groups before producing output. - - use crate::test::exec::MockExec; - - // Force the first batch to spill while leaving room for spill replay. - let pool_size = 8 * 1024; // 8 KiB - let batch_size = 2; - let memory_pool = Arc::new(FairSpillPool::new(pool_size)); - let task_ctx = Arc::new( - TaskContext::default() - .with_session_config( - SessionConfig::new().with_batch_size(batch_size).set_bool( - "datafusion.execution.enable_migration_aggregate", - true, - ), - ) - .with_runtime(Arc::new( - RuntimeEnvBuilder::new() - .with_memory_pool(memory_pool) - .build()?, - )), - ); - - let soft_limit = 1024; - - let schema = - Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, false)])); - - let input_batches = vec![ - RecordBatch::try_new( - Arc::clone(&schema), - // Fewer groups than the soft limit, but enough to force a spill. - vec![Arc::new(UInt32Array::from_iter_values(0..512))], - ) - .map_err(DataFusionError::from), - RecordBatch::try_new( - Arc::clone(&schema), - // Reach the soft limit after spilling, with overlapping groups. - vec![Arc::new(UInt32Array::from_iter_values(256..1280))], - ) - .map_err(DataFusionError::from), - // Reading past the soft limit must fail the test, even if those - // additional rows would not change the aggregate's output. - Err(DataFusionError::Execution( - "input must not be polled after reaching the soft limit".to_string(), - )), - ]; - - let group_by = - PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); - - let input = Arc::new( - MockExec::new(input_batches, Arc::clone(&schema)) - .with_use_task(false) - // Avoid inspecting the intentional input error (in `vec!`) while - // computing statistics before execution. - .with_unknown_statistics(), - ); - - let single_aggregate = Arc::new( - AggregateExec::try_new( - AggregateMode::Single, - group_by, - vec![], - vec![], - input, - Arc::clone(&schema), - )? - .with_limit_options(Some(LimitOptions::new(soft_limit))), - ); - - let stream = single_aggregate.execute_typed(0, &task_ctx)?; - assert!(matches!(stream, StreamType::SingleHash(_))); - let stream: SendableRecordBatchStream = stream.into(); - let output = collect(stream).await?; - - let mut values: Vec = output - .iter() - .flat_map(|batch| { - batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .iter() - .copied() - }) - .collect(); - values.sort_unstable(); - - //The merged output must contain each value in 0..1280 - // exactly once. - assert_eq!(values, (0..1280).collect::>()); - - let metrics = single_aggregate.metrics().expect("metrics must exist"); - - assert!( - metrics.spill_count().unwrap_or(0) > 0, - "test must actually spill" - ); - - assert!( - metrics.spilled_rows().unwrap_or(0) > 0, - "test must spill group rows", - ); - - Ok(()) - } - #[tokio::test] async fn limited_distinct_aggregate_uses_migrated_hash_streams() -> Result<()> { let schema = diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 5b20505e10c96..cfa708c51129b 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, @@ -105,8 +133,8 @@ pub(crate) struct SingleHashAggregateStream { /// state for emitting output batches. state: Option, - /// When set, there are no aggregate expressions: AggregateExec routes - /// limited non-DISTINCT aggregates to a different stream. + /// See the "Optimization: DISTINCT LIMIT Soft Limit" section in + /// [`SingleHashAggregateStream`] for details. group_values_soft_limit: Option, } @@ -454,23 +482,16 @@ impl SingleHashAggregateStream { return Self::break_with_err(e); } - // Soft limit optimization: - // - // Stop reading input once the in-memory table contains enough distinct - // groups to satisfy the soft limit. - // - // When a limit is present, AggregateExec routes only unordered, - // unfiltered DISTINCT aggregates to this stream. - // - // With no aggregate expressions, additional input can only match existing - // groups or add new ones; it cannot change any existing group's output. - // Since there is no ordering requirement and we already have enough - // distinct groups, we can finish reading as if the input were exhausted. - // - // Reuse the input-exhausted transition to merge any existing spills - // before producing output. The downstream limit operator enforces - // the exact output row count. - if self.hit_soft_group_limit(&hash_table) { + // 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); } @@ -777,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.