You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "kumarUjjawal (via GitHub)" <gi...@apache.org> on 2026/05/27 15:25:24 UTC
[PR] Fix TopK DISTINCT aggregation preserving NULLs [datafusion]
kumarUjjawal opened a new pull request, #22571:
URL: https://github.com/apache/datafusion/pull/22571
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
- Closes #22554.
## Rationale for this change
TopK aggregation dropped NULL group keys for ordered DISTINCT queries.
For example, `SELECT DISTINCT v FROM t ORDER BY v ASC NULLS FIRST LIMIT 1` could return an empty string instead of NULL when TopK aggregation was enabled.
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
## What changes are included in this PR?
This PR preserves NULL group keys for DISTINCT TopK aggregation by tracking whether a NULL group key was seen separately from the heap.
The heap still only stores non-NULL values. This avoids making the TopK heap implementations handle NULL values directly.
The stream also now marks itself done after emitting, so NULL-only DISTINCT results are emitted once and do not repeat.
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
## Are these changes tested?
Yes
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
## Are there any user-facing changes?
No API Change
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix TopK DISTINCT aggregation preserving NULLs [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #22571:
URL: https://github.com/apache/datafusion/pull/22571
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix TopK DISTINCT aggregation preserving NULLs [datafusion]
Posted by "avantgardnerio (via GitHub)" <gi...@apache.org>.
avantgardnerio commented on PR #22571:
URL: https://github.com/apache/datafusion/pull/22571#issuecomment-4565373184
> > The heap still only stores non-NULL values. This avoids making the TopK heap implementations handle NULL values directly.
>
> Also, perhaps @avantgardnerio you can help review this PR as I think you are familiar with the original logic
After some consideration, I think a separate bool is the way to go. The heap was generic over arrow::Primitive types, which does not make a solution allowing NULL straight-forward. We can probably make it generic over Option<arrow::Primitive>, but that's probably less performant? And it certainly stores more information than necessary.
The boolean seems like a pragmatic decision given performance concerns , which is why we do all this craziness in the first place. One note though: in the hot loop it's probably better to just always write the value rather than a conditional branch. Since we write much more than read, the check could move to the read side.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix TopK DISTINCT aggregation preserving NULLs [datafusion]
Posted by "kumarUjjawal (via GitHub)" <gi...@apache.org>.
kumarUjjawal commented on PR #22571:
URL: https://github.com/apache/datafusion/pull/22571#issuecomment-4560787640
> Looks good to me overall @kumarUjjawal -- i had a few questions
>
> I see you say
>
> > The heap still only stores non-NULL values. This avoids making the TopK heap implementations handle NULL values directly.
>
> I don't have any sense of how complicated this would be / not be -- did you try it?
>
> Also, perhaps @avantgardnerio you can help review this PR as I think you are familiar with the original logic
> I don't have any sense of how complicated this would be / not be -- did you try it?
I did try that as my first approach but I find it broad and since the only question I needed to answer was did we see a NULL group key and since SQL DISTINCT has only one NULL group, a boolean seemed good choice. Let me know what you think, I can go back to the other approach.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Fix TopK DISTINCT aggregation preserving NULLs [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #22571:
URL: https://github.com/apache/datafusion/pull/22571#discussion_r3312899353
##########
datafusion/physical-plan/src/aggregates/topk_stream.rs:
##########
@@ -209,17 +225,32 @@ impl Stream for GroupedTopKAggregateStream {
// Release the input pipeline's resources before emitting.
let input_schema = self.input.schema();
self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
- if self.priority_map.is_empty() {
+ if self.priority_map.is_empty() && !self.null_group_seen {
trace!("partition {} emit None", self.partition);
+ self.done = true;
return Poll::Ready(None);
}
let batch = {
let _timer = emitting_time.timer();
- let mut cols = self.priority_map.emit()?;
+ let mut cols = if self.priority_map.is_empty() {
+ vec![]
+ } else {
+ self.priority_map.emit()?
+ };
// For DISTINCT case (no aggregate expressions), only use the group key column
// since the schema only has one field and key/value are the same
- if self.aggregate_arguments.is_empty() {
+ if self.is_distinct() {
Review Comment:
can we try and encapsulate some of this logic in a helper function perhaps to try and keep this code easy to read?
It would also perhaps help to add some comments here about why concat is needed
##########
datafusion/physical-plan/src/aggregates/topk_stream.rs:
##########
@@ -138,6 +148,9 @@ impl GroupedTopKAggregateStream {
let has_nulls = vals.null_count() > 0;
for row_idx in 0..len {
if has_nulls && vals.is_null(row_idx) {
+ if self.is_distinct() {
+ self.null_group_seen = true;
Review Comment:
this is in the (hot) inner loop and I worry about the performance implications of this change
Is it enough to check outside the loop? something like
```rust
let has_nulls = vals.null_count() > 0;
if has_nulls && self.is_distinct() {
self.null_group_seen = true;
}
for row_idx in 0..len {
...
}
##########
datafusion/physical-plan/src/aggregates/topk_stream.rs:
##########
@@ -128,6 +134,10 @@ impl RecordBatchStream for GroupedTopKAggregateStream {
}
impl GroupedTopKAggregateStream {
+ fn is_distinct(&self) -> bool {
+ self.aggregate_arguments.is_empty()
Review Comment:
I think you can get no aggregates for just a normal `SELECT x,y,z GROUP BY x,y,z` type query, Though maybe that has the same semantics
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org