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        let upstream_fragment_id = node.get_upstream_fragment_id();
39        let upstream_actors = actor_context
40            .initial_upstream_actors
41            .get(&upstream_fragment_id);
42        if upstream_actors.is_none() && !node.allow_empty_upstream {
43            return Ok(None);
44        }
45
46        let inputs: Vec<_> = try_join_all(
47            upstream_actors
48                .into_iter()
49                .flat_map(|upstream_actors| upstream_actors.actors.iter())
50                .map(|upstream_actor| {
51                    new_input(
52                        &local_barrier_manager,
53                        executor_stats.clone(),
54                        actor_context.id,
55                        actor_context.fragment_id,
56                        upstream_actor,
57                        upstream_fragment_id,
58                        actor_context.config.clone(),
59                    )
60                }),
61        )
62        .await?;
63
64        // If there's always only one upstream, we can use `ReceiverExecutor`. Note that it can't
65        // scale to multiple upstreams. An initially empty merge must stay dynamic to accept a
66        // later MergeUpdate, even for dispatcher kinds normally optimized to a singleton.
67        let always_single_input = !node.allow_empty_upstream
68            && !inputs.is_empty()
69            && match node.get_upstream_dispatcher_type()? {
70                DispatcherType::Unspecified => unreachable!(),
71                DispatcherType::Hash | DispatcherType::Broadcast => false,
72                // There could be arbitrary number of upstreams with simple dispatcher.
73                DispatcherType::Simple => false,
74                // There should be always only one upstream with no-shuffle dispatcher.
75                DispatcherType::NoShuffle => true,
76            };
77
78        let upstreams = if always_single_input {
79            MergeExecutorUpstream::Singleton(Itertools::exactly_one(inputs.into_iter()).unwrap())
80        } else {
81            MergeExecutorUpstream::Merge(MergeExecutor::new_merge_upstream(
82                inputs,
83                &executor_stats,
84                &actor_context,
85                chunk_size,
86                info.schema.clone(),
87            ))
88        };
89
90        Ok(Some(MergeExecutorInput::new(
91            upstreams,
92            actor_context,
93            upstream_fragment_id,
94            local_barrier_manager,
95            executor_stats,
96            info,
97        )))
98    }
99}
100
101impl_stream_node_body!(Merge(MergeNode) => MergeExecutorBuilder);
102
103impl ExecutorBuilder for MergeExecutorBuilder {
104    type Node = MergeNode;
105
106    async fn new_boxed_executor(
107        params: ExecutorParams,
108        node: &Self::Node,
109        _store: impl StateStore,
110    ) -> StreamResult<Executor> {
111        let actor_id = params.actor_context.id;
112        let fragment_id = params.actor_context.fragment_id;
113        let barrier_rx = params.local_barrier_manager.subscribe_barrier(actor_id);
114        Ok(Self::new_input(
115            params.local_barrier_manager,
116            params.executor_stats,
117            params.actor_context,
118            params.info,
119            node,
120            params.config.developer.chunk_size,
121        )
122        .await?
123        .ok_or_else(|| {
124            anyhow!(
125                "no upstream actors found for actor {} in fragment {}",
126                actor_id,
127                fragment_id
128            )
129        })?
130        .into_executor(barrier_rx))
131    }
132}
risingwave_stream/from_proto/merge.rs

risingwave_stream/from_proto/
merge.rs
