risingwave_expr_impl/scalar/
in_.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::collections::HashSet;
16use std::fmt::Debug;
17use std::sync::Arc;
18
19use risingwave_common::array::{ArrayBuilder, ArrayRef, BoolArrayBuilder, DataChunk};
20use risingwave_common::bail;
21use risingwave_common::row::OwnedRow;
22use risingwave_common::types::{DataType, Datum, Scalar, ToOwnedDatum};
23use risingwave_common::util::iter_util::ZipEqFast;
24use risingwave_expr::expr::{
25    AsyncExpression, AsyncExpressionBoxExt, BoxedExpression, ExpressionInfo, SyncExpression,
26    SyncExpressionBoxExt,
27};
28use risingwave_expr::{Result, build_function};
29
30#[derive(Debug)]
31pub struct InExpression<E> {
32    left: E,
33    set: HashSet<Datum>,
34    return_type: DataType,
35}
36
37impl<E> InExpression<E> {
38    pub fn new(left: E, data: impl Iterator<Item = Datum>, return_type: DataType) -> Self {
39        Self {
40            left,
41            set: data.collect(),
42            return_type,
43        }
44    }
45
46    // Returns true if datum exists in set, null if datum is null or datum does not exist in set
47    // but null does, and false if neither datum nor null exists in set.
48    fn exists(&self, datum: &Datum) -> Option<bool> {
49        if datum.is_none() {
50            None
51        } else if self.set.contains(datum) {
52            Some(true)
53        } else if self.set.contains(&None) {
54            None
55        } else {
56            Some(false)
57        }
58    }
59}
60
61impl<E: ExpressionInfo> ExpressionInfo for InExpression<E> {
62    fn return_type(&self) -> DataType {
63        self.return_type.clone()
64    }
65}
66
67macro_rules! eval_in {
68    ($mode:ident, $this:expr, $input:expr) => {{
69        let input_array = risingwave_expr::forward!($mode, $this.left, eval($input))?;
70        let mut output_array = BoolArrayBuilder::new(input_array.len());
71        for (data, vis) in input_array.iter().zip_eq_fast($input.visibility().iter()) {
72            if vis {
73                // TODO: avoid `to_owned_datum()`
74                let ret = $this.exists(&data.to_owned_datum());
75                output_array.append(ret);
76            } else {
77                output_array.append(None);
78            }
79        }
80        Ok(Arc::new(output_array.finish().into()))
81    }};
82}
83
84macro_rules! eval_row_in {
85    ($mode:ident, $this:expr, $input:expr) => {{
86        let data = risingwave_expr::forward!($mode, $this.left, eval_row($input))?;
87        let ret = $this.exists(&data);
88        Ok(ret.map(|b| b.to_scalar_value()))
89    }};
90}
91
92impl<E: SyncExpression> SyncExpression for InExpression<E> {
93    fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
94        eval_in!(sync, self, input)
95    }
96
97    fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
98        eval_row_in!(sync, self, input)
99    }
100}
101
102impl<E: AsyncExpression> AsyncExpression for InExpression<E> {
103    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
104        eval_in!(async, self, input)
105    }
106
107    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
108        eval_row_in!(async, self, input)
109    }
110}
111
112#[build_function("in(any, ...) -> boolean")]
113fn build(return_type: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
114    let mut iter = children.into_iter();
115    let left_expr = iter.next().unwrap();
116    let mut data = Vec::with_capacity(iter.size_hint().0);
117    let data_chunk = DataChunk::new_dummy(1);
118    for child in iter {
119        let BoxedExpression::Sync(child) = child else {
120            bail!("IN set expression must be sync")
121        };
122        let array = child.eval(&data_chunk)?;
123        let datum = array.value_at(0).to_owned_datum();
124        data.push(datum);
125    }
126    Ok(match left_expr {
127        BoxedExpression::Sync(left_expr) => {
128            InExpression::new(left_expr, data.into_iter(), return_type).boxed()
129        }
130        left_expr @ BoxedExpression::Async(_) => {
131            InExpression::new(left_expr, data.into_iter(), return_type).boxed()
132        }
133    })
134}
135
136#[cfg(test)]
137mod tests {
138    use risingwave_common::array::DataChunk;
139    use risingwave_common::row::Row;
140    use risingwave_common::test_prelude::DataChunkTestExt;
141    use risingwave_common::types::ToOwnedDatum;
142    use risingwave_common::util::iter_util::ZipEqDebug;
143    use risingwave_expr::expr::build_from_pretty;
144
145    #[tokio::test]
146    async fn test_in_expr() {
147        let expr = build_from_pretty("(in:boolean $0:varchar abc:varchar def:varchar)");
148        let (input, expected) = DataChunk::from_pretty(
149            "T   B
150             abc t
151             a   f
152             def t
153             abc t
154             .   .",
155        )
156        .split_column_at(1);
157
158        // test eval
159        let output = expr.eval(&input).await.unwrap();
160        assert_eq!(&output, expected.column_at(0));
161
162        // test eval_row
163        for (row, expected) in input.rows().zip_eq_debug(expected.rows()) {
164            let result = expr.eval_row(&row.to_owned_row()).await.unwrap();
165            assert_eq!(result, expected.datum_at(0).to_owned_datum());
166        }
167    }
168
169    #[tokio::test]
170    async fn test_in_expr_null() {
171        let expr = build_from_pretty("(in:boolean $0:varchar abc:varchar null:varchar)");
172        let (input, expected) = DataChunk::from_pretty(
173            "T   B
174             abc t
175             a   .
176             .   .",
177        )
178        .split_column_at(1);
179
180        // test eval
181        let output = expr.eval(&input).await.unwrap();
182        assert_eq!(&output, expected.column_at(0));
183
184        // test eval_row
185        for (row, expected) in input.rows().zip_eq_debug(expected.rows()) {
186            let result = expr.eval_row(&row.to_owned_row()).await.unwrap();
187            assert_eq!(result, expected.datum_at(0).to_owned_datum());
188        }
189    }
190}
risingwave_expr_impl/scalar/in_.rs

risingwave_expr_impl/scalar/
in_.rs
