risingwave_stream/from_proto/
merge.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::sync::Arc;
16
17use anyhow::anyhow;
18use futures::future::try_join_all;
19use risingwave_pb::stream_plan::{DispatcherType, MergeNode};
20
21use super::*;
22use crate::executor::exchange::input::new_input;
23use crate::executor::monitor::StreamingMetrics;
24use crate::executor::{ActorContextRef, MergeExecutor, MergeExecutorInput, MergeExecutorUpstream};
25use crate::task::LocalBarrierManager;
26
27pub struct MergeExecutorBuilder;
28
29impl MergeExecutorBuilder {
30    pub(crate) async fn new_input(
31        local_barrier_manager: LocalBarrierManager,
32        executor_stats: Arc<StreamingMetrics>,
33        actor_context: ActorContextRef,
34        info: ExecutorInfo,
35        node: &MergeNode,
36        chunk_size: usize,
37    ) -> StreamResult<Option<MergeExecutorInput>> {
38        Self::new_input_inner(
39            local_barrier_manager,
40            executor_stats,
41            actor_context,
42            info,
43            node,
44            chunk_size,
45            false,
46        )
47        .await
48    }
49
50    /// Creates a merge input that may start empty and can be updated in place later.
51    pub(crate) async fn new_input_allow_empty_upstream(
52        local_barrier_manager: LocalBarrierManager,
53        executor_stats: Arc<StreamingMetrics>,
54        actor_context: ActorContextRef,
55        info: ExecutorInfo,
56        node: &MergeNode,
57        chunk_size: usize,
58    ) -> StreamResult<MergeExecutorInput> {
59        Ok(Self::new_input_inner(
60            local_barrier_manager,
61            executor_stats,
62            actor_context,
63            info,
64            node,
65            chunk_size,
66            true,
67        )
68        .await?
69        .expect("allowing an empty upstream always creates a merge input"))
70    }
71
72    async fn new_input_inner(
73        local_barrier_manager: LocalBarrierManager,
74        executor_stats: Arc<StreamingMetrics>,
75        actor_context: ActorContextRef,
76        info: ExecutorInfo,
77        node: &MergeNode,
78        chunk_size: usize,
79        allow_empty_upstream: bool,
80    ) -> StreamResult<Option<MergeExecutorInput>> {
81        let upstream_fragment_id = node.get_upstream_fragment_id();
82        let upstream_actors = actor_context
83            .initial_upstream_actors
84            .get(&upstream_fragment_id);
85        if upstream_actors.is_none() && !allow_empty_upstream {
86            return Ok(None);
87        }
88
89        let inputs: Vec<_> = try_join_all(
90            upstream_actors
91                .into_iter()
92                .flat_map(|upstream_actors| upstream_actors.actors.iter())
93                .map(|upstream_actor| {
94                    new_input(
95                        &local_barrier_manager,
96                        executor_stats.clone(),
97                        actor_context.id,
98                        actor_context.fragment_id,
99                        upstream_actor,
100                        upstream_fragment_id,
101                        actor_context.config.clone(),
102                    )
103                }),
104        )
105        .await?;
106
107        // If there's always only one upstream, we can use `ReceiverExecutor`. Note that it can't
108        // scale to multiple upstreams. A merge that allows an empty upstream must not use the
109        // singleton optimization because it may transition to or from an empty input set.
110        let always_single_input = !allow_empty_upstream
111            && !inputs.is_empty()
112            && match node.get_upstream_dispatcher_type()? {
113                DispatcherType::Unspecified => unreachable!(),
114                DispatcherType::Hash | DispatcherType::Broadcast => false,
115                // There could be arbitrary number of upstreams with simple dispatcher.
116                DispatcherType::Simple => false,
117                // There should be always only one upstream with no-shuffle dispatcher.
118                DispatcherType::NoShuffle => true,
119            };
120
121        let upstreams = if always_single_input {
122            MergeExecutorUpstream::Singleton(Itertools::exactly_one(inputs.into_iter()).unwrap())
123        } else {
124            MergeExecutorUpstream::Merge(MergeExecutor::new_merge_upstream(
125                inputs,
126                &executor_stats,
127                &actor_context,
128                chunk_size,
129                info.schema.clone(),
130            ))
131        };
132
133        Ok(Some(MergeExecutorInput::new(
134            upstreams,
135            actor_context,
136            upstream_fragment_id,
137            local_barrier_manager,
138            executor_stats,
139            info,
140        )))
141    }
142}
143
144impl_stream_node_body!(Merge(MergeNode) => MergeExecutorBuilder);
145
146impl ExecutorBuilder for MergeExecutorBuilder {
147    type Node = MergeNode;
148
149    async fn new_boxed_executor(
150        params: ExecutorParams,
151        node: &Self::Node,
152        _store: impl StateStore,
153    ) -> StreamResult<Executor> {
154        let actor_id = params.actor_context.id;
155        let fragment_id = params.actor_context.fragment_id;
156        Ok(Self::new_input(
157            params.local_barrier_manager,
158            params.executor_stats,
159            params.actor_context,
160            params.info,
161            node,
162            params.config.developer.chunk_size,
163        )
164        .await?
165        .ok_or_else(|| {
166            anyhow!(
167                "no upstream actors found for actor {} in fragment {}",
168                actor_id,
169                fragment_id
170            )
171        })?
172        .into_executor())
173    }
174}
risingwave_stream/from_proto/merge.rs

risingwave_stream/from_proto/
merge.rs
