risingwave_expr_impl/scalar/
array_transform.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::sync::Arc;
16
17use risingwave_common::array::{ArrayRef, DataChunk};
18use risingwave_common::row::OwnedRow;
19use risingwave_common::types::{DataType, Datum, ListValue, ScalarImpl};
20use risingwave_expr::expr::{
21    AsyncExpression, AsyncExpressionBoxExt, BoxedExpression, ExpressionInfo, SyncExpression,
22    SyncExpressionBoxExt,
23};
24use risingwave_expr::{Result, build_function};
25
26#[derive(Debug)]
27struct ArrayTransformExpression<E> {
28    array: E,
29    lambda: E,
30}
31
32impl<E: ExpressionInfo> ExpressionInfo for ArrayTransformExpression<E> {
33    fn return_type(&self) -> DataType {
34        DataType::list(self.lambda.return_type())
35    }
36}
37
38impl<E: SyncExpression> SyncExpression for ArrayTransformExpression<E> {
39    fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
40        let lambda_input = self.array.eval(input)?;
41        let lambda_input = Arc::unwrap_or_clone(lambda_input).into_list();
42        let new_list = lambda_input.map_inner_sync(|flatten_input| {
43            let flatten_len = flatten_input.len();
44            let chunk = DataChunk::new(vec![Arc::new(flatten_input)], flatten_len);
45            self.lambda.eval(&chunk).map(Arc::unwrap_or_clone)
46        })?;
47        Ok(Arc::new(new_list.into()))
48    }
49
50    fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
51        let lambda_input = self.array.eval_row(input)?;
52        let lambda_input = lambda_input.map(ScalarImpl::into_list);
53        if let Some(lambda_input) = lambda_input {
54            let len = lambda_input.len();
55            let chunk = DataChunk::new(vec![Arc::new(lambda_input.into_array())], len);
56            let new_vals = self.lambda.eval(&chunk)?;
57            let new_list = ListValue::new(Arc::unwrap_or_clone(new_vals));
58            Ok(Some(new_list.into()))
59        } else {
60            Ok(None)
61        }
62    }
63}
64
65impl<E: AsyncExpression> AsyncExpression for ArrayTransformExpression<E> {
66    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
67        let lambda_input = self.array.eval(input).await?;
68        let lambda_input = Arc::unwrap_or_clone(lambda_input).into_list();
69        let new_list = lambda_input
70            .map_inner(|flatten_input| async move {
71                let flatten_len = flatten_input.len();
72                let chunk = DataChunk::new(vec![Arc::new(flatten_input)], flatten_len);
73                self.lambda.eval(&chunk).await.map(Arc::unwrap_or_clone)
74            })
75            .await?;
76        Ok(Arc::new(new_list.into()))
77    }
78
79    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
80        let lambda_input = self.array.eval_row(input).await?;
81        let lambda_input = lambda_input.map(ScalarImpl::into_list);
82        if let Some(lambda_input) = lambda_input {
83            let len = lambda_input.len();
84            let chunk = DataChunk::new(vec![Arc::new(lambda_input.into_array())], len);
85            let new_vals = self.lambda.eval(&chunk).await?;
86            let new_list = ListValue::new(Arc::unwrap_or_clone(new_vals));
87            Ok(Some(new_list.into()))
88        } else {
89            Ok(None)
90        }
91    }
92}
93
94#[build_function("array_transform(anyarray, any) -> anyarray")]
95fn build(_: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
96    let [array, lambda] = <[BoxedExpression; 2]>::try_from(children).unwrap();
97    Ok(match (array, lambda) {
98        (BoxedExpression::Sync(array), BoxedExpression::Sync(lambda)) => {
99            ArrayTransformExpression { array, lambda }.boxed()
100        }
101        (array, lambda) => ArrayTransformExpression { array, lambda }.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_array_transform() {
115        let expr =
116            build_from_pretty("(array_transform:int4[] $0:int4[] (multiply:int4 $0:int4 2:int4))");
117        let (input, expected) = DataChunk::from_pretty(
118            "i[]     i[]
119             {1,2,3} {2,4,6}",
120        )
121        .split_column_at(1);
122
123        // test eval
124        let output = expr.eval(&input).await.unwrap();
125        assert_eq!(&output, expected.column_at(0));
126
127        // test eval_row
128        for (row, expected) in input.rows().zip_eq_debug(expected.rows()) {
129            let result = expr.eval_row(&row.to_owned_row()).await.unwrap();
130            assert_eq!(result, expected.datum_at(0).to_owned_datum());
131        }
132    }
133}
risingwave_expr_impl/scalar/array_transform.rs

risingwave_expr_impl/scalar/
array_transform.rs
