risingwave_frontend/optimizer/plan_visitor/
relation_collector_visitor.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::HashSet;
16
17use risingwave_common::catalog::TableId;
18
19use super::{DefaultBehavior, DefaultValue};
20use crate::PlanRef;
21use crate::optimizer::plan_node::{BatchSource, LogicalScan, StreamSource, StreamTableScan};
22use crate::optimizer::plan_visitor::PlanVisitor;
23
24/// TODO(rc): maybe we should rename this to `DependencyCollectorVisitor`.
25#[derive(Debug, Clone, Default)]
26pub struct RelationCollectorVisitor {
27    relations: HashSet<TableId>,
28}
29
30impl RelationCollectorVisitor {
31    fn new_with(relations: HashSet<TableId>) -> Self {
32        Self { relations }
33    }
34
35    /// `collect_with` will collect all the relations in the plan with some default ones, which are
36    /// collected during the binding phase. Note that during visit the collected relations might be
37    /// duplicated with the default ones. The collection is necessary, because implicit dependencies
38    /// on indices can only be discovered after plan is built.
39    pub fn collect_with(relations: HashSet<TableId>, plan: PlanRef) -> HashSet<TableId> {
40        let mut visitor = Self::new_with(relations);
41        visitor.visit(plan);
42        visitor.relations
43    }
44}
45
46impl PlanVisitor for RelationCollectorVisitor {
47    type Result = ();
48
49    type DefaultBehavior = impl DefaultBehavior<Self::Result>;
50
51    fn default_behavior() -> Self::DefaultBehavior {
52        DefaultValue
53    }
54
55    fn visit_batch_seq_scan(&mut self, plan: &crate::optimizer::plan_node::BatchSeqScan) {
56        self.relations.insert(plan.core().table_desc.table_id);
57    }
58
59    fn visit_logical_scan(&mut self, plan: &LogicalScan) {
60        self.relations.insert(plan.table_desc().table_id);
61    }
62
63    fn visit_stream_table_scan(&mut self, plan: &StreamTableScan) {
64        let logical = plan.core();
65        self.relations.insert(logical.table_desc.table_id);
66    }
67
68    fn visit_batch_source(&mut self, plan: &BatchSource) {
69        if let Some(catalog) = plan.source_catalog() {
70            self.relations.insert(catalog.id.into());
71        }
72    }
73
74    fn visit_stream_source(&mut self, plan: &StreamSource) {
75        if let Some(catalog) = plan.source_catalog() {
76            self.relations.insert(catalog.id.into());
77        }
78    }
79}
risingwave_frontend/optimizer/plan_visitor/relation_collector_visitor.rs

risingwave_frontend/optimizer/plan_visitor/
relation_collector_visitor.rs
