risingwave_frontend/optimizer/plan_rewriter/
mod.rs

1// Copyright 2025 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
15mod plan_cloner;
16mod share_source_rewriter;
17
18use itertools::Itertools;
19use paste::paste;
20pub use plan_cloner::*;
21pub use share_source_rewriter::*;
22
23use crate::for_all_plan_nodes;
24use crate::optimizer::plan_node::*;
25
26macro_rules! def_rewrite {
27    ( $convention:ident, Share ) => {
28        paste! {
29            /// When we use the plan rewriter, we need to take care of the share operator,
30            /// because our plan is a DAG rather than a tree.
31            /// Make sure this method can keep the shape of DAG.
32            fn [<rewrite_ $convention:snake _ share>](&mut self, plan: &[<$convention Share>]) -> PlanRef;
33        }
34    };
35
36    ( $convention:ident, $name:ident ) => {
37        paste! {
38            #[doc = "Visit [`" [<$convention $name>] "`] , the function should rewrite the inputs."]
39            fn [<rewrite_ $convention:snake _ $name:snake>](&mut self, plan: &[<$convention $name>]) -> PlanRef {
40                let new_inputs = plan
41                    .inputs()
42                    .into_iter()
43                    .map(|input| self.rewrite(input.clone()))
44                    .collect_vec();
45                plan.clone_with_inputs(&new_inputs)
46            }
47        }
48    };
49}
50
51/// Define `PlanRewriter` trait.
52macro_rules! def_rewriter {
53    ($({ $convention:ident, $name:ident }),*) => {
54
55        /// it's kind of like a [`PlanVisitor<PlanRef>`](super::plan_visitor::PlanVisitor), but with default behaviour of each rewrite method
56        pub trait PlanRewriter {
57            paste! {
58                fn rewrite(&mut self, plan: PlanRef) -> PlanRef{
59                    use risingwave_common::util::recursive::{tracker, Recurse};
60                    use crate::session::current::notice_to_user;
61
62                    tracker!().recurse(|t| {
63                        if t.depth_reaches(PLAN_DEPTH_THRESHOLD) {
64                            notice_to_user(PLAN_TOO_DEEP_NOTICE);
65                        }
66
67                        match plan.node_type() {
68                            $(
69                                PlanNodeType::[<$convention $name>] => self.[<rewrite_ $convention:snake _ $name:snake>](plan.downcast_ref::<[<$convention $name>]>().unwrap()),
70                            )*
71                        }
72                    })
73                }
74
75                $(
76                    def_rewrite! {$convention, $name}
77                )*
78            }
79        }
80    };
81}
82for_all_plan_nodes! { def_rewriter }
risingwave_frontend/optimizer/plan_rewriter/mod.rs

risingwave_frontend/optimizer/plan_rewriter/
mod.rs
