risingwave_frontend/optimizer/plan_visitor/
share_parent_counter.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;
16
17use super::{DefaultBehavior, DefaultValue, LogicalPlanVisitor};
18use crate::optimizer::ShareId;
19use crate::optimizer::plan_node::{LogicalShare, PlanTreeNodeUnary, ShareNode};
20use crate::optimizer::plan_visitor::PlanVisitor;
21
22#[derive(Debug, Clone, Default)]
23pub struct ShareParentCounter {
24    /// Share identity to parent number mapping.
25    parent_counter: HashMap<ShareId, usize>,
26}
27
28impl ShareParentCounter {
29    pub fn get_parent_num(&self, share: &LogicalShare) -> usize {
30        self.get_parent_num_by_id(share.share_id())
31    }
32
33    pub fn get_parent_num_by_id(&self, share_id: ShareId) -> usize {
34        *self
35            .parent_counter
36            .get(&share_id)
37            .expect("share must exist")
38    }
39}
40
41impl LogicalPlanVisitor for ShareParentCounter {
42    type Result = ();
43
44    type DefaultBehavior = impl DefaultBehavior<Self::Result>;
45
46    fn default_behavior() -> Self::DefaultBehavior {
47        DefaultValue
48    }
49
50    fn visit_logical_share(&mut self, share: &LogicalShare) {
51        let v = self
52            .parent_counter
53            .entry(share.share_id())
54            .and_modify(|counter| *counter += 1)
55            .or_insert(1);
56        if *v == 1 {
57            self.visit(share.input())
58        }
59    }
60}
risingwave_frontend/optimizer/plan_visitor/share_parent_counter.rs

risingwave_frontend/optimizer/plan_visitor/
share_parent_counter.rs
