datafusion_physical_optimizer/
topk_aggregation.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! An optimizer rule that detects aggregate operations that could use a limited bucket count
19
20use std::sync::Arc;
21
22use crate::PhysicalOptimizerRule;
23use arrow::datatypes::DataType;
24use datafusion_common::config::ConfigOptions;
25use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
26use datafusion_common::Result;
27use datafusion_physical_expr::expressions::Column;
28use datafusion_physical_plan::aggregates::AggregateExec;
29use datafusion_physical_plan::execution_plan::CardinalityEffect;
30use datafusion_physical_plan::projection::ProjectionExec;
31use datafusion_physical_plan::sorts::sort::SortExec;
32use datafusion_physical_plan::ExecutionPlan;
33use itertools::Itertools;
34
35/// An optimizer rule that passes a `limit` hint to aggregations if the whole result is not needed
36#[derive(Debug)]
37pub struct TopKAggregation {}
38
39impl TopKAggregation {
40    /// Create a new `LimitAggregation`
41    pub fn new() -> Self {
42        Self {}
43    }
44
45    fn transform_agg(
46        aggr: &AggregateExec,
47        order_by: &str,
48        order_desc: bool,
49        limit: usize,
50    ) -> Option<Arc<dyn ExecutionPlan>> {
51        // ensure the sort direction matches aggregate function
52        let (field, desc) = aggr.get_minmax_desc()?;
53        if desc != order_desc {
54            return None;
55        }
56        let group_key = aggr.group_expr().expr().iter().exactly_one().ok()?;
57        let kt = group_key.0.data_type(&aggr.input().schema()).ok()?;
58        if !kt.is_primitive()
59            && kt != DataType::Utf8
60            && kt != DataType::Utf8View
61            && kt != DataType::LargeUtf8
62        {
63            return None;
64        }
65        if aggr.filter_expr().iter().any(|e| e.is_some()) {
66            return None;
67        }
68
69        // ensure the sort is on the same field as the aggregate output
70        if order_by != field.name() {
71            return None;
72        }
73
74        // We found what we want: clone, copy the limit down, and return modified node
75        let new_aggr = AggregateExec::try_new(
76            *aggr.mode(),
77            aggr.group_expr().clone(),
78            aggr.aggr_expr().to_vec(),
79            aggr.filter_expr().to_vec(),
80            Arc::clone(aggr.input()),
81            aggr.input_schema(),
82        )
83        .expect("Unable to copy Aggregate!")
84        .with_limit(Some(limit));
85        Some(Arc::new(new_aggr))
86    }
87
88    fn transform_sort(plan: &Arc<dyn ExecutionPlan>) -> Option<Arc<dyn ExecutionPlan>> {
89        let sort = plan.as_any().downcast_ref::<SortExec>()?;
90
91        let children = sort.children();
92        let child = children.into_iter().exactly_one().ok()?;
93        let order = sort.properties().output_ordering()?;
94        let order = order.iter().exactly_one().ok()?;
95        let order_desc = order.options.descending;
96        let order = order.expr.as_any().downcast_ref::<Column>()?;
97        let mut cur_col_name = order.name().to_string();
98        let limit = sort.fetch()?;
99
100        let mut cardinality_preserved = true;
101        let closure = |plan: Arc<dyn ExecutionPlan>| {
102            if !cardinality_preserved {
103                return Ok(Transformed::no(plan));
104            }
105            if let Some(aggr) = plan.as_any().downcast_ref::<AggregateExec>() {
106                // either we run into an Aggregate and transform it
107                match Self::transform_agg(aggr, &cur_col_name, order_desc, limit) {
108                    None => cardinality_preserved = false,
109                    Some(plan) => return Ok(Transformed::yes(plan)),
110                }
111            } else if let Some(proj) = plan.as_any().downcast_ref::<ProjectionExec>() {
112                // track renames due to successive projections
113                for proj_expr in proj.expr() {
114                    let Some(src_col) = proj_expr.expr.as_any().downcast_ref::<Column>()
115                    else {
116                        continue;
117                    };
118                    if proj_expr.alias == cur_col_name {
119                        cur_col_name = src_col.name().to_string();
120                    }
121                }
122            } else {
123                // or we continue down through types that don't reduce cardinality
124                match plan.cardinality_effect() {
125                    CardinalityEffect::Equal | CardinalityEffect::GreaterEqual => {}
126                    CardinalityEffect::Unknown | CardinalityEffect::LowerEqual => {
127                        cardinality_preserved = false;
128                    }
129                }
130            }
131            Ok(Transformed::no(plan))
132        };
133        let child = Arc::clone(child).transform_down(closure).data().ok()?;
134        let sort = SortExec::new(sort.expr().clone(), child)
135            .with_fetch(sort.fetch())
136            .with_preserve_partitioning(sort.preserve_partitioning());
137        Some(Arc::new(sort))
138    }
139}
140
141impl Default for TopKAggregation {
142    fn default() -> Self {
143        Self::new()
144    }
145}
146
147impl PhysicalOptimizerRule for TopKAggregation {
148    fn optimize(
149        &self,
150        plan: Arc<dyn ExecutionPlan>,
151        config: &ConfigOptions,
152    ) -> Result<Arc<dyn ExecutionPlan>> {
153        if config.optimizer.enable_topk_aggregation {
154            plan.transform_down(|plan| {
155                Ok(if let Some(plan) = TopKAggregation::transform_sort(&plan) {
156                    Transformed::yes(plan)
157                } else {
158                    Transformed::no(plan)
159                })
160            })
161            .data()
162        } else {
163            Ok(plan)
164        }
165    }
166
167    fn name(&self) -> &str {
168        "LimitAggregation"
169    }
170
171    fn schema_check(&self) -> bool {
172        true
173    }
174}
175
176// see `aggregate.slt` for tests
datafusion_physical_optimizer/topk_aggregation.rs

datafusion_physical_optimizer/
topk_aggregation.rs
