risingwave_expr_impl/scalar/
coalesce.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::BitAnd;
16use std::sync::Arc;
17
18use risingwave_common::array::{ArrayRef, DataChunk};
19use risingwave_common::row::OwnedRow;
20use risingwave_common::types::{DataType, Datum};
21use risingwave_expr::expr::{
22    AsyncExpression, AsyncExpressionBoxExt, BoxedExpression, ExpressionInfo, SyncExpression,
23    SyncExpressionBoxExt, try_into_sync_exprs,
24};
25use risingwave_expr::{Result, build_function};
26
27#[derive(Debug)]
28pub struct CoalesceExpression<E> {
29    return_type: DataType,
30    children: Vec<E>,
31}
32
33impl<E: ExpressionInfo> ExpressionInfo for CoalesceExpression<E> {
34    fn return_type(&self) -> DataType {
35        self.return_type.clone()
36    }
37}
38
39macro_rules! eval_coalesce {
40    ($mode:ident, $this:expr, $input:expr) => {{
41        let init_vis = $input.visibility();
42        let mut input = $input.clone();
43        let len = input.capacity();
44        let mut selection: Vec<Option<usize>> = vec![None; len];
45        let mut children_array = Vec::with_capacity($this.children.len());
46        for (child_idx, child) in $this.children.iter().enumerate() {
47            let res = risingwave_expr::forward!($mode, child, eval(&input))?;
48            let res_bitmap = res.null_bitmap();
49            let orig_vis = input.visibility();
50            for pos in orig_vis.bitand(res_bitmap).iter_ones() {
51                selection[pos] = Some(child_idx);
52            }
53            let new_vis = orig_vis & !res_bitmap;
54            input.set_visibility(new_vis);
55            children_array.push(res);
56        }
57        let mut builder = $this.return_type.create_array_builder(len);
58        for (i, sel) in selection.iter().enumerate() {
59            if init_vis.is_set(i)
60                && let Some(child_idx) = sel
61            {
62                builder.append(children_array[*child_idx].value_at(i));
63            } else {
64                builder.append_null()
65            }
66        }
67        Ok(Arc::new(builder.finish()))
68    }};
69}
70
71macro_rules! eval_row_coalesce {
72    ($mode:ident, $this:expr, $input:expr) => {{
73        for child in &$this.children {
74            let datum = risingwave_expr::forward!($mode, child, eval_row($input))?;
75            if datum.is_some() {
76                return Ok(datum);
77            }
78        }
79        Ok(None)
80    }};
81}
82
83impl<E: SyncExpression> SyncExpression for CoalesceExpression<E> {
84    fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
85        eval_coalesce!(sync, self, input)
86    }
87
88    fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
89        eval_row_coalesce!(sync, self, input)
90    }
91}
92
93impl<E: AsyncExpression> AsyncExpression for CoalesceExpression<E> {
94    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
95        eval_coalesce!(async, self, input)
96    }
97
98    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
99        eval_row_coalesce!(async, self, input)
100    }
101}
102
103#[build_function("coalesce(...) -> any", type_infer = "unreachable")]
104fn build(return_type: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
105    match try_into_sync_exprs(children) {
106        Ok(children) => Ok(CoalesceExpression {
107            return_type,
108            children,
109        }
110        .boxed()),
111        Err(children) => Ok(CoalesceExpression {
112            return_type,
113            children,
114        }
115        .boxed()),
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use risingwave_common::array::DataChunk;
122    use risingwave_common::row::Row;
123    use risingwave_common::test_prelude::DataChunkTestExt;
124    use risingwave_common::types::ToOwnedDatum;
125    use risingwave_common::util::iter_util::ZipEqDebug;
126    use risingwave_expr::expr::build_from_pretty;
127
128    #[tokio::test]
129    async fn test_coalesce_expr() {
130        let expr = build_from_pretty("(coalesce:int4 $0:int4 $1:int4 $2:int4)");
131        let (input, expected) = DataChunk::from_pretty(
132            "i i i i
133             1 . . 1
134             . 2 . 2
135             . . 3 3
136             . . . .",
137        )
138        .split_column_at(3);
139
140        // test eval
141        let output = expr.eval(&input).await.unwrap();
142        assert_eq!(&output, expected.column_at(0));
143
144        // test eval_row
145        for (row, expected) in input.rows().zip_eq_debug(expected.rows()) {
146            let result = expr.eval_row(&row.to_owned_row()).await.unwrap();
147            assert_eq!(result, expected.datum_at(0).to_owned_datum());
148        }
149    }
150}
risingwave_expr_impl/scalar/coalesce.rs

risingwave_expr_impl/scalar/
coalesce.rs
