risingwave_expr/expr/
expr_input_ref.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::ops::Index;
16
17use risingwave_common::array::{ArrayRef, DataChunk};
18use risingwave_common::row::OwnedRow;
19use risingwave_common::types::{DataType, Datum};
20use risingwave_pb::expr::ExprNode;
21
22use super::{BoxedExpression, Build};
23use crate::Result;
24use crate::expr::{ExpressionInfo, SyncExpression};
25
26/// A reference to a column in input relation.
27#[derive(Debug, Clone)]
28pub struct InputRefExpression {
29    return_type: DataType,
30    idx: usize,
31}
32
33impl ExpressionInfo for InputRefExpression {
34    fn return_type(&self) -> DataType {
35        self.return_type.clone()
36    }
37
38    fn input_ref_index(&self) -> Option<usize> {
39        Some(self.idx)
40    }
41}
42
43impl SyncExpression for InputRefExpression {
44    fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
45        Ok(input.column_at(self.idx).clone())
46    }
47
48    fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
49        let cell = input.index(self.idx).as_ref().cloned();
50        Ok(cell)
51    }
52}
53
54impl InputRefExpression {
55    pub fn new(return_type: DataType, idx: usize) -> Self {
56        InputRefExpression { return_type, idx }
57    }
58
59    /// Create an [`InputRefExpression`] from a protobuf expression.
60    ///
61    /// Panics if the protobuf expression is not an input reference.
62    pub fn from_prost(prost: &ExprNode) -> Self {
63        let ret_type = DataType::from(prost.get_return_type().unwrap());
64        let input_col_idx = prost.get_rex_node().unwrap().as_input_ref().unwrap();
65
66        Self {
67            return_type: ret_type,
68            idx: *input_col_idx as _,
69        }
70    }
71
72    pub fn index(&self) -> usize {
73        self.idx
74    }
75
76    pub fn eval_immut(&self, input: &DataChunk) -> Result<ArrayRef> {
77        Ok(input.column_at(self.idx).clone())
78    }
79}
80
81impl Build for InputRefExpression {
82    fn build(
83        prost: &ExprNode,
84        _build_child: impl Fn(&ExprNode) -> Result<BoxedExpression>,
85    ) -> Result<Self> {
86        Ok(Self::from_prost(prost))
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use risingwave_common::row::OwnedRow;
93    use risingwave_common::types::{DataType, Datum};
94
95    use crate::expr::{InputRefExpression, SyncExpression};
96
97    #[tokio::test]
98    async fn test_eval_row_input_ref() {
99        let datums: Vec<Datum> = vec![Some(1.into()), Some(2.into()), None];
100        let input_row = OwnedRow::new(datums.clone());
101
102        for (i, expected) in datums.iter().enumerate() {
103            let expr = InputRefExpression::new(DataType::Int32, i);
104            let result = expr.eval_row(&input_row).unwrap();
105            assert_eq!(*expected, result);
106        }
107    }
108}
risingwave_expr/expr/expr_input_ref.rs

risingwave_expr/expr/
expr_input_ref.rs
