fix: preserve NULL groups in TopK aggregation

Grouped TopK aggregation previously dropped groups whose MIN/MAX inputs were all NULL. Track bounded all-NULL candidates alongside valued groups and emit them for the final sort to rank and truncate. Handle NULL-to-value transitions, fully reuse freed hash-table slots, and keep NULL bookkeeping off the common no-NULL hot path.

Nullable MIN/MAX with NULLS FIRST is not monotonic for bounded aggregation: a group can move from NULL to a worse non-NULL rank. Skip the TopK pushdown in that case to preserve exact results while retaining it for NULLS LAST and non-nullable inputs.

Closes #23440, closes #22190.

Author: u70b3

datafusion/physical-optimizer/src/topk_aggregation.rs

@@ -46,6 +46,7 @@ impl TopKAggregation {
         aggr: &AggregateExec,
         order_by: &str,
         order_desc: bool,
+        nulls_first: bool,
         limit: usize,
     ) -> Option<Arc<dyn ExecutionPlan>> {
         // Current only support single group key
@@ -66,6 +67,26 @@ impl TopKAggregation {
 
         // Check if this is ordering by an aggregate function (MIN/MAX)
         if let Some((field, desc)) = aggr.get_minmax_desc() {
+            // A nullable MIN/MAX starts as NULL and becomes non-NULL when the
+            // group sees its first value. With NULLS FIRST that transition
+            // worsens the group's rank, so a bounded aggregation cannot safely
+            // discard other NULL groups. Use regular aggregation for exact
+            // results. Non-nullable inputs never take this transition and can
+            // still use TopK.
+            let input_nullable = aggr
+                .aggr_expr()
+                .iter()
+                .exactly_one()
+                .ok()?
+                .expressions()
+                .into_iter()
+                .exactly_one()
+                .ok()?
+                .nullable(aggr.input_schema.as_ref())
+                .ok()?;
+            if nulls_first && input_nullable {
+                return None;
+            }
             // ensure the sort direction matches aggregate function
             if desc != order_desc {
                 return None;
@@ -100,6 +121,7 @@ impl TopKAggregation {
         let order = sort.properties().output_ordering()?;
         let order = order.iter().exactly_one().ok()?;
         let order_desc = order.options.descending;
+        let nulls_first = order.options.nulls_first;
         let order = order.expr.downcast_ref::<Column>()?;
         let mut cur_col_name = order.name().to_string();
         let limit = sort.fetch()?;
@@ -111,7 +133,13 @@ impl TopKAggregation {
             }
             if let Some(aggr) = plan.downcast_ref::<AggregateExec>() {
                 // either we run into an Aggregate and transform it
-                match Self::transform_agg(aggr, &cur_col_name, order_desc, limit) {
+                match Self::transform_agg(
+                    aggr,
+                    &cur_col_name,
+                    order_desc,
+                    nulls_first,
+                    limit,
+                ) {
                     None => cardinality_preserved = false,
                     Some(plan) => return Ok(Transformed::yes(plan)),
                 }

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

@@ -149,11 +149,26 @@ impl GroupedTopKAggregateStream {
         if has_nulls && self.is_group_by_only() {
             self.null_group_seen = true;
         }
+        // Keep the common no-NULL path free of NULL bookkeeping. Once a NULL
+        // group exists, use the NULL-aware path until it has been resolved.
+        let track_null_groups = !self.is_group_by_only()
+            && (has_nulls || self.priority_map.has_null_groups());
         for row_idx in 0..len {
             if has_nulls && vals.is_null(row_idx) {
+                // MIN/MAX ignore NULL inputs, but a group whose values are all
+                // NULL must still be emitted with a NULL aggregate value, so
+                // track it. (GROUP BY-only aggregations handle NULL group keys
+                // via `null_group_seen` instead.)
+                if !self.is_group_by_only() {
+                    self.priority_map.insert_null(row_idx);
+                }
                 continue;
             }
-            self.priority_map.insert(row_idx)?;
+            if track_null_groups {
+                self.priority_map.insert_with_null_groups(row_idx)?;
+            } else {
+                self.priority_map.insert(row_idx)?;
+            }
         }
         Ok(())
     }