Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 29 additions & 1 deletion datafusion/physical-optimizer/src/topk_aggregation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;
Expand Down Expand Up @@ -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()?;
Expand All @@ -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)),
}
Expand Down
17 changes: 16 additions & 1 deletion datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(())
}
Expand Down
Loading