risingwave_expr_impl/scalar/
field.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 anyhow::anyhow;
16use risingwave_common::array::{ArrayImpl, ArrayRef, DataChunk};
17use risingwave_common::row::OwnedRow;
18use risingwave_common::types::{DataType, Datum, ScalarImpl};
19use risingwave_expr::expr::{
20    AsyncExpression, AsyncExpressionBoxExt, BoxedExpression, ExpressionInfo, SyncExpression,
21    SyncExpressionBoxExt,
22};
23use risingwave_expr::{Result, build_function};
24
25/// `FieldExpression` access a field from a struct.
26#[derive(Debug)]
27pub struct FieldExpression<E> {
28    return_type: DataType,
29    input: E,
30    index: usize,
31}
32
33impl<E: ExpressionInfo> ExpressionInfo for FieldExpression<E> {
34    fn return_type(&self) -> DataType {
35        self.return_type.clone()
36    }
37}
38
39macro_rules! eval_field {
40    ($mode:ident, $this:expr, $input:expr) => {{
41        let array = risingwave_expr::forward!($mode, $this.input, eval($input))?;
42        if let ArrayImpl::Struct(struct_array) = array.as_ref() {
43            Ok(struct_array.field_at($this.index).clone())
44        } else {
45            Err(anyhow!("expects a struct array ref").into())
46        }
47    }};
48}
49
50macro_rules! eval_row_field {
51    ($mode:ident, $this:expr, $input:expr) => {{
52        let struct_datum = risingwave_expr::forward!($mode, $this.input, eval_row($input))?;
53        struct_datum
54            .map(|s| match s {
55                ScalarImpl::Struct(v) => Ok(v.fields()[$this.index].clone()),
56                _ => Err(anyhow!("expects a struct array ref").into()),
57            })
58            .transpose()
59            .map(|x| x.flatten())
60    }};
61}
62
63impl<E: SyncExpression> SyncExpression for FieldExpression<E> {
64    fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
65        eval_field!(sync, self, input)
66    }
67
68    fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
69        eval_row_field!(sync, self, input)
70    }
71}
72
73impl<E: AsyncExpression> AsyncExpression for FieldExpression<E> {
74    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
75        eval_field!(async, self, input)
76    }
77
78    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
79        eval_row_field!(async, self, input)
80    }
81}
82
83#[build_function("field(struct, int4) -> any", type_infer = "unreachable")]
84fn build(return_type: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
85    // Field `func_call_node` have 2 child nodes, the first is Field `FuncCall` or
86    // `InputRef`, the second is i32 `Literal`.
87    let [input, index]: [_; 2] = children.try_into().unwrap();
88    let index = index.eval_const()?.unwrap().into_int32() as usize;
89    Ok(match input {
90        BoxedExpression::Sync(input) => FieldExpression {
91            return_type,
92            input,
93            index,
94        }
95        .boxed(),
96        input @ BoxedExpression::Async(_) => FieldExpression {
97            return_type,
98            input,
99            index,
100        }
101        .boxed(),
102    })
103}
104
105#[cfg(test)]
106mod tests {
107    use risingwave_common::array::{DataChunk, DataChunkTestExt};
108    use risingwave_common::row::Row;
109    use risingwave_common::types::ToOwnedDatum;
110    use risingwave_common::util::iter_util::ZipEqDebug;
111    use risingwave_expr::expr::build_from_pretty;
112
113    #[tokio::test]
114    async fn test_field_expr() {
115        let expr = build_from_pretty("(field:int4 $0:struct<a_int4,b_float4> 0:int4)");
116        let (input, expected) = DataChunk::from_pretty(
117            "<i,f>   i
118             (1,2.0) 1
119             (2,2.0) 2
120             (3,2.0) 3",
121        )
122        .split_column_at(1);
123
124        // test eval
125        let output = expr.eval(&input).await.unwrap();
126        assert_eq!(&output, expected.column_at(0));
127
128        // test eval_row
129        for (row, expected) in input.rows().zip_eq_debug(expected.rows()) {
130            let result = expr.eval_row(&row.to_owned_row()).await.unwrap();
131            assert_eq!(result, expected.datum_at(0).to_owned_datum());
132        }
133    }
134}
risingwave_expr_impl/scalar/field.rs

risingwave_expr_impl/scalar/
field.rs
