This is an automated email from the ASF dual-hosted git repository. kumarUjjawal pushed a commit to branch fix/topk-distinct-aggregation in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 0d228d9269771723caef939b391ebbc12dcaa556 Author: Kumar Ujjawal <[email protected]> AuthorDate: Wed May 27 20:46:39 2026 +0530 Fix TopK DISTINCT aggregation preserving NULLs --- .../physical_optimizer/aggregate_statistics.rs | 84 +++++++++++++++++ .../physical-plan/src/aggregates/topk_stream.rs | 40 +++++++- .../sqllogictest/test_files/aggregates_topk.slt | 105 +++++++++++++++++++++ 3 files changed, 225 insertions(+), 4 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..65a5ea1a71 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 (only tracked for DISTINCT queries) + 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_distinct(&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(); @@ -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; + } continue; } self.priority_map.insert(row_idx)?; @@ -153,6 +166,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) { @@ -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() { cols.truncate(1); + if self.null_group_seen { + 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 { + cols[0] = + concat(&[cols[0].as_ref(), null_arr.as_ref()])?; + } + } } RecordBatch::try_new(Arc::clone(&self.schema), cols)? }; @@ -232,6 +263,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..c45c047c86 100644 --- a/datafusion/sqllogictest/test_files/aggregates_topk.slt +++ b/datafusion/sqllogictest/test_files/aggregates_topk.slt @@ -456,6 +456,111 @@ 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'); + +# 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: [email protected] For additional commands, e-mail: [email protected]
