risingwave_batch_executors/executor/
project.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 futures::{Stream, StreamExt};
16use itertools::Itertools;
17use risingwave_common::array::DataChunk;
18use risingwave_common::catalog::{Field, Schema};
19use risingwave_expr::expr::{BoxedExpression, build_batch_expr_from_prost};
20use risingwave_pb::batch_plan::plan_node::NodeBody;
21
22use crate::error::{BatchError, Result};
23use crate::executor::{
24    BoxedDataChunkStream, BoxedExecutor, BoxedExecutorBuilder, Executor, ExecutorBuilder,
25};
26
27pub struct ProjectExecutor {
28    expr: Vec<BoxedExpression>,
29    child: BoxedExecutor,
30    schema: Schema,
31    identity: String,
32}
33
34impl Executor for ProjectExecutor {
35    fn schema(&self) -> &Schema {
36        &self.schema
37    }
38
39    fn identity(&self) -> &str {
40        &self.identity
41    }
42
43    fn execute(self: Box<Self>) -> BoxedDataChunkStream {
44        (*self).do_execute().boxed()
45    }
46}
47
48impl ProjectExecutor {
49    fn do_execute(self) -> impl Stream<Item = Result<DataChunk>> + 'static {
50        let Self { expr, child, .. } = self;
51        child
52            .execute()
53            .map(move |data_chunk| {
54                let expr = expr.clone();
55                async move {
56                    let data_chunk = data_chunk?;
57                    let arrays = {
58                        let expr_futs = expr.iter().map(|expr| expr.eval(&data_chunk));
59                        futures::future::join_all(expr_futs)
60                            .await
61                            .into_iter()
62                            .try_collect()?
63                    };
64                    let (_, vis) = data_chunk.into_parts();
65                    Ok::<_, BatchError>(DataChunk::new(arrays, vis))
66                }
67            })
68            .buffered(16)
69    }
70}
71
72impl BoxedExecutorBuilder for ProjectExecutor {
73    async fn new_boxed_executor(
74        source: &ExecutorBuilder<'_>,
75        inputs: Vec<BoxedExecutor>,
76    ) -> Result<BoxedExecutor> {
77        let [child]: [_; 1] = inputs.try_into().unwrap();
78
79        let project_node = try_match_expand!(
80            source.plan_node().get_node_body().unwrap(),
81            NodeBody::Project
82        )?;
83
84        let project_exprs: Vec<_> = project_node
85            .get_select_list()
86            .iter()
87            .map(build_batch_expr_from_prost)
88            .try_collect()?;
89
90        let fields = project_exprs
91            .iter()
92            .map(|expr| Field::unnamed(expr.return_type()))
93            .collect::<Vec<Field>>();
94
95        Ok(Box::new(Self {
96            expr: project_exprs,
97            child,
98            schema: Schema { fields },
99            identity: source.plan_node().get_identity().clone(),
100        }))
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use risingwave_common::array::{Array, I32Array};
107    use risingwave_common::test_prelude::*;
108    use risingwave_common::types::DataType;
109    use risingwave_expr::expr::{InputRefExpression, LiteralExpression};
110
111    use super::*;
112    use crate::executor::ValuesExecutor;
113    use crate::executor::test_utils::MockExecutor;
114    use crate::*;
115
116    const CHUNK_SIZE: usize = 1024;
117
118    #[tokio::test]
119    async fn test_project_executor() -> Result<()> {
120        let chunk = DataChunk::from_pretty(
121            "
122            i     i
123            1     7
124            2     8
125            33333 66666
126            4     4
127            5     3
128        ",
129        );
130
131        let expr1 = InputRefExpression::new(DataType::Int32, 0);
132        let expr_vec: Vec<BoxedExpression> = vec![expr1.into()];
133
134        let schema = schema_unnamed! { DataType::Int32, DataType::Int32 };
135        let mut mock_executor = MockExecutor::new(schema);
136        mock_executor.add(chunk);
137
138        let fields = expr_vec
139            .iter()
140            .map(|expr| Field::unnamed(expr.return_type()))
141            .collect::<Vec<Field>>();
142
143        let proj_executor = Box::new(ProjectExecutor {
144            expr: expr_vec,
145            child: Box::new(mock_executor),
146            schema: Schema { fields },
147            identity: "ProjectExecutor".to_owned(),
148        });
149
150        let fields = &proj_executor.schema().fields;
151        assert_eq!(fields[0].data_type, DataType::Int32);
152
153        let mut stream = proj_executor.execute();
154        let result_chunk = stream.next().await.unwrap().unwrap();
155        assert_eq!(result_chunk.dimension(), 1);
156        assert_eq!(
157            result_chunk
158                .column_at(0)
159                .as_int32()
160                .iter()
161                .collect::<Vec<_>>(),
162            vec![Some(1), Some(2), Some(33333), Some(4), Some(5)]
163        );
164        Ok(())
165    }
166
167    #[tokio::test]
168    async fn test_project_dummy_chunk() {
169        let literal = LiteralExpression::new(DataType::Int32, Some(1_i32.into()));
170
171        let values_executor2: Box<dyn Executor> = Box::new(ValuesExecutor::new(
172            vec![vec![]], // One single row with no column.
173            Schema::default(),
174            "ValuesExecutor".to_owned(),
175            CHUNK_SIZE,
176        ));
177
178        let proj_executor = Box::new(ProjectExecutor {
179            expr: vec![literal.into()],
180            child: values_executor2,
181            schema: schema_unnamed!(DataType::Int32),
182            identity: "ProjectExecutor2".to_owned(),
183        });
184        let mut stream = proj_executor.execute();
185        let chunk = stream.next().await.unwrap().unwrap();
186        assert_eq!(*chunk.column_at(0), I32Array::from_iter([1]).into_ref());
187    }
188}
risingwave_batch_executors/executor/project.rs

risingwave_batch_executors/executor/
project.rs
