risingwave_frontend/optimizer/plan_node/
merge_eq_nodes.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::hash::Hash;
17
18use super::{
19    EndoPlan, LogicalPlanRef as PlanRef, LogicalShare, PlanTreeNodeUnary, ShareNode, VisitPlan,
20};
21use crate::optimizer::{ShareId, plan_visitor};
22use crate::utils::{Endo, Visit};
23
24pub trait Semantics<V: Hash + Eq> {
25    fn semantics(&self) -> V;
26}
27
28impl Semantics<PlanRef> for PlanRef {
29    fn semantics(&self) -> PlanRef {
30        self.clone()
31    }
32}
33
34impl PlanRef {
35    pub fn common_subplan_sharing<V: Hash + Eq>(self) -> PlanRef
36    where
37        PlanRef: Semantics<V>,
38    {
39        Merger::default().apply(self)
40    }
41}
42
43struct Merger<V: Hash + Eq> {
44    cache: HashMap<V, LogicalShare>,
45}
46
47impl<V: Hash + Eq> Default for Merger<V> {
48    fn default() -> Self {
49        Merger {
50            cache: Default::default(),
51        }
52    }
53}
54
55impl<V: Hash + Eq> Endo<PlanRef> for Merger<V>
56where
57    PlanRef: Semantics<V>,
58{
59    fn apply(&mut self, t: PlanRef) -> PlanRef {
60        let semantics = t.semantics();
61        let share = self.cache.get(&semantics).cloned().unwrap_or_else(|| {
62            let share = LogicalShare::new(self.tree_apply(t));
63            self.cache.entry(semantics).or_insert(share).clone()
64        });
65        share.into()
66    }
67}
68
69impl PlanRef {
70    pub fn prune_share(&self) -> PlanRef {
71        let mut counter = Counter::default();
72        counter.visit(self);
73        counter.to_pruner().apply(self.clone())
74    }
75}
76
77#[derive(Default)]
78struct Counter {
79    counts: HashMap<ShareId, u64>,
80}
81
82impl Counter {
83    fn to_pruner(&self) -> Pruner<'_> {
84        Pruner {
85            counts: &self.counts,
86            cache: HashMap::new(),
87        }
88    }
89}
90
91impl VisitPlan for Counter {
92    fn visited<F>(&mut self, plan: &PlanRef, mut f: F)
93    where
94        F: FnMut(&mut Self),
95    {
96        let share_id = plan
97            .as_logical_share()
98            .expect("dag cache is only used for shares")
99            .share_id();
100        if self.counts.get(&share_id).is_none_or(|c| *c <= 1) {
101            f(self);
102        }
103    }
104}
105
106impl Visit<PlanRef> for Counter {
107    fn visit(&mut self, t: &PlanRef) {
108        if let Some(s) = t.as_logical_share() {
109            self.counts
110                .entry(s.share_id())
111                .and_modify(|c| *c += 1)
112                .or_insert(1);
113        }
114        self.dag_visit(t);
115    }
116}
117
118struct Pruner<'a> {
119    counts: &'a HashMap<ShareId, u64>,
120    cache: HashMap<ShareId, PlanRef>,
121}
122
123impl EndoPlan for Pruner<'_> {
124    fn cached<F>(&mut self, plan: PlanRef, mut f: F) -> PlanRef
125    where
126        F: FnMut(&mut Self) -> PlanRef,
127    {
128        let share_id = plan
129            .as_logical_share()
130            .expect("dag cache is only used for shares")
131            .share_id();
132        self.cache.get(&share_id).cloned().unwrap_or_else(|| {
133            let res = f(self);
134            self.cache.entry(share_id).or_insert(res).clone()
135        })
136    }
137}
138
139impl Endo<PlanRef> for Pruner<'_> {
140    fn pre(&mut self, t: PlanRef) -> PlanRef {
141        let prunable = |s: &&LogicalShare| {
142            // Prune if share node has only one parent
143            // or it just shares a scan
144            // or it doesn't share any scan or source.
145            *self
146                .counts
147                .get(&s.share_id())
148                .expect("Unprocessed shared node.")
149                == 1
150                || s.input().as_logical_scan().is_some()
151                || !(plan_visitor::has_logical_scan(s.input())
152                    || plan_visitor::has_logical_source(s.input()))
153        };
154        t.as_logical_share()
155            .filter(prunable)
156            .map_or(t.clone(), |s| self.pre(s.input()))
157    }
158
159    fn apply(&mut self, t: PlanRef) -> PlanRef {
160        self.dag_apply(t)
161    }
162}
risingwave_frontend/optimizer/plan_node/merge_eq_nodes.rs

risingwave_frontend/optimizer/plan_node/
merge_eq_nodes.rs
