fix(physical-plan): honor distinct soft limits in SingleHashAggregateStream

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::<usize>(),
+            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<u32> = output
+            .iter()
+            .flat_map(|batch| {
+                batch
+                    .column(0)
+                    .as_any()
+                    .downcast_ref::<UInt32Array>()
+                    .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::<Vec<u32>>());
+
+        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 =

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<SingleHashAggregateState>,
+
+    /// When set, there are no aggregate expressions: AggregateExec routes
+    /// limited non-DISTINCT aggregates to a different stream.
+    group_values_soft_limit: Option<usize>,
 }
 
 /// 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<SingleMarker>,
+    ) -> 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<SingleMarker>,
+        spill_context: Option<Box<SingleSpillContext>>,
+    ) -> 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),
                 }
             }
         }