risingwave_frontend/optimizer/plan_visitor/
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
15use paste::paste;
16mod apply_visitor;
17pub use apply_visitor::*;
18mod plan_correlated_id_finder;
19pub use plan_correlated_id_finder::*;
20mod share_parent_counter;
21pub use share_parent_counter::*;
22
23#[cfg(debug_assertions)]
24mod input_ref_validator;
25#[cfg(debug_assertions)]
26pub use input_ref_validator::*;
27
28mod execution_mode_decider;
29pub use execution_mode_decider::*;
30mod temporal_join_validator;
31pub use temporal_join_validator::*;
32mod relation_collector_visitor;
33mod sys_table_visitor;
34pub use relation_collector_visitor::*;
35pub use sys_table_visitor::*;
36mod side_effect_visitor;
37pub use side_effect_visitor::*;
38mod cardinality_visitor;
39pub use cardinality_visitor::*;
40mod jsonb_stream_key_checker;
41pub use jsonb_stream_key_checker::*;
42mod distributed_dml_visitor;
43mod read_storage_table_visitor;
44mod rw_timestamp_validator;
45pub use distributed_dml_visitor::*;
46pub use read_storage_table_visitor::*;
47pub use rw_timestamp_validator::*;
48
49use crate::for_all_plan_nodes;
50use crate::optimizer::plan_node::*;
51
52/// The behavior for the default implementations of `visit_xxx`.
53pub trait DefaultBehavior<R> {
54    /// Apply this behavior to the plan node with the given results.
55    fn apply(&self, results: impl IntoIterator<Item = R>) -> R;
56}
57
58/// Visit all input nodes, merge the results with a function.
59/// - If there's no input node, return the default value of the result type.
60/// - If there's only a single input node, directly return its result.
61pub struct Merge<F>(F);
62
63impl<F, R> DefaultBehavior<R> for Merge<F>
64where
65    F: Fn(R, R) -> R,
66    R: Default,
67{
68    fn apply(&self, results: impl IntoIterator<Item = R>) -> R {
69        results.into_iter().reduce(&self.0).unwrap_or_default()
70    }
71}
72
73/// Visit all input nodes, return the default value of the result type.
74pub struct DefaultValue;
75
76impl<R> DefaultBehavior<R> for DefaultValue
77where
78    R: Default,
79{
80    fn apply(&self, results: impl IntoIterator<Item = R>) -> R {
81        let _ = results.into_iter().count(); // consume the iterator
82        R::default()
83    }
84}
85
86/// Define `PlanVisitor` trait.
87macro_rules! def_visitor {
88    ($({ $convention:ident, $name:ident }),*) => {
89        /// The visitor for plan nodes. visit all inputs and return the ret value of the left most input,
90        /// and leaf node returns `R::default()`
91        pub trait PlanVisitor {
92            type Result: Default;
93            type DefaultBehavior: DefaultBehavior<Self::Result>;
94
95            /// The behavior for the default implementations of `visit_xxx`.
96            fn default_behavior() -> Self::DefaultBehavior;
97
98            paste! {
99                fn visit(&mut self, plan: PlanRef) -> Self::Result {
100                    use risingwave_common::util::recursive::{tracker, Recurse};
101                    use crate::session::current::notice_to_user;
102
103                    tracker!().recurse(|t| {
104                        if t.depth_reaches(PLAN_DEPTH_THRESHOLD) {
105                            notice_to_user(PLAN_TOO_DEEP_NOTICE);
106                        }
107
108                        match plan.node_type() {
109                            $(
110                                PlanNodeType::[<$convention $name>] => self.[<visit_ $convention:snake _ $name:snake>](plan.downcast_ref::<[<$convention $name>]>().unwrap()),
111                            )*
112                        }
113                    })
114                }
115
116                $(
117                    #[doc = "Visit [`" [<$convention $name>] "`] , the function should visit the inputs."]
118                    fn [<visit_ $convention:snake _ $name:snake>](&mut self, plan: &[<$convention $name>]) -> Self::Result {
119                        let results = plan.inputs().into_iter().map(|input| self.visit(input));
120                        Self::default_behavior().apply(results)
121                    }
122                )*
123            }
124        }
125    }
126}
127
128for_all_plan_nodes! { def_visitor }
129
130macro_rules! impl_has_variant {
131    ( $($variant:ty),* ) => {
132        paste! {
133            $(
134                pub fn [<has_ $variant:snake _where>]<P>(plan: PlanRef, pred: P) -> bool
135                where
136                    P: FnMut(&$variant) -> bool,
137                {
138                    struct HasWhere<P> {
139                        pred: P,
140                    }
141
142                    impl<P> PlanVisitor for HasWhere<P>
143                    where
144                        P: FnMut(&$variant) -> bool,
145                    {
146                        type Result = bool;
147                        type DefaultBehavior = impl DefaultBehavior<Self::Result>;
148
149                        fn default_behavior() -> Self::DefaultBehavior {
150                            Merge(|a, b| a | b)
151                        }
152
153                        fn [<visit_ $variant:snake>](&mut self, node: &$variant) -> Self::Result {
154                            (self.pred)(node)
155                        }
156                    }
157
158                    let mut visitor = HasWhere { pred };
159                    visitor.visit(plan)
160                }
161
162                #[allow(dead_code)]
163                pub fn [<has_ $variant:snake>](plan: PlanRef) -> bool {
164                    [<has_ $variant:snake _where>](plan, |_| true)
165                }
166            )*
167        }
168    };
169}
170
171impl_has_variant! {
172    LogicalApply,
173    LogicalMaxOneRow,
174    LogicalOverWindow,
175    LogicalScan,
176    LogicalSource,
177    BatchExchange,
178    BatchSeqScan,
179    BatchSource,
180    BatchInsert,
181    BatchDelete,
182    BatchUpdate
183}
risingwave_frontend/optimizer/plan_visitor/mod.rs

risingwave_frontend/optimizer/plan_visitor/
mod.rs
