risingwave_batch_executors/executor/
values.rs

1// Copyright 2024 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 futures_async_stream::try_stream;
18use itertools::Itertools;
19use risingwave_common::array::DataChunk;
20use risingwave_common::catalog::{Field, Schema};
21use risingwave_common::util::iter_util::ZipEqFast;
22use risingwave_expr::expr::{BoxedExpression, build_from_prost};
23use risingwave_pb::batch_plan::plan_node::NodeBody;
24
25use crate::error::{BatchError, Result};
26use crate::executor::{
27    BoxedDataChunkStream, BoxedExecutor, BoxedExecutorBuilder, Executor, ExecutorBuilder,
28};
29
30/// [`ValuesExecutor`] implements Values executor.
31pub struct ValuesExecutor {
32    rows: vec::IntoIter<Vec<BoxedExpression>>,
33    schema: Schema,
34    identity: String,
35    chunk_size: usize,
36}
37
38impl ValuesExecutor {
39    #[cfg(test)]
40    pub(crate) fn new(
41        rows: Vec<Vec<BoxedExpression>>,
42        schema: Schema,
43        identity: String,
44        chunk_size: usize,
45    ) -> Self {
46        Self {
47            rows: rows.into_iter(),
48            schema,
49            identity,
50            chunk_size,
51        }
52    }
53}
54
55impl Executor for ValuesExecutor {
56    fn schema(&self) -> &Schema {
57        &self.schema
58    }
59
60    fn identity(&self) -> &str {
61        &self.identity
62    }
63
64    fn execute(self: Box<Self>) -> BoxedDataChunkStream {
65        self.do_execute()
66    }
67}
68
69impl ValuesExecutor {
70    #[try_stream(boxed, ok = DataChunk, error = BatchError)]
71    async fn do_execute(mut self: Box<Self>) {
72        if !self.rows.is_empty() {
73            let cardinality = self.rows.len();
74            ensure!(cardinality > 0);
75
76            while !self.rows.is_empty() {
77                // We need a one row chunk rather than an empty chunk because constant
78                // expression's eval result is same size as input chunk
79                // cardinality.
80                let one_row_chunk = DataChunk::new_dummy(1);
81
82                let chunk_size = self.chunk_size.min(self.rows.len());
83                let mut array_builders = self.schema.create_array_builders(chunk_size);
84                for row in self.rows.by_ref().take(chunk_size) {
85                    for (expr, builder) in row.into_iter().zip_eq_fast(&mut array_builders) {
86                        let out = expr.eval(&one_row_chunk).await?;
87                        builder.append_array(&out);
88                    }
89                }
90
91                let columns: Vec<_> = array_builders
92                    .into_iter()
93                    .map(|b| b.finish().into())
94                    .collect();
95
96                let chunk = DataChunk::new(columns, chunk_size);
97
98                yield chunk
99            }
100        }
101    }
102}
103
104impl BoxedExecutorBuilder for ValuesExecutor {
105    async fn new_boxed_executor(
106        source: &ExecutorBuilder<'_>,
107        inputs: Vec<BoxedExecutor>,
108    ) -> Result<BoxedExecutor> {
109        ensure!(inputs.is_empty(), "ValuesExecutor should have no child!");
110        let value_node = try_match_expand!(
111            source.plan_node().get_node_body().unwrap(),
112            NodeBody::Values
113        )?;
114
115        let mut rows: Vec<Vec<BoxedExpression>> = Vec::with_capacity(value_node.get_tuples().len());
116        for row in value_node.get_tuples() {
117            let expr_row: Vec<_> = row.get_cells().iter().map(build_from_prost).try_collect()?;
118            rows.push(expr_row);
119        }
120
121        let fields = value_node
122            .get_fields()
123            .iter()
124            .map(Field::from)
125            .collect::<Vec<Field>>();
126
127        Ok(Box::new(Self {
128            rows: rows.into_iter(),
129            schema: Schema { fields },
130            identity: source.plan_node().get_identity().clone(),
131            chunk_size: source.context().get_config().developer.chunk_size,
132        }))
133    }
134}
135
136#[cfg(test)]
137mod tests {
138
139    use futures::stream::StreamExt;
140    use risingwave_common::array::{
141        Array, ArrayImpl, I16Array, I32Array, I64Array, StructArray, StructValue,
142    };
143    use risingwave_common::catalog::{Field, Schema};
144    use risingwave_common::types::{DataType, ScalarImpl, StructType};
145    use risingwave_expr::expr::{BoxedExpression, LiteralExpression};
146
147    use crate::executor::{Executor, ValuesExecutor};
148
149    const CHUNK_SIZE: usize = 1024;
150
151    #[tokio::test]
152    async fn test_values_executor() {
153        let value = StructValue::new(vec![Some(1.into()), Some(2.into()), Some(3.into())]);
154        let exprs: Vec<BoxedExpression> = vec![
155            LiteralExpression::new(DataType::Int16, Some(ScalarImpl::Int16(1))).into(),
156            LiteralExpression::new(DataType::Int32, Some(ScalarImpl::Int32(2))).into(),
157            LiteralExpression::new(DataType::Int64, Some(ScalarImpl::Int64(3))).into(),
158            LiteralExpression::new(
159                StructType::unnamed(vec![DataType::Int32, DataType::Int32, DataType::Int32]).into(),
160                Some(ScalarImpl::Struct(value)),
161            )
162            .into(),
163        ];
164
165        let fields = exprs
166            .iter() // for each column
167            .map(|col| Field::unnamed(col.return_type()))
168            .collect::<Vec<Field>>();
169
170        let values_executor = Box::new(ValuesExecutor {
171            rows: vec![exprs].into_iter(),
172            schema: Schema { fields },
173            identity: "ValuesExecutor2".to_owned(),
174            chunk_size: CHUNK_SIZE,
175        });
176
177        let fields = &values_executor.schema().fields;
178        assert_eq!(fields[0].data_type, DataType::Int16);
179        assert_eq!(fields[1].data_type, DataType::Int32);
180        assert_eq!(fields[2].data_type, DataType::Int64);
181        assert_eq!(
182            fields[3].data_type,
183            StructType::unnamed(vec![DataType::Int32, DataType::Int32, DataType::Int32]).into()
184        );
185
186        let mut stream = values_executor.execute();
187        let result = stream.next().await.unwrap();
188        let array: ArrayImpl = StructArray::new(
189            StructType::unnamed(vec![DataType::Int32, DataType::Int32, DataType::Int32]),
190            vec![
191                I32Array::from_iter([1]).into_ref(),
192                I32Array::from_iter([2]).into_ref(),
193                I32Array::from_iter([3]).into_ref(),
194            ],
195            [true].into_iter().collect(),
196        )
197        .into();
198
199        if let Ok(result) = result {
200            assert_eq!(*result.column_at(0), I16Array::from_iter([1]).into_ref());
201            assert_eq!(*result.column_at(1), I32Array::from_iter([2]).into_ref());
202            assert_eq!(*result.column_at(2), I64Array::from_iter([3]).into_ref());
203            assert_eq!(*result.column_at(3), array.into());
204        }
205    }
206
207    #[tokio::test]
208    async fn test_chunk_split_size() {
209        let rows = [
210            LiteralExpression::new(DataType::Int32, Some(ScalarImpl::Int32(1))).into(),
211            LiteralExpression::new(DataType::Int32, Some(ScalarImpl::Int32(2))).into(),
212            LiteralExpression::new(DataType::Int32, Some(ScalarImpl::Int32(3))).into(),
213            LiteralExpression::new(DataType::Int32, Some(ScalarImpl::Int32(4))).into(),
214        ]
215        .into_iter()
216        .map(|expr| vec![expr])
217        .collect::<Vec<_>>();
218
219        let fields = vec![Field::unnamed(DataType::Int32)];
220
221        let values_executor = Box::new(ValuesExecutor::new(
222            rows,
223            Schema { fields },
224            "ValuesExecutor2".to_owned(),
225            3,
226        ));
227        let mut stream = values_executor.execute();
228        assert_eq!(stream.next().await.unwrap().unwrap().cardinality(), 3);
229        assert_eq!(stream.next().await.unwrap().unwrap().cardinality(), 1);
230        assert!(stream.next().await.is_none());
231    }
232
233    // Handle the possible case of ValuesNode([[]])
234    #[tokio::test]
235    async fn test_no_column_values_executor() {
236        let values_executor = Box::new(ValuesExecutor::new(
237            vec![vec![]],
238            Schema::default(),
239            "ValuesExecutor2".to_owned(),
240            CHUNK_SIZE,
241        ));
242        let mut stream = values_executor.execute();
243
244        let result = stream.next().await.unwrap().unwrap();
245        assert_eq!(result.cardinality(), 1);
246        assert_eq!(result.dimension(), 0);
247
248        assert!(stream.next().await.is_none());
249    }
250}
risingwave_batch_executors/executor/values.rs

risingwave_batch_executors/executor/
values.rs
