risingwave_stream/executor/
values.rs

1// Copyright 2023 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::vec;
16
17use risingwave_common::array::{DataChunk, Op};
18use risingwave_common::ensure;
19use risingwave_common::util::iter_util::ZipEqFast;
20use risingwave_expr::expr::NonStrictExpression;
21use tokio::sync::mpsc::UnboundedReceiver;
22
23use crate::executor::prelude::*;
24use crate::task::CreateMviewProgressReporter;
25
26const DEFAULT_CHUNK_SIZE: usize = 1024;
27
28/// Execute `values` in stream. As is a leaf, current workaround holds a `barrier_executor`.
29/// May refactor with `BarrierRecvExecutor` in the near future.
30pub struct ValuesExecutor {
31    ctx: ActorContextRef,
32
33    schema: Schema,
34    // Receiver of barrier channel.
35    barrier_receiver: UnboundedReceiver<Barrier>,
36    progress: CreateMviewProgressReporter,
37
38    rows: vec::IntoIter<Vec<NonStrictExpression>>,
39}
40
41impl ValuesExecutor {
42    pub fn new(
43        ctx: ActorContextRef,
44        schema: Schema,
45        progress: CreateMviewProgressReporter,
46        rows: Vec<Vec<NonStrictExpression>>,
47        barrier_receiver: UnboundedReceiver<Barrier>,
48    ) -> Self {
49        Self {
50            ctx,
51            schema,
52            progress,
53            barrier_receiver,
54            rows: rows.into_iter(),
55        }
56    }
57
58    #[try_stream(ok = Message, error = StreamExecutorError)]
59    async fn execute_inner(self) {
60        let Self {
61            schema,
62            mut progress,
63            mut barrier_receiver,
64            mut rows,
65            ..
66        } = self;
67        let barrier = barrier_receiver
68            .recv()
69            .instrument_await("values_executor_recv_first_barrier")
70            .await
71            .unwrap();
72
73        let emit = barrier.is_newly_added(self.ctx.id);
74        let paused_on_startup = barrier.is_pause_on_startup();
75
76        yield Message::Barrier(barrier);
77
78        // If it's failover, do not evaluate rows (assume they have been yielded)
79        if emit {
80            if paused_on_startup {
81                // Wait for the data stream to be resumed before yielding the chunks.
82                while let Some(barrier) = barrier_receiver.recv().await {
83                    let is_resume = barrier.is_resume();
84                    yield Message::Barrier(barrier);
85                    if is_resume {
86                        break;
87                    }
88                }
89            }
90
91            let cardinality = schema.len();
92            ensure!(cardinality > 0);
93            while !rows.is_empty() {
94                // We need a one row chunk rather than an empty chunk because constant
95                // expression's eval result is same size as input chunk
96                // cardinality.
97                let one_row_chunk = DataChunk::new_dummy(1);
98
99                let chunk_size = DEFAULT_CHUNK_SIZE.min(rows.len());
100                let mut array_builders = schema.create_array_builders(chunk_size);
101                for row in rows.by_ref().take(chunk_size) {
102                    for (expr, builder) in row.into_iter().zip_eq_fast(&mut array_builders) {
103                        let out = expr.eval_infallible(&one_row_chunk).await;
104                        builder.append_array(&out);
105                    }
106                }
107
108                let columns: Vec<_> = array_builders
109                    .into_iter()
110                    .map(|b| b.finish().into())
111                    .collect();
112
113                let chunk = DataChunk::new(columns, chunk_size);
114                let ops = vec![Op::Insert; chunk_size];
115
116                let stream_chunk = StreamChunk::from_parts(ops, chunk);
117                yield Message::Chunk(stream_chunk);
118            }
119        }
120
121        let mut finish_reported = !emit;
122
123        while let Some(barrier) = barrier_receiver.recv().await {
124            if !finish_reported {
125                progress.finish(barrier.epoch, 0);
126                finish_reported = true;
127            }
128            yield Message::Barrier(barrier);
129        }
130    }
131}
132
133impl Execute for ValuesExecutor {
134    fn execute(self: Box<Self>) -> BoxedMessageStream {
135        self.execute_inner().boxed()
136    }
137}
138
139#[cfg(test)]
140mod tests {
141
142    use futures::StreamExt;
143    use risingwave_common::array::{
144        Array, ArrayImpl, I16Array, I32Array, I64Array, StructArray, StructValue,
145    };
146    use risingwave_common::catalog::{Field, Schema};
147    use risingwave_common::types::{DataType, ScalarImpl, StructType};
148    use risingwave_common::util::epoch::test_epoch;
149    use risingwave_expr::expr::{BoxedExpression, LiteralExpression, NonStrictExpression};
150    use tokio::sync::mpsc::unbounded_channel;
151
152    use super::ValuesExecutor;
153    use crate::executor::test_utils::StreamExecutorTestExt;
154    use crate::executor::{ActorContext, AddMutation, Barrier, Execute, Mutation};
155    use crate::task::CreateMviewProgressReporter;
156    use crate::task::barrier_test_utils::LocalBarrierTestEnv;
157
158    #[tokio::test]
159    async fn test_values() {
160        let test_env = LocalBarrierTestEnv::for_test().await;
161        let barrier_manager = test_env.local_barrier_manager.clone();
162        let progress = CreateMviewProgressReporter::for_test(barrier_manager);
163        let actor_id = progress.actor_id();
164        let (tx, barrier_receiver) = unbounded_channel();
165        let value = StructValue::new(vec![Some(1.into()), Some(2.into()), Some(3.into())]);
166        let exprs: Vec<BoxedExpression> = vec![
167            LiteralExpression::new(DataType::Int16, Some(ScalarImpl::Int16(1))).into(),
168            LiteralExpression::new(DataType::Int32, Some(ScalarImpl::Int32(2))).into(),
169            LiteralExpression::new(DataType::Int64, Some(ScalarImpl::Int64(3))).into(),
170            LiteralExpression::new(
171                StructType::unnamed(vec![DataType::Int32, DataType::Int32, DataType::Int32]).into(),
172                Some(ScalarImpl::Struct(value)),
173            )
174            .into(),
175            LiteralExpression::new(DataType::Int64, Some(ScalarImpl::Int64(0))).into(),
176        ];
177        let schema = exprs
178            .iter() // for each column
179            .map(|col| Field::unnamed(col.return_type()))
180            .collect::<Schema>();
181        let values_executor_struct = ValuesExecutor::new(
182            ActorContext::for_test(actor_id),
183            schema,
184            progress,
185            vec![
186                exprs
187                    .into_iter()
188                    .map(NonStrictExpression::for_test)
189                    .collect(),
190            ],
191            barrier_receiver,
192        );
193        let mut values_executor = Box::new(values_executor_struct).execute();
194
195        // Init barrier
196        let first_message =
197            Barrier::new_test_barrier(test_epoch(1)).with_mutation(Mutation::Add(AddMutation {
198                added_actors: maplit::hashset! {actor_id},
199                ..Default::default()
200            }));
201        tx.send(first_message).unwrap();
202
203        assert!(matches!(
204            values_executor.next_unwrap_ready_barrier().unwrap(),
205            Barrier { .. }
206        ));
207
208        // Consume the barrier
209        let values_msg = values_executor.next().await.unwrap().unwrap();
210
211        let result = values_msg
212            .into_chunk()
213            .unwrap()
214            .compact_vis()
215            .data_chunk()
216            .to_owned();
217
218        let array: ArrayImpl = StructArray::new(
219            StructType::unnamed(vec![DataType::Int32, DataType::Int32, DataType::Int32]),
220            vec![
221                I32Array::from_iter([1]).into_ref(),
222                I32Array::from_iter([2]).into_ref(),
223                I32Array::from_iter([3]).into_ref(),
224            ],
225            [true].into_iter().collect(),
226        )
227        .into();
228
229        assert_eq!(*result.column_at(0), I16Array::from_iter([1]).into_ref());
230        assert_eq!(*result.column_at(1), I32Array::from_iter([2]).into_ref());
231        assert_eq!(*result.column_at(2), I64Array::from_iter([3]).into_ref());
232        assert_eq!(*result.column_at(3), array.into());
233        assert_eq!(*result.column_at(4), I64Array::from_iter([0]).into_ref());
234
235        // ValueExecutor should simply forward following barriers
236        tx.send(Barrier::new_test_barrier(test_epoch(2))).unwrap();
237
238        assert!(matches!(
239            values_executor.next_unwrap_ready_barrier().unwrap(),
240            Barrier { .. }
241        ));
242
243        tx.send(Barrier::new_test_barrier(test_epoch(3))).unwrap();
244
245        assert!(matches!(
246            values_executor.next_unwrap_ready_barrier().unwrap(),
247            Barrier { .. }
248        ));
249    }
250}
risingwave_stream/executor/values.rs

risingwave_stream/executor/
values.rs
