risingwave_expr/expr/
and_or.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
15//! For expression that only accept two nullable arguments as input.
16
17use std::sync::Arc;
18
19use risingwave_common::array::*;
20use risingwave_common::row::OwnedRow;
21use risingwave_common::types::{DataType, Datum, Scalar};
22use risingwave_expr_macro::build_function;
23use risingwave_pb::expr::expr_node::Type;
24
25use super::{
26    AsyncExpression, AsyncExpressionBoxExt, BoxedExpression, ExpressionInfo, SyncExpression,
27    SyncExpressionBoxExt,
28};
29use crate::Result;
30
31/// This is just an implementation detail. The semantic is not guaranteed at SQL level because
32/// optimizer may have rearranged the boolean expressions. #6202
33#[derive(Debug)]
34pub struct BinaryShortCircuitExpression<E> {
35    expr_ia1: E,
36    expr_ia2: E,
37    expr_type: Type,
38}
39
40impl<E: ExpressionInfo> ExpressionInfo for BinaryShortCircuitExpression<E> {
41    fn return_type(&self) -> DataType {
42        DataType::Boolean
43    }
44}
45
46macro_rules! eval_short_circuit {
47    ($mode:ident, $this:expr, $input:expr) => {{
48        let left = forward!($mode, $this.expr_ia1, eval($input))?;
49        let left = left.as_bool();
50
51        let res_vis = match $this.expr_type {
52            // For `Or` operator, if res of left part is not null and is true, we do not want to
53            // calculate right part because the result must be true.
54            Type::Or => !left.to_bitmap(),
55            // For `And` operator, If res of left part is not null and is false, we do not want
56            // to calculate right part because the result must be false.
57            Type::And => left.data() | !left.null_bitmap(),
58            _ => unimplemented!(),
59        };
60        let new_vis = $input.visibility() & res_vis;
61        let mut input1 = $input.clone();
62        input1.set_visibility(new_vis);
63
64        let right = forward!($mode, $this.expr_ia2, eval(&input1))?;
65        let right = right.as_bool();
66        assert_eq!(left.len(), right.len());
67
68        let mut bitmap = $input.visibility() & left.null_bitmap() & right.null_bitmap();
69
70        let c = match $this.expr_type {
71            Type::Or => {
72                let data = left.to_bitmap() | right.to_bitmap();
73                bitmap |= &data; // is_true || is_true
74                BoolArray::new(data, bitmap)
75            }
76            Type::And => {
77                let data = left.to_bitmap() & right.to_bitmap();
78                bitmap |= !left.data() & left.null_bitmap(); // is_false
79                bitmap |= !right.data() & right.null_bitmap(); // is_false
80                BoolArray::new(data, bitmap)
81            }
82            _ => unimplemented!(),
83        };
84        Ok(Arc::new(c.into()))
85    }};
86}
87
88macro_rules! eval_row_short_circuit {
89    ($mode:ident, $this:expr, $input:expr) => {{
90        let ret_ia1 = forward!($mode, $this.expr_ia1, eval_row($input))?.map(|x| x.into_bool());
91        match $this.expr_type {
92            Type::Or if ret_ia1 == Some(true) => return Ok(Some(true.to_scalar_value())),
93            Type::And if ret_ia1 == Some(false) => return Ok(Some(false.to_scalar_value())),
94            _ => {}
95        }
96        let ret_ia2 = forward!($mode, $this.expr_ia2, eval_row($input))?.map(|x| x.into_bool());
97        match $this.expr_type {
98            Type::Or => Ok(or(ret_ia1, ret_ia2).map(|x| x.to_scalar_value())),
99            Type::And => Ok(and(ret_ia1, ret_ia2).map(|x| x.to_scalar_value())),
100            _ => unimplemented!(),
101        }
102    }};
103}
104
105impl<E: SyncExpression> SyncExpression for BinaryShortCircuitExpression<E> {
106    fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
107        eval_short_circuit!(sync, self, input)
108    }
109
110    fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
111        eval_row_short_circuit!(sync, self, input)
112    }
113}
114
115impl<E: AsyncExpression> AsyncExpression for BinaryShortCircuitExpression<E> {
116    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
117        eval_short_circuit!(async, self, input)
118    }
119
120    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
121        eval_row_short_circuit!(async, self, input)
122    }
123}
124
125#[build_function("and(boolean, boolean) -> boolean")]
126fn build_and_expr(_: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
127    build_binary_short_circuit(children, Type::And)
128}
129
130#[build_function("or(boolean, boolean) -> boolean")]
131fn build_or_expr(_: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
132    build_binary_short_circuit(children, Type::Or)
133}
134
135fn build_binary_short_circuit(
136    children: Vec<BoxedExpression>,
137    expr_type: Type,
138) -> Result<BoxedExpression> {
139    let [left, right]: [_; 2] = children.try_into().unwrap();
140    Ok(match (left, right) {
141        (BoxedExpression::Sync(left), BoxedExpression::Sync(right)) => {
142            BinaryShortCircuitExpression {
143                expr_ia1: left,
144                expr_ia2: right,
145                expr_type,
146            }
147            .boxed()
148        }
149        (left, right) => BinaryShortCircuitExpression {
150            expr_ia1: left,
151            expr_ia2: right,
152            expr_type,
153        }
154        .boxed(),
155    })
156}
157
158// #[function("and(boolean, boolean) -> boolean")]
159fn and(l: Option<bool>, r: Option<bool>) -> Option<bool> {
160    match (l, r) {
161        (Some(lb), Some(lr)) => Some(lb & lr),
162        (Some(true), None) => None,
163        (None, Some(true)) => None,
164        (Some(false), None) => Some(false),
165        (None, Some(false)) => Some(false),
166        (None, None) => None,
167    }
168}
169
170// #[function("or(boolean, boolean) -> boolean")]
171fn or(l: Option<bool>, r: Option<bool>) -> Option<bool> {
172    match (l, r) {
173        (Some(lb), Some(lr)) => Some(lb | lr),
174        (Some(true), None) => Some(true),
175        (None, Some(true)) => Some(true),
176        (Some(false), None) => None,
177        (None, Some(false)) => None,
178        (None, None) => None,
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185    use crate::expr::build_from_pretty;
186
187    #[tokio::test]
188    async fn test_and() {
189        let (input, target) = DataChunk::from_pretty(
190            "
191            B B B
192            t t t
193            t f f
194            t . .
195            f t f
196            f f f
197            f . f
198            . t .
199            . f f
200            . . .
201        ",
202        )
203        .split_column_at(2);
204        let expr = build_from_pretty("(and:boolean $0:boolean $1:boolean)");
205        let result = expr.eval(&input).await.unwrap();
206        assert_eq!(&result, target.column_at(0));
207    }
208
209    #[tokio::test]
210    async fn test_or() {
211        let (input, target) = DataChunk::from_pretty(
212            "
213            B B B
214            t t t
215            t f t
216            t . t
217            f t t
218            f f f
219            f . .
220            . t t
221            . f .
222            . . .
223        ",
224        )
225        .split_column_at(2);
226        let expr = build_from_pretty("(or:boolean $0:boolean $1:boolean)");
227        let result = expr.eval(&input).await.unwrap();
228        assert_eq!(&result, target.column_at(0));
229    }
230
231    #[test]
232    fn test_and_() {
233        assert_eq!(Some(true), and(Some(true), Some(true)));
234        assert_eq!(Some(false), and(Some(true), Some(false)));
235        assert_eq!(Some(false), and(Some(false), Some(false)));
236        assert_eq!(None, and(Some(true), None));
237        assert_eq!(Some(false), and(Some(false), None));
238        assert_eq!(None, and(None, None));
239    }
240
241    #[test]
242    fn test_or_() {
243        assert_eq!(Some(true), or(Some(true), Some(true)));
244        assert_eq!(Some(true), or(Some(true), Some(false)));
245        assert_eq!(Some(false), or(Some(false), Some(false)));
246        assert_eq!(Some(true), or(Some(true), None));
247        assert_eq!(None, or(Some(false), None));
248        assert_eq!(None, or(None, None));
249    }
250}
risingwave_expr/expr/and_or.rs

risingwave_expr/expr/
and_or.rs
