You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/05/29 20:11:01 UTC
(datafusion) branch gh-readonly-queue/main/pr-22571-d8c4588285c7db993935292c3f0bae6df3a171d0 created (now e7a2e052c3)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-22571-d8c4588285c7db993935292c3f0bae6df3a171d0
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at e7a2e052c3 Fix TopK DISTINCT aggregation preserving NULLs (#22571)
This branch includes the following new commits:
new e7a2e052c3 Fix TopK DISTINCT aggregation preserving NULLs (#22571)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: Fix TopK DISTINCT aggregation preserving NULLs (#22571)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22571-d8c4588285c7db993935292c3f0bae6df3a171d0
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit e7a2e052c319de23db36a3e209640413dccf7f0c
Author: Kumar Ujjawal <uj...@gmail.com>
AuthorDate: Sat May 30 01:40:40 2026 +0530
Fix TopK DISTINCT aggregation preserving NULLs (#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.
-->
---
.../physical_optimizer/aggregate_statistics.rs | 84 ++++++++++++++
.../physical-plan/src/aggregates/topk_stream.rs | 66 +++++++++--
.../sqllogictest/test_files/aggregates_topk.slt | 121 +++++++++++++++++++++
3 files changed, 261 insertions(+), 10 deletions(-)
diff --git a/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs b/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs
index 808e163b08..0fa60ae20d 100644
--- a/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs
+++ b/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs
@@ -553,3 +553,87 @@ async fn test_count_distinct_optimization() -> Result<()> {
Ok(())
}
+
+/// Regression test for https://github.com/apache/datafusion/issues/22554
+///
+/// TopK aggregation for DISTINCT queries was unconditionally dropping NULL
+/// group keys, producing wrong results with NULLS FIRST / NULLS LAST ordering.
+#[tokio::test]
+async fn topk_distinct_preserves_nulls() -> Result<()> {
+ let ctx = SessionContext::new_with_config(SessionConfig::new());
+
+ let batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new("v", DataType::Utf8, true)])),
+ vec![Arc::new(StringArray::from(vec![None, Some(""), Some("a")]))],
+ )?;
+ let table = MemTable::try_new(batch.schema(), vec![vec![batch]])?;
+ ctx.register_table("t", Arc::new(table))?;
+
+ // ASC NULLS FIRST LIMIT 1 → NULL should come first
+ let result = ctx
+ .sql("SELECT DISTINCT v FROM t ORDER BY v ASC NULLS FIRST LIMIT 1")
+ .await?
+ .collect()
+ .await?;
+ assert_batches_eq!(&["+---+", "| v |", "+---+", "| |", "+---+"], &result);
+ assert!(result[0].column(0).is_null(0), "first row should be NULL");
+
+ // ASC NULLS FIRST LIMIT 2 → NULL, then empty string
+ let result = ctx
+ .sql("SELECT DISTINCT v FROM t ORDER BY v ASC NULLS FIRST LIMIT 2")
+ .await?
+ .collect()
+ .await?;
+ assert_eq!(result[0].num_rows(), 2);
+ assert!(result[0].column(0).is_null(0));
+ assert!(!result[0].column(0).is_null(1));
+
+ // ASC NULLS LAST LIMIT 1 → empty string (smallest non-null)
+ let result = ctx
+ .sql("SELECT DISTINCT v FROM t ORDER BY v ASC NULLS LAST LIMIT 1")
+ .await?
+ .collect()
+ .await?;
+ assert!(
+ !result[0].column(0).is_null(0),
+ "first row should NOT be NULL"
+ );
+
+ // Full result with NULLS LAST should include NULL at end
+ let result = ctx
+ .sql("SELECT DISTINCT v FROM t ORDER BY v ASC NULLS LAST LIMIT 3")
+ .await?
+ .collect()
+ .await?;
+ assert_eq!(result[0].num_rows(), 3);
+ assert!(result[0].column(0).is_null(2), "last row should be NULL");
+
+ // Integer column
+ let batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, true)])),
+ vec![Arc::new(Int64Array::from(vec![None, Some(3), Some(1)]))],
+ )?;
+ let table = MemTable::try_new(batch.schema(), vec![vec![batch]])?;
+ ctx.register_table("t_int", Arc::new(table))?;
+
+ let result = ctx
+ .sql("SELECT DISTINCT v FROM t_int ORDER BY v ASC NULLS FIRST LIMIT 1")
+ .await?
+ .collect()
+ .await?;
+ assert!(
+ result[0].column(0).is_null(0),
+ "integer NULL should be first"
+ );
+
+ let result = ctx
+ .sql("SELECT DISTINCT v FROM t_int ORDER BY v DESC NULLS LAST LIMIT 2")
+ .await?
+ .collect()
+ .await?;
+ assert_eq!(result[0].num_rows(), 2);
+ assert!(!result[0].column(0).is_null(0));
+ assert!(!result[0].column(0).is_null(1));
+
+ Ok(())
+}
diff --git a/datafusion/physical-plan/src/aggregates/topk_stream.rs b/datafusion/physical-plan/src/aggregates/topk_stream.rs
index 9128844f1d..97f4662c11 100644
--- a/datafusion/physical-plan/src/aggregates/topk_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/topk_stream.rs
@@ -28,7 +28,8 @@ use crate::aggregates::{
use crate::metrics::BaselineMetrics;
use crate::stream::EmptyRecordBatchStream;
use crate::{RecordBatchStream, SendableRecordBatchStream};
-use arrow::array::{Array, ArrayRef, RecordBatch};
+use arrow::array::{Array, ArrayRef, RecordBatch, new_null_array};
+use arrow::compute::concat;
use arrow::datatypes::SchemaRef;
use arrow::util::pretty::print_batches;
use datafusion_common::Result;
@@ -46,6 +47,7 @@ pub struct GroupedTopKAggregateStream {
partition: usize,
row_count: usize,
started: bool,
+ done: bool,
schema: SchemaRef,
input: SendableRecordBatchStream,
baseline_metrics: BaselineMetrics,
@@ -53,6 +55,8 @@ pub struct GroupedTopKAggregateStream {
aggregate_arguments: Vec<Vec<Arc<dyn PhysicalExpr>>>,
group_by: Arc<PhysicalGroupBy>,
priority_map: PriorityMap,
+ /// Whether a NULL group key has been seen for a group-by-only aggregation.
+ null_group_seen: bool,
}
impl GroupedTopKAggregateStream {
@@ -109,6 +113,7 @@ impl GroupedTopKAggregateStream {
Ok(GroupedTopKAggregateStream {
partition,
started: false,
+ done: false,
row_count: 0,
schema: agg_schema,
input,
@@ -117,6 +122,7 @@ impl GroupedTopKAggregateStream {
aggregate_arguments,
group_by,
priority_map,
+ null_group_seen: false,
})
}
}
@@ -128,6 +134,10 @@ impl RecordBatchStream for GroupedTopKAggregateStream {
}
impl GroupedTopKAggregateStream {
+ fn is_group_by_only(&self) -> bool {
+ self.aggregate_arguments.is_empty()
+ }
+
fn intern(&mut self, ids: &ArrayRef, vals: &ArrayRef) -> Result<()> {
let _timer = self.group_by_metrics.time_calculating_group_ids.timer();
@@ -136,6 +146,9 @@ impl GroupedTopKAggregateStream {
.set_batch(Arc::clone(ids), Arc::clone(vals));
let has_nulls = vals.null_count() > 0;
+ if has_nulls && self.is_group_by_only() {
+ self.null_group_seen = true;
+ }
for row_idx in 0..len {
if has_nulls && vals.is_null(row_idx) {
continue;
@@ -144,6 +157,39 @@ impl GroupedTopKAggregateStream {
}
Ok(())
}
+
+ fn emit_columns(&mut self) -> Result<Vec<ArrayRef>> {
+ let mut cols = if self.priority_map.is_empty() {
+ vec![]
+ } else {
+ self.priority_map.emit()?
+ };
+
+ // GROUP BY-only aggregation covers DISTINCT-like queries. The group
+ // key and heap value are the same column, but the output schema has
+ // only the group key.
+ if self.is_group_by_only() {
+ cols.truncate(1);
+ if self.null_group_seen {
+ self.append_null_group(&mut cols)?;
+ }
+ }
+
+ Ok(cols)
+ }
+
+ fn append_null_group(&self, cols: &mut Vec<ArrayRef>) -> Result<()> {
+ let dt = self.schema.field(0).data_type();
+ let null_arr = new_null_array(dt, 1);
+ if cols.is_empty() {
+ cols.push(null_arr);
+ } else {
+ // NULL group keys are tracked outside the heap, so append a
+ // one-row NULL array to the emitted non-NULL group key column.
+ cols[0] = concat(&[cols[0].as_ref(), null_arr.as_ref()])?;
+ }
+ Ok(())
+ }
}
impl Stream for GroupedTopKAggregateStream {
@@ -153,6 +199,9 @@ impl Stream for GroupedTopKAggregateStream {
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Self::Item>> {
+ if self.done {
+ return Poll::Ready(None);
+ }
let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
let emitting_time = self.group_by_metrics.emitting_time.clone();
while let Poll::Ready(res) = self.input.poll_next_unpin(cx) {
@@ -185,8 +234,8 @@ impl Stream for GroupedTopKAggregateStream {
"Exactly 1 group value required"
);
let group_by_values = Arc::clone(&group_by_values[0][0]);
- let input_values = if self.aggregate_arguments.is_empty() {
- // DISTINCT case: use group key as both key and value
+ let input_values = if self.is_group_by_only() {
+ // GROUP BY-only case: use group key as both key and value
Arc::clone(&group_by_values)
} else {
// MIN/MAX case: evaluate aggregate expressions
@@ -209,18 +258,14 @@ 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()?;
- // 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() {
- cols.truncate(1);
- }
+ let cols = self.emit_columns()?;
RecordBatch::try_new(Arc::clone(&self.schema), cols)?
};
let batch = batch.record_output(&self.baseline_metrics);
@@ -232,6 +277,7 @@ impl Stream for GroupedTopKAggregateStream {
if log::log_enabled!(Level::Trace) {
print_batches(std::slice::from_ref(&batch))?;
}
+ self.done = true;
return Poll::Ready(Some(Ok(batch)));
}
// inner had error, return to caller
diff --git a/datafusion/sqllogictest/test_files/aggregates_topk.slt b/datafusion/sqllogictest/test_files/aggregates_topk.slt
index 19ead8965e..81c85c433b 100644
--- a/datafusion/sqllogictest/test_files/aggregates_topk.slt
+++ b/datafusion/sqllogictest/test_files/aggregates_topk.slt
@@ -456,6 +456,127 @@ select count(*) from (select category from values_table group by category order
----
3
+# Test DISTINCT with NULLs and NULLS FIRST ordering (issue #22554)
+statement ok
+create table nullable_vals (v varchar) as values (NULL), (''), ('a'), ('b');
+
+# Verify this regression test exercises the TopK aggregation path
+query TT
+explain select distinct v from nullable_vals order by v asc nulls first limit 1;
+----
+logical_plan
+01)Sort: nullable_vals.v ASC NULLS FIRST, fetch=1
+02)--Aggregate: groupBy=[[nullable_vals.v]], aggr=[[]]
+03)----TableScan: nullable_vals projection=[v]
+physical_plan
+01)SortPreservingMergeExec: [v@0 ASC], fetch=1
+02)--SortExec: TopK(fetch=1), expr=[v@0 ASC], preserve_partitioning=[true]
+03)----AggregateExec: mode=FinalPartitioned, gby=[v@0 as v], aggr=[], lim=[1]
+04)------RepartitionExec: partitioning=Hash([v@0], 4), input_partitions=1
+05)--------AggregateExec: mode=Partial, gby=[v@0 as v], aggr=[], lim=[1]
+06)----------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# NULLS FIRST: NULL should be the first row returned by LIMIT
+query T
+select distinct v from nullable_vals order by v asc nulls first limit 1;
+----
+NULL
+
+query T
+select distinct v from nullable_vals order by v asc nulls first limit 2;
+----
+NULL
+(empty)
+
+query T
+select distinct v from nullable_vals order by v asc nulls first limit 3;
+----
+NULL
+(empty)
+a
+
+# NULLS LAST: non-null values come first
+query T
+select distinct v from nullable_vals order by v asc nulls last limit 1;
+----
+(empty)
+
+query T
+select distinct v from nullable_vals order by v asc nulls last limit 4;
+----
+(empty)
+a
+b
+NULL
+
+# DESC NULLS FIRST: NULL comes first
+query T
+select distinct v from nullable_vals order by v desc nulls first limit 1;
+----
+NULL
+
+# DESC NULLS LAST: NULL comes last
+query T
+select distinct v from nullable_vals order by v desc nulls last limit 1;
+----
+b
+
+query T
+select distinct v from nullable_vals order by v desc nulls last limit 4;
+----
+b
+a
+(empty)
+NULL
+
+# Test with integer column containing NULLs
+statement ok
+create table nullable_ints (v int) as values (NULL), (3), (1), (2);
+
+query I
+select distinct v from nullable_ints order by v asc nulls first limit 1;
+----
+NULL
+
+query I
+select distinct v from nullable_ints order by v asc nulls first limit 3;
+----
+NULL
+1
+2
+
+query I
+select distinct v from nullable_ints order by v desc nulls last limit 2;
+----
+3
+2
+
+query I
+select distinct v from nullable_ints order by v asc nulls last limit 4;
+----
+1
+2
+3
+NULL
+
+# Test with all-NULL column
+statement ok
+create table all_nulls (v varchar) as values (NULL), (NULL);
+
+query T
+select distinct v from all_nulls order by v asc nulls first limit 1;
+----
+NULL
+
+statement ok
+drop table nullable_vals;
+
+statement ok
+drop table nullable_ints;
+
+statement ok
+drop table all_nulls;
+
statement ok
drop table values_table;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org