risingwave_frontend/optimizer/plan_node/
stream_source.rs

1// Copyright 2022 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::rc::Rc;
16
17use itertools::Itertools;
18use pretty_xmlish::{Pretty, XmlNode};
19use risingwave_common::catalog::{ColumnCatalog, Field};
20use risingwave_common::types::DataType;
21use risingwave_pb::stream_plan::stream_node::PbNodeBody;
22use risingwave_pb::stream_plan::{Columns, PbStreamSource, SourceNode};
23
24use super::stream::prelude::*;
25use super::utils::{Distill, TableCatalogBuilder, childless_record};
26use super::{ExprRewritable, PlanBase, StreamNode, generic};
27use crate::TableCatalog;
28use crate::catalog::source_catalog::SourceCatalog;
29use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
30use crate::optimizer::plan_node::utils::column_names_pretty;
31use crate::optimizer::property::{Distribution, MonotonicityMap, WatermarkColumns};
32use crate::stream_fragmenter::BuildFragmentGraphState;
33
34/// [`StreamSource`] represents a table/connector source at the very beginning of the graph.
35#[derive(Debug, Clone, PartialEq, Eq, Hash)]
36pub struct StreamSource {
37    pub base: PlanBase<Stream>,
38    pub(crate) core: generic::Source,
39    /// Downstream columns are used by list node to know which columns are needed.
40    /// For example, iceberg list node will use this info to plan files to read.
41    pub(crate) downstream_columns: Option<Vec<ColumnCatalog>>,
42}
43
44impl StreamSource {
45    pub fn new(core: generic::Source) -> Self {
46        let base = PlanBase::new_stream_with_core(
47            &core,
48            Distribution::SomeShard,
49            core.stream_kind(),
50            false,
51            WatermarkColumns::new(),
52            MonotonicityMap::new(),
53        );
54        Self {
55            base,
56            core,
57            downstream_columns: None,
58        }
59    }
60
61    pub fn source_catalog(&self) -> Option<Rc<SourceCatalog>> {
62        self.core.catalog.clone()
63    }
64
65    fn infer_internal_table_catalog(&self) -> TableCatalog {
66        if !self.core.is_iceberg_connector() {
67            generic::Source::infer_internal_table_catalog(false)
68        } else {
69            // Iceberg List stores a single value without a primary key.
70            let mut builder = TableCatalogBuilder::default();
71            let field = if self.core.is_iceberg_update_source() {
72                Field {
73                    data_type: DataType::Jsonb,
74                    name: "list_state".to_owned(),
75                }
76            } else {
77                Field {
78                    data_type: DataType::Int64,
79                    name: "last_snapshot".to_owned(),
80                }
81            };
82            builder.add_column(&field);
83            builder.build(vec![], 0)
84        }
85    }
86}
87
88impl_plan_tree_node_for_leaf! { Stream, StreamSource }
89
90impl Distill for StreamSource {
91    fn distill<'a>(&self) -> XmlNode<'a> {
92        let fields = if let Some(catalog) = self.source_catalog() {
93            let src = Pretty::from(catalog.name.clone());
94            let col = column_names_pretty(self.schema());
95            vec![("source", src), ("columns", col)]
96        } else {
97            vec![]
98        };
99        childless_record("StreamSource", fields)
100    }
101}
102
103impl StreamNode for StreamSource {
104    fn to_stream_prost_body(&self, state: &mut BuildFragmentGraphState) -> PbNodeBody {
105        let source_catalog = self.source_catalog();
106
107        let source_inner = source_catalog.map(|source_catalog| {
108            let (with_properties, secret_refs) =
109                source_catalog.with_properties.clone().into_parts();
110            PbStreamSource {
111                source_id: source_catalog.id,
112                source_name: source_catalog.name.clone(),
113                state_table: Some(
114                    self.infer_internal_table_catalog()
115                        .with_id(state.gen_table_id_wrapped())
116                        .to_internal_table_prost(),
117                ),
118                info: Some(source_catalog.info.clone()),
119                row_id_index: self.core.row_id_index.map(|index| index as _),
120                columns: self
121                    .core
122                    .column_catalog
123                    .iter()
124                    .map(|c| c.to_protobuf())
125                    .collect_vec(),
126                with_properties,
127                rate_limit: source_catalog.rate_limit,
128                secret_refs,
129                downstream_columns: self.downstream_columns.as_ref().map(|cols| Columns {
130                    columns: cols.iter().map(|c| c.to_protobuf()).collect_vec(),
131                }),
132                refresh_mode: source_catalog.refresh_mode,
133                associated_table_id: None, // fill the actual associated table id in `BuildingFragment::fill_job`
134                iceberg_fetch_state_table: None,
135            }
136        });
137        PbNodeBody::Source(Box::new(SourceNode { source_inner }))
138    }
139}
140
141impl ExprRewritable<Stream> for StreamSource {}
142
143impl ExprVisitable for StreamSource {}
risingwave_frontend/optimizer/plan_node/stream_source.rs

risingwave_frontend/optimizer/plan_node/
stream_source.rs
