risingwave_frontend/optimizer/plan_rewriter/
mod.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
// Copyright 2024 RisingWave Labs
//
// Licensed 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.

mod plan_cloner;
mod share_source_rewriter;

use itertools::Itertools;
use paste::paste;
pub use plan_cloner::*;
pub use share_source_rewriter::*;

use crate::for_all_plan_nodes;
use crate::optimizer::plan_node::*;

macro_rules! def_rewrite {
    ( $convention:ident, Share ) => {
        paste! {
            /// When we use the plan rewriter, we need to take care of the share operator,
            /// because our plan is a DAG rather than a tree.
            /// Make sure this method can keep the shape of DAG.
            fn [<rewrite_ $convention:snake _ share>](&mut self, plan: &[<$convention Share>]) -> PlanRef;
        }
    };

    ( $convention:ident, $name:ident ) => {
        paste! {
            #[doc = "Visit [`" [<$convention $name>] "`] , the function should rewrite the inputs."]
            fn [<rewrite_ $convention:snake _ $name:snake>](&mut self, plan: &[<$convention $name>]) -> PlanRef {
                let new_inputs = plan
                    .inputs()
                    .into_iter()
                    .map(|input| self.rewrite(input.clone()))
                    .collect_vec();
                plan.clone_with_inputs(&new_inputs)
            }
        }
    };
}

/// Define `PlanRewriter` trait.
macro_rules! def_rewriter {
    ($({ $convention:ident, $name:ident }),*) => {

        /// it's kind of like a [`PlanVisitor<PlanRef>`](super::plan_visitor::PlanVisitor), but with default behaviour of each rewrite method
        pub trait PlanRewriter {
            paste! {
                fn rewrite(&mut self, plan: PlanRef) -> PlanRef{
                    use risingwave_common::util::recursive::{tracker, Recurse};
                    use crate::session::current::notice_to_user;

                    tracker!().recurse(|t| {
                        if t.depth_reaches(PLAN_DEPTH_THRESHOLD) {
                            notice_to_user(PLAN_TOO_DEEP_NOTICE);
                        }

                        match plan.node_type() {
                            $(
                                PlanNodeType::[<$convention $name>] => self.[<rewrite_ $convention:snake _ $name:snake>](plan.downcast_ref::<[<$convention $name>]>().unwrap()),
                            )*
                        }
                    })
                }

                $(
                    def_rewrite! {$convention, $name}
                )*
            }
        }
    };
}
for_all_plan_nodes! { def_rewriter }