risingwave_frontend/optimizer/plan_visitor/
mod.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 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 locality_backfill_scan_estimator;
44mod locality_provider_counter;
45mod rw_timestamp_validator;
46mod sole_sys_table_visitor;
47pub use distributed_dml_visitor::*;
48pub use locality_backfill_scan_estimator::*;
49pub use locality_provider_counter::*;
50pub use rw_timestamp_validator::*;
51pub use sole_sys_table_visitor::*;
52
53#[cfg(feature = "datafusion")]
54mod datafusion_execute_checker;
55#[cfg(feature = "datafusion")]
56pub use datafusion_execute_checker::*;
57#[cfg(feature = "datafusion")]
58mod datafusion_plan_converter;
59#[cfg(feature = "datafusion")]
60pub use datafusion_plan_converter::*;
61
62use crate::for_each_convention_all_plan_nodes;
63use crate::optimizer::plan_node::*;
64
65/// The behavior for the default implementations of `visit_xxx`.
66pub trait DefaultBehavior<R> {
67    /// Apply this behavior to the plan node with the given results.
68    fn apply(&self, results: impl IntoIterator<Item = R>) -> R;
69}
70
71/// Visit all input nodes, merge the results with a function.
72/// - If there's no input node, return the default value of the result type.
73/// - If there's only a single input node, directly return its result.
74pub struct Merge<F>(F);
75
76impl<F, R> DefaultBehavior<R> for Merge<F>
77where
78    F: Fn(R, R) -> R,
79    R: Default,
80{
81    fn apply(&self, results: impl IntoIterator<Item = R>) -> R {
82        results.into_iter().reduce(&self.0).unwrap_or_default()
83    }
84}
85
86/// Visit all input nodes, return the default value of the result type.
87pub struct DefaultValue;
88
89impl<R> DefaultBehavior<R> for DefaultValue
90where
91    R: Default,
92{
93    fn apply(&self, results: impl IntoIterator<Item = R>) -> R {
94        let _ = results.into_iter().count(); // consume the iterator
95        R::default()
96    }
97}
98
99pub trait PlanVisitor<C: ConventionMarker> {
100    type Result;
101    fn visit(&mut self, plan: PlanRef<C>) -> Self::Result;
102}
103
104/// Define `PlanVisitor` trait.
105macro_rules! def_visitor {
106    ({
107        $( $convention:ident, { $( $name:ident ),* }),*
108    }) => {
109        paste! {
110            $(
111                /// The visitor for plan nodes. visit all inputs and return the ret value of the left most input,
112                /// and leaf node returns `R::default()`
113                pub trait [<$convention  PlanVisitor>] {
114                    type Result;
115                    type DefaultBehavior: DefaultBehavior<Self::Result>;
116
117                    /// The behavior for the default implementations of `visit_xxx`.
118                    fn default_behavior() -> Self::DefaultBehavior;
119
120                    fn [<visit_ $convention:snake>](&mut self, plan: PlanRef<$convention>) -> Self::Result {
121                        use risingwave_common::util::recursive::{tracker, Recurse};
122                        use crate::session::current::notice_to_user;
123
124                        tracker!().recurse(|t| {
125                            if t.depth_reaches(PLAN_DEPTH_THRESHOLD) {
126                                notice_to_user(PLAN_TOO_DEEP_NOTICE);
127                            }
128
129                            match plan.node_type() {
130                                $(
131                                    [<$convention PlanNodeType>]::[<$convention $name>] => self.[<visit_ $convention:snake _ $name:snake>](plan.downcast_ref::<[<$convention $name>]>().unwrap()),
132                                )*
133                            }
134                        })
135                    }
136
137                    $(
138                        #[doc = "Visit [`" [<$convention $name>] "`] , the function should visit the inputs."]
139                        fn [<visit_ $convention:snake _ $name:snake>](&mut self, plan: &[<$convention $name>]) -> Self::Result {
140                            let results = plan.inputs().into_iter().map(|input| self.[<visit_ $convention:snake>](input));
141                            Self::default_behavior().apply(results)
142                        }
143                    )*
144
145                }
146
147                impl<V: [<$convention  PlanVisitor>]> PlanVisitor<$convention> for V {
148                    type Result = V::Result;
149                    fn visit(&mut self, plan: PlanRef<$convention>) -> Self::Result {
150                        self.[<visit_ $convention:snake>](plan)
151                    }
152                }
153            )*
154        }
155    }
156}
157
158for_each_convention_all_plan_nodes! { def_visitor }
159
160macro_rules! impl_has_variant {
161    ( $({$convention:ident $variant_name:ident}),* ) => {
162        paste! {
163            $(
164                pub fn [<has_ $convention:snake _ $variant_name:snake _where>]<P>(plan: PlanRef<$convention>, pred: P) -> bool
165                where
166                    P: FnMut(&[<$convention $variant_name>]) -> bool,
167                {
168                    struct HasWhere<P> {
169                        pred: P,
170                    }
171
172                    impl<P> [<$convention PlanVisitor>] for HasWhere<P>
173                    where
174                        P: FnMut(&[<$convention $variant_name>]) -> bool,
175                    {
176                        type Result = bool;
177                        type DefaultBehavior = impl DefaultBehavior<Self::Result>;
178
179                        fn default_behavior() -> Self::DefaultBehavior {
180                            Merge(|a, b| a | b)
181                        }
182
183                        fn [<visit_ $convention:snake _ $variant_name:snake>](&mut self, node: &[<$convention $variant_name>]) -> Self::Result {
184                            (self.pred)(node)
185                        }
186                    }
187
188                    let mut visitor = HasWhere { pred };
189                    visitor.visit(plan)
190                }
191
192                #[allow(dead_code)]
193                pub fn [<has_ $convention:snake _ $variant_name:snake>](plan: PlanRef<$convention>) -> bool {
194                    [<has_ $convention:snake _$variant_name:snake _where>](plan, |_| true)
195                }
196            )*
197        }
198    };
199}
200
201impl_has_variant! {
202    {Logical Apply},
203    {Logical MaxOneRow},
204    {Logical OverWindow},
205    {Logical Scan},
206    {Logical Source},
207    {Batch Exchange},
208    {Batch SeqScan},
209    {Batch Source},
210    {Batch Insert},
211    {Batch Delete},
212    {Batch Update}
213}
risingwave_frontend/optimizer/plan_visitor/mod.rs

risingwave_frontend/optimizer/plan_visitor/
mod.rs
