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

risingwave_frontend/optimizer/plan_visitor/
share_parent_counter.rs
