risingwave_frontend/optimizer/
heuristic_optimizer.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::collections::hash_map::Entry;
17use std::fmt;
18
19use itertools::Itertools;
20
21use super::ApplyResult;
22#[cfg(debug_assertions)]
23use crate::Explain;
24use crate::error::Result;
25use crate::optimizer::plan_node::{ConventionMarker, ShareNode};
26use crate::optimizer::rule::BoxedRule;
27use crate::optimizer::{PlanRef, PlanTreeNode, ShareId};
28
29/// Traverse order of [`HeuristicOptimizer`]
30pub enum ApplyOrder {
31    TopDown,
32    BottomUp,
33}
34
35// TODO: we should have a builder of HeuristicOptimizer here
36/// A rule-based heuristic optimizer, which traverses every plan nodes and tries to
37/// apply each rule on them.
38pub struct HeuristicOptimizer<'a, C: ConventionMarker> {
39    apply_order: &'a ApplyOrder,
40    rules: &'a [BoxedRule<C>],
41    stats: Stats,
42    share_cache: HashMap<ShareId, (PlanRef<C>, bool)>,
43}
44
45impl<'a, C: ConventionMarker> HeuristicOptimizer<'a, C> {
46    pub fn new(apply_order: &'a ApplyOrder, rules: &'a [BoxedRule<C>]) -> Self {
47        Self {
48            apply_order,
49            rules,
50            stats: Stats::new(),
51            share_cache: HashMap::new(),
52        }
53    }
54
55    fn optimize_node(&mut self, mut plan: PlanRef<C>) -> Result<(PlanRef<C>, bool)> {
56        let mut changed = false;
57        for rule in self.rules {
58            match rule.apply(plan.clone()) {
59                ApplyResult::Ok(applied) => {
60                    #[cfg(debug_assertions)]
61                    Self::check_equivalent_plan(rule.description(), &plan, &applied);
62
63                    plan = applied;
64                    changed = true;
65                    self.stats.count_rule(rule);
66                }
67                ApplyResult::NotApplicable => {}
68                ApplyResult::Err(error) => return Err(error),
69            }
70        }
71        Ok((plan, changed))
72    }
73
74    fn optimize_inputs(&mut self, plan: PlanRef<C>) -> Result<(PlanRef<C>, bool)> {
75        let optimized_inputs: Vec<_> = plan
76            .inputs()
77            .into_iter()
78            .map(|sub_tree| self.optimize_recursively(sub_tree))
79            .try_collect()?;
80        let changed = optimized_inputs.iter().any(|(_, changed)| *changed);
81
82        if changed {
83            let inputs = optimized_inputs
84                .into_iter()
85                .map(|(input, _)| input)
86                .collect_vec();
87            Ok((plan.clone_root_with_inputs(&inputs), true))
88        } else {
89            Ok((plan, false))
90        }
91    }
92
93    fn optimize_recursively(&mut self, plan: PlanRef<C>) -> Result<(PlanRef<C>, bool)> {
94        if let Some(share_id) = plan.as_share_node().map(ShareNode::share_id) {
95            if let Some((cached, changed)) = self.share_cache.get(&share_id) {
96                return Ok((cached.clone(), *changed));
97            }
98            let (optimized, changed) = self.optimize_uncached(plan)?;
99            self.share_cache
100                .insert(share_id, (optimized.clone(), changed));
101            return Ok((optimized, changed));
102        }
103
104        self.optimize_uncached(plan)
105    }
106
107    fn optimize_uncached(&mut self, plan: PlanRef<C>) -> Result<(PlanRef<C>, bool)> {
108        match self.apply_order {
109            ApplyOrder::TopDown => {
110                let (plan, node_changed) = self.optimize_node(plan)?;
111                let (plan, inputs_changed) = self.optimize_inputs(plan)?;
112                Ok((plan, node_changed || inputs_changed))
113            }
114            ApplyOrder::BottomUp => {
115                let (plan, inputs_changed) = self.optimize_inputs(plan)?;
116                let (plan, node_changed) = self.optimize_node(plan)?;
117                Ok((plan, inputs_changed || node_changed))
118            }
119        }
120    }
121
122    pub fn optimize(&mut self, plan: PlanRef<C>) -> Result<PlanRef<C>> {
123        self.share_cache.clear();
124        self.optimize_recursively(plan).map(|(plan, _)| plan)
125    }
126
127    pub fn get_stats(&self) -> &Stats {
128        &self.stats
129    }
130
131    #[cfg(debug_assertions)]
132    pub fn check_equivalent_plan(
133        rule_desc: &str,
134        input_plan: &PlanRef<C>,
135        output_plan: &PlanRef<C>,
136    ) {
137        use crate::optimizer::plan_node::generic::GenericPlanRef;
138        if !input_plan.schema().type_eq(output_plan.schema()) {
139            panic!(
140                "{} fails to generate equivalent plan.\nInput schema: {:?}\nInput plan: \n{}\nOutput schema: {:?}\nOutput plan: \n{}\nSQL: {}",
141                rule_desc,
142                input_plan.schema(),
143                input_plan.explain_to_string(),
144                output_plan.schema(),
145                output_plan.explain_to_string(),
146                output_plan.ctx().sql()
147            );
148        }
149    }
150}
151
152pub struct Stats {
153    total_applied: usize,
154    rule_counter: HashMap<String, u32>,
155}
156
157impl Stats {
158    pub fn new() -> Self {
159        Self {
160            rule_counter: HashMap::new(),
161            total_applied: 0,
162        }
163    }
164
165    pub fn count_rule(&mut self, rule: &BoxedRule<impl ConventionMarker>) {
166        self.total_applied += 1;
167        match self.rule_counter.entry(rule.description().to_owned()) {
168            Entry::Occupied(mut entry) => {
169                *entry.get_mut() += 1;
170            }
171            Entry::Vacant(entry) => {
172                entry.insert(1);
173            }
174        }
175    }
176
177    pub fn has_applied_rule(&self) -> bool {
178        self.total_applied != 0
179    }
180
181    pub fn total_applied(&self) -> usize {
182        self.total_applied
183    }
184}
185
186impl fmt::Display for Stats {
187    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
188        for (rule, count) in &self.rule_counter {
189            writeln!(f, "apply {} {} time(s)", rule, count)?;
190        }
191        Ok(())
192    }
193}
risingwave_frontend/optimizer/heuristic_optimizer.rs

risingwave_frontend/optimizer/
heuristic_optimizer.rs
