fix(physical-plan): honor distinct soft limits in SingleHashAggregateStream by TinyMurky · Pull Request #25158 · apache/datafusion

TinyMurky · 2026-09-10T15:55:58Z

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.

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.

…Stream ## Which issue does this PR close? - Closes [apache#24980](apache#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.

codecov-commenter · 2026-09-10T16:36:11Z

Codecov Report

❌ Patch coverage is 90.12346% with 16 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.90%. Comparing base (5747e87) to head (f7a7acb).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/aggregates/mod.rs	88.18%	6 Missing and 9 partials ⚠️
...sion/physical-plan/src/aggregates/single_stream.rs	97.14%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25158      +/-   ##
==========================================
- Coverage   81.90%   81.90%   -0.01%     
==========================================
  Files        1132     1132              
  Lines      420155   420296     +141     
  Branches   420155   420296     +141     
==========================================
+ Hits       344126   344239     +113     
- Misses      55742    55756      +14     
- Partials    20287    20301      +14

☔ 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 · 2026-09-11T03:05:17Z

                    return Self::break_with_err(e);
                }

+                // Soft limit optimization:


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

2010YOUY01 · 2026-09-11T03:08:00Z

+                // 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) {


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;

}

2010YOUY01 · 2026-09-11T03:09:23Z

                    return Self::break_with_err(e);
                }

+                // Soft limit optimization:


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

2010YOUY01 · 2026-09-11T03:15:41Z

        Ok(())
    }

+    #[tokio::test]


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)

2010YOUY01 · 2026-09-11T03:17:16Z

+    /// When set, there are no aggregate expressions: AggregateExec routes
+    /// limited non-DISTINCT aggregates to a different stream.


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

TinyMurky · 2026-09-11T03:44:18Z

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

github-actions Bot added the physical-plan Changes to the physical-plan crate label Sep 10, 2026

TinyMurky mentioned this pull request Sep 10, 2026

SingleHashAggregateStream ignores the distinct soft limit (lim=[n]) and reads all input #24980

Open

2010YOUY01 reviewed Sep 11, 2026

View reviewed changes

Navigation Menu

Uh oh!

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

fix(physical-plan): honor distinct soft limits in SingleHashAggregateStream#25158
TinyMurky wants to merge 1 commit into
apache:mainfrom
TinyMurky:SingleHashAggregateStream-add-soft-limit

TinyMurky commented Sep 10, 2026 •

edited

Loading

Uh oh!

codecov-commenter commented Sep 10, 2026

Uh oh!

2010YOUY01 left a comment

Uh oh!

2010YOUY01 Sep 11, 2026

Uh oh!

2010YOUY01 Sep 11, 2026

Uh oh!

2010YOUY01 Sep 11, 2026

Uh oh!

2010YOUY01 Sep 11, 2026

Uh oh!

2010YOUY01 Sep 11, 2026

Uh oh!

TinyMurky commented Sep 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

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

		/// When set, there are no aggregate expressions: AggregateExec routes
		/// limited non-DISTINCT aggregates to a different stream.

Uh oh!

Conversation

TinyMurky commented Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

codecov-commenter commented Sep 10, 2026

Codecov Report

Uh oh!

2010YOUY01 left a comment

Choose a reason for hiding this comment

Uh oh!

2010YOUY01 Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

2010YOUY01 Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

2010YOUY01 Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

2010YOUY01 Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

2010YOUY01 Sep 11, 2026

Choose a reason for hiding this comment

Uh oh!

TinyMurky commented Sep 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

TinyMurky commented Sep 10, 2026 •

edited

Loading