datafusion_physical_optimizer/

topk_aggregation.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! An optimizer rule that detects aggregate operations that could use a limited bucket count

use std::sync::Arc;

use crate::PhysicalOptimizerRule;
use datafusion_common::Result;
use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
use datafusion_physical_expr::expressions::Column;
use datafusion_physical_plan::ExecutionPlan;
use datafusion_physical_plan::aggregates::LimitOptions;
use datafusion_physical_plan::aggregates::{AggregateExec, topk_types_supported};
use datafusion_physical_plan::execution_plan::CardinalityEffect;
use datafusion_physical_plan::projection::ProjectionExec;
use datafusion_physical_plan::sorts::sort::SortExec;
use itertools::Itertools;

/// An optimizer rule that passes a `limit` hint to aggregations if the whole result is not needed
#[derive(Debug)]
pub struct TopKAggregation {}

impl TopKAggregation {
    /// Create a new `LimitAggregation`
    pub fn new() -> Self {
        Self {}
    }

    fn transform_agg(
        aggr: &AggregateExec,
        order_by: &str,
        order_desc: bool,
        nulls_first: bool,
        limit: usize,
    ) -> Option<Arc<dyn ExecutionPlan>> {
        // Current only support single group key
        let (group_key, group_key_alias) =
            aggr.group_expr().expr().iter().exactly_one().ok()?;
        let kt = group_key.data_type(&aggr.input().schema()).ok()?;
        let vt = if let Some((field, _)) = aggr.get_minmax_desc() {
            field.data_type().clone()
        } else {
            kt.clone()
        };
        if !topk_types_supported(&kt, &vt) {
            return None;
        }
        if aggr.filter_expr().iter().any(|e| e.is_some()) {
            return None;
        }

        // 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;
            }
            // ensure the sort is on the same field as the aggregate output
            if order_by != field.name() {
                return None;
            }
        } else if aggr.aggr_expr().is_empty() {
            // This is a GROUP BY without aggregates, check if ordering is on the group key itself
            if order_by != group_key_alias {
                return None;
            }
        } else {
            // Has aggregates but not MIN/MAX, or doesn't DISTINCT
            return None;
        }

        // We found what we want: clone, copy the limit down, and return modified node
        let new_aggr = AggregateExec::with_new_limit_options(
            aggr,
            Some(LimitOptions::new_with_order(limit, order_desc)),
        );
        Some(Arc::new(new_aggr))
    }

    fn transform_sort(plan: &Arc<dyn ExecutionPlan>) -> Option<Arc<dyn ExecutionPlan>> {
        let sort = plan.downcast_ref::<SortExec>()?;

        let children = sort.children();
        let child = children.into_iter().exactly_one().ok()?;
        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()?;

        let mut cardinality_preserved = true;
        let closure = |plan: Arc<dyn ExecutionPlan>| {
            if !cardinality_preserved {
                return Ok(Transformed::no(plan));
            }
            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,
                    nulls_first,
                    limit,
                ) {
                    None => cardinality_preserved = false,
                    Some(plan) => return Ok(Transformed::yes(plan)),
                }
            } else if let Some(proj) = plan.downcast_ref::<ProjectionExec>() {
                // track renames due to successive projections
                for proj_expr in proj.expr() {
                    let Some(src_col) = proj_expr.expr.downcast_ref::<Column>() else {
                        continue;
                    };
                    if proj_expr.alias == cur_col_name {
                        cur_col_name = src_col.name().to_string();
                    }
                }
            } else {
                // or we continue down through types that don't reduce cardinality
                match plan.cardinality_effect() {
                    CardinalityEffect::Equal | CardinalityEffect::GreaterEqual => {}
                    CardinalityEffect::Unknown | CardinalityEffect::LowerEqual => {
                        cardinality_preserved = false;
                    }
                }
            }
            Ok(Transformed::no(plan))
        };
        let child = Arc::clone(child).transform_down(closure).data().ok()?;
        let sort = SortExec::new(sort.expr().clone(), child)
            .with_fetch(sort.fetch())
            .with_preserve_partitioning(sort.preserve_partitioning());
        Some(Arc::new(sort))
    }
}

impl Default for TopKAggregation {
    fn default() -> Self {
        Self::new()
    }
}

impl PhysicalOptimizerRule for TopKAggregation {
    fn optimize(
        &self,
        plan: Arc<dyn ExecutionPlan>,
        config: &ConfigOptions,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        if config.optimizer.enable_topk_aggregation {
            plan.transform_down(|plan| {
                Ok(if let Some(plan) = TopKAggregation::transform_sort(&plan) {
                    Transformed::yes(plan)
                } else {
                    Transformed::no(plan)
                })
            })
            .data()
        } else {
            Ok(plan)
        }
    }

    fn name(&self) -> &str {
        "LimitAggregation"
    }

    fn schema_check(&self) -> bool {
        true
    }
}

// see `aggregate.slt` for tests