risingwave_frontend/optimizer/plan_node/
stream_share.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 pretty_xmlish::XmlNode;
16use risingwave_pb::stream_plan::PbStreamNode;
17use risingwave_pb::stream_plan::stream_node::PbNodeBody;
18
19use super::stream::prelude::*;
20use super::utils::Distill;
21use super::{
22    ExprRewritable, PlanTreeNodeUnary, ShareNode, StreamExchange, StreamNode,
23    StreamPlanRef as PlanRef, generic,
24};
25use crate::Explain;
26use crate::optimizer::ShareId;
27use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
28use crate::optimizer::plan_node::generic::Share;
29use crate::optimizer::plan_node::{LogicalShare, PlanBase, PlanTreeNode};
30use crate::scheduler::SchedulerResult;
31use crate::stream_fragmenter::BuildFragmentGraphState;
32
33/// `StreamShare` will be translated into an `ExchangeNode` based on its distribution finally.
34#[derive(Debug, Clone, PartialEq, Eq, Hash)]
35pub struct StreamShare {
36    pub base: PlanBase<Stream>,
37    core: generic::Share<PlanRef>,
38}
39
40impl StreamShare {
41    pub fn new(core: generic::Share<PlanRef>) -> Self {
42        let input = core.input();
43        let dist = input.distribution().clone();
44        // Filter executor won't change the append-only behavior of the stream.
45        let base = PlanBase::new_stream_share(
46            &core,
47            dist,
48            input.stream_kind(),
49            input.emit_on_window_close(),
50            input.watermark_columns().clone(),
51            input.columns_monotonicity().clone(),
52        );
53
54        StreamShare { base, core }
55    }
56
57    pub fn new_from_input(input: PlanRef) -> Self {
58        let ctx = input.ctx();
59        Self::new(ctx.register_stream_share(input))
60    }
61}
62
63impl Distill for StreamShare {
64    fn distill<'a>(&self) -> XmlNode<'a> {
65        LogicalShare::pretty_fields(&self.base, "StreamShare")
66    }
67}
68
69impl PlanTreeNodeUnary<Stream> for StreamShare {
70    fn input(&self) -> PlanRef {
71        self.core.input()
72    }
73
74    fn clone_with_input(&self, _input: PlanRef) -> Self {
75        unreachable!("shared node should be handled specially in PlanRef::clone_with_input")
76    }
77}
78
79impl ShareNode<Stream> for StreamShare {
80    fn share_id(&self) -> ShareId {
81        self.core.share_id()
82    }
83
84    fn new_share(core: Share<PlanRef>) -> PlanRef {
85        Self::new(core).into()
86    }
87
88    fn replace_input(&self, plan: PlanRef) -> PlanRef {
89        debug_assert!(
90            self.schema().type_eq(plan.schema()),
91            "replacing a stream share input must preserve its schema"
92        );
93        self.ctx()
94            .update_stream_share(self.share_id(), plan.clone());
95        Self::new(self.core.with_input(plan)).into()
96    }
97
98    fn fork_with_input(&self, plan: PlanRef) -> PlanRef {
99        Self::new_from_input(plan).into()
100    }
101}
102
103impl_plan_tree_node_for_unary! { Stream, StreamShare }
104
105impl StreamNode for StreamShare {
106    fn to_stream_prost_body(&self, _state: &mut BuildFragmentGraphState) -> PbNodeBody {
107        unreachable!(
108            "stream scan cannot be converted into a prost body -- call `adhoc_to_stream_prost` instead."
109        )
110    }
111}
112
113impl StreamShare {
114    pub fn adhoc_to_stream_prost(
115        &self,
116        state: &mut BuildFragmentGraphState,
117    ) -> SchedulerResult<PbStreamNode> {
118        let operator_id = self.base.id().to_stream_node_operator_id();
119
120        match state.get_share_stream_node(operator_id) {
121            None => {
122                let node_body =
123                    StreamExchange::new_no_shuffle(self.input()).to_stream_prost_body(state);
124
125                let input = self
126                    .inputs()
127                    .into_iter()
128                    .map(|plan| plan.to_stream_prost(state))
129                    .try_collect()?;
130
131                let stream_node = PbStreamNode {
132                    input,
133                    identity: self.distill_to_string(),
134                    node_body: Some(node_body),
135                    operator_id: self.id().to_stream_node_operator_id(),
136                    stream_key: self
137                        .stream_key()
138                        .unwrap_or_else(|| panic!("should always have a stream key in the stream plan but not, sub plan: {}",
139                       PlanRef::from(self.clone()).explain_to_string()))
140                        .iter()
141                        .map(|x| *x as u32)
142                        .collect(),
143                    fields: self.schema().to_prost(),
144                    stream_kind: self.stream_kind().to_protobuf() as i32,
145                };
146
147                state.add_share_stream_node(operator_id, stream_node.clone());
148                Ok(stream_node)
149            }
150
151            Some(stream_node) => Ok(stream_node.clone()),
152        }
153    }
154}
155
156impl ExprRewritable<Stream> for StreamShare {}
157
158impl ExprVisitable for StreamShare {}
risingwave_frontend/optimizer/plan_node/stream_share.rs

risingwave_frontend/optimizer/plan_node/
stream_share.rs
