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

[TinyMurky]

Which issue does this PR close?

• Closes SingleHashAggregateStream ignores the distinct soft limit (lim=[n]) and reads all input #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.

If spilling has already occurred, the optimization is skipped to keep the spill and replay path simple. In that case, the stream consumes the remaining input and follows the existing spill merge path.

What changes are included in this PR?

• Add a distinct soft-limit check to SingleHashAggregateStream after processing each input batch.

• Stop consuming input when the soft limit is reached
  before spilling.

• Reuse the input-exhausted transition when the soft limit is reached, preserving spill merging and output preparation.
• Document how the optimization applies to single-stage
  distinct aggregation.

• Update the AggregateExec::limit_options documentation to list SingleHash as supporting distinct soft limits.
• Add an end-to-end SQL test that verifies:
  • the query is planned with a single-stage
    AggregateExec;
  • the distinct soft limit is pushed into the
    aggregate;
  • the aggregate stops consuming its input after
    reaching the soft limit;
  • the query returns the correct number of rows.

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.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 97.36842% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 81.93%. Comparing base (f8cc678) to head (e92aee7).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...sion/physical-plan/src/aggregates/single_stream.rs	97.36%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25158      +/-   ##
==========================================
- Coverage   81.93%   81.93%   -0.01%     
==========================================
  Files        1133     1133              
  Lines      423529   423546      +17     
  Branches   423529   423546      +17     
==========================================
+ Hits       347028   347038      +10     
- Misses      55910    55916       +6     
- Partials    20591    20592       +1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[2010YOUY01]

Thank you for working on this. I've left several suggestions.

[2010YOUY01]

Can we move this comment to SingleHashAggregateStream, and here we can comment 'see comments at xxx for details'

Additionally we can follow the comment pattern in (first explain how the SQL get optimized to soft limit, and next the internal early termination mechanism)

• datafusion/datafusion/physical-plan/src/aggregates/hash_stream.rs

  Line 77 in 7e5f40a

  /// # Optimization: DISTINCT LIMIT Soft Limit

[TinyMurky]

I added documentation to SingleHashAggregateStream explaining how the soft-limit optimization is planned and how early termination works:

datafusion/datafusion/physical-plan/src/aggregates/single_stream.rs

Lines 92 to 118 in e92aee7

	/// # 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.

I also changed the field documentation to:

datafusion/datafusion/physical-plan/src/aggregates/single_stream.rs

Lines 136 to 138 in e92aee7

	/// See the "Optimization: DISTINCT LIMIT Soft Limit" section in
	/// [`SingleHashAggregateStream`] for details.
	group_values_soft_limit: Option<usize>,

to

    /// See the "Optimization: DISTINCT LIMIT Soft Limit" section in
    /// [`SingleHashAggregateStream`] for details.

[2010YOUY01]

I suggest to skip this optimization if we have spilled before

Here is the pattern to follow, and also the explanaiton

datafusion/datafusion/physical-plan/src/aggregates/hash_stream.rs

Lines 857 to 865 in 7e5f40a

	// 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());
	if self.hit_soft_group_limit(hash_table) && !spilled {
	break;
	}

[TinyMurky]

I added the skip of spilling.

datafusion/datafusion/physical-plan/src/aggregates/single_stream.rs

Lines 485 to 498 in e92aee7

	// 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);
	}

[2010YOUY01]

and we can update control flow comment at poll_next to briefly mention this optimization change

[TinyMurky]

I changed the comment of poll_next as following:

datafusion/datafusion/physical-plan/src/aggregates/single_stream.rs

Lines 793 to 805 in e92aee7

	/// ReadingInput
	/// -> ReadingInput
	/// Aggregate one raw input batch. If it fits in memory, continue with
	/// the next input batch.
	/// -> Spilling
	/// The table cannot reserve enough memory. Move all current states into
	/// one fully group-key-sorted spill run.
	/// -> ProducingOutput
	/// 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.

[2010YOUY01]

I recommend to write this test differently (follow the pattern in 40a6454#diff-02af0439a3df656429990b220b80e50d8df259ce45c47e008460c1ca3781aca3)

The main difference is

• Try to exercise this feature end-to-end, from select distinct query, and get it optimized to aggregate with soft limit
• Also assert the internal metric of AggregateExec, otherwise we can't ensure if this soft limit optimization is applied -- limit can also be enforced by the downstream LimitExec operator.

(I think only such e2e test is enough, we don't have to test it individually on AggregateExec, since this optimization is only useful from such SQL patterns, and should not be directly used on the AggregateExec)

[TinyMurky]

I removed the test in single_stream.rs, and add an e2e test in datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs

datafusion/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs

Lines 222 to 339 in e92aee7

	// 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<usize>,
	input_rows: usize,
	output_rows: usize,
	}
	fn collect_aggregate_runtime_metrics(
	plan: &Arc<dyn ExecutionPlan>,
	metrics: &mut Vec<AggregateRuntimeMetric>,
	) {
	if let Some(agg) = plan.downcast_ref::<AggregateExec>() {
	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<dyn ExecutionPlan>,
	) -> Vec<AggregateRuntimeMetric> {
	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::<usize>(),
	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(())
	}

[2010YOUY01]

Maybe we can also change it to 'see top comments for details'.

[TinyMurky]

I changed the comment of group_values_soft_limit to

datafusion/datafusion/physical-plan/src/aggregates/single_stream.rs

Lines 136 to 138 in e92aee7

	/// See the "Optimization: DISTINCT LIMIT Soft Limit" section in
	/// [`SingleHashAggregateStream`] for details.
	group_values_soft_limit: Option<usize>,

[TinyMurky]

Thanks for the review!
I will start to work on them!

[TinyMurky]

@2010YOUY01 Hi! Sorry for the delay. I've pushed a new commit addressing your review comments. When you have a chance, could you please take another look? Thanks!!