risingwave_expr/expr/
and_or.rs

1// Copyright 2025 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::{BoxedExpression, Expression};
26use crate::Result;
27
28/// This is just an implementation detail. The semantic is not guaranteed at SQL level because
29/// optimizer may have rearranged the boolean expressions. #6202
30#[derive(Debug)]
31pub struct BinaryShortCircuitExpression {
32    expr_ia1: BoxedExpression,
33    expr_ia2: BoxedExpression,
34    expr_type: Type,
35}
36
37#[async_trait::async_trait]
38impl Expression for BinaryShortCircuitExpression {
39    fn return_type(&self) -> DataType {
40        DataType::Boolean
41    }
42
43    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
44        let left = self.expr_ia1.eval(input).await?;
45        let left = left.as_bool();
46
47        let res_vis = match self.expr_type {
48            // For `Or` operator, if res of left part is not null and is true, we do not want to
49            // calculate right part because the result must be true.
50            Type::Or => !left.to_bitmap(),
51            // For `And` operator, If res of left part is not null and is false, we do not want
52            // to calculate right part because the result must be false.
53            Type::And => left.data() | !left.null_bitmap(),
54            _ => unimplemented!(),
55        };
56        let new_vis = input.visibility() & res_vis;
57        let mut input1 = input.clone();
58        input1.set_visibility(new_vis);
59
60        let right = self.expr_ia2.eval(&input1).await?;
61        let right = right.as_bool();
62        assert_eq!(left.len(), right.len());
63
64        let mut bitmap = input.visibility() & left.null_bitmap() & right.null_bitmap();
65
66        let c = match self.expr_type {
67            Type::Or => {
68                let data = left.to_bitmap() | right.to_bitmap();
69                bitmap |= &data; // is_true || is_true
70                BoolArray::new(data, bitmap)
71            }
72            Type::And => {
73                let data = left.to_bitmap() & right.to_bitmap();
74                bitmap |= !left.data() & left.null_bitmap(); // is_false
75                bitmap |= !right.data() & right.null_bitmap(); // is_false
76                BoolArray::new(data, bitmap)
77            }
78            _ => unimplemented!(),
79        };
80        Ok(Arc::new(c.into()))
81    }
82
83    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
84        let ret_ia1 = self.expr_ia1.eval_row(input).await?.map(|x| x.into_bool());
85        match self.expr_type {
86            Type::Or if ret_ia1 == Some(true) => return Ok(Some(true.to_scalar_value())),
87            Type::And if ret_ia1 == Some(false) => return Ok(Some(false.to_scalar_value())),
88            _ => {}
89        }
90        let ret_ia2 = self.expr_ia2.eval_row(input).await?.map(|x| x.into_bool());
91        match self.expr_type {
92            Type::Or => Ok(or(ret_ia1, ret_ia2).map(|x| x.to_scalar_value())),
93            Type::And => Ok(and(ret_ia1, ret_ia2).map(|x| x.to_scalar_value())),
94            _ => unimplemented!(),
95        }
96    }
97}
98
99#[build_function("and(boolean, boolean) -> boolean")]
100fn build_and_expr(_: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
101    let mut iter = children.into_iter();
102    Ok(Box::new(BinaryShortCircuitExpression {
103        expr_ia1: iter.next().unwrap(),
104        expr_ia2: iter.next().unwrap(),
105        expr_type: Type::And,
106    }))
107}
108
109#[build_function("or(boolean, boolean) -> boolean")]
110fn build_or_expr(_: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
111    let mut iter = children.into_iter();
112    Ok(Box::new(BinaryShortCircuitExpression {
113        expr_ia1: iter.next().unwrap(),
114        expr_ia2: iter.next().unwrap(),
115        expr_type: Type::Or,
116    }))
117}
118
119// #[function("and(boolean, boolean) -> boolean")]
120fn and(l: Option<bool>, r: Option<bool>) -> Option<bool> {
121    match (l, r) {
122        (Some(lb), Some(lr)) => Some(lb & lr),
123        (Some(true), None) => None,
124        (None, Some(true)) => None,
125        (Some(false), None) => Some(false),
126        (None, Some(false)) => Some(false),
127        (None, None) => None,
128    }
129}
130
131// #[function("or(boolean, boolean) -> boolean")]
132fn or(l: Option<bool>, r: Option<bool>) -> Option<bool> {
133    match (l, r) {
134        (Some(lb), Some(lr)) => Some(lb | lr),
135        (Some(true), None) => Some(true),
136        (None, Some(true)) => Some(true),
137        (Some(false), None) => None,
138        (None, Some(false)) => None,
139        (None, None) => None,
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    use crate::expr::build_from_pretty;
147
148    #[tokio::test]
149    async fn test_and() {
150        let (input, target) = DataChunk::from_pretty(
151            "
152            B B B
153            t t t
154            t f f
155            t . .
156            f t f
157            f f f
158            f . f
159            . t .
160            . f f
161            . . .
162        ",
163        )
164        .split_column_at(2);
165        let expr = build_from_pretty("(and:boolean $0:boolean $1:boolean)");
166        let result = expr.eval(&input).await.unwrap();
167        assert_eq!(&result, target.column_at(0));
168    }
169
170    #[tokio::test]
171    async fn test_or() {
172        let (input, target) = DataChunk::from_pretty(
173            "
174            B B B
175            t t t
176            t f t
177            t . t
178            f t t
179            f f f
180            f . .
181            . t t
182            . f .
183            . . .
184        ",
185        )
186        .split_column_at(2);
187        let expr = build_from_pretty("(or:boolean $0:boolean $1:boolean)");
188        let result = expr.eval(&input).await.unwrap();
189        assert_eq!(&result, target.column_at(0));
190    }
191
192    #[test]
193    fn test_and_() {
194        assert_eq!(Some(true), and(Some(true), Some(true)));
195        assert_eq!(Some(false), and(Some(true), Some(false)));
196        assert_eq!(Some(false), and(Some(false), Some(false)));
197        assert_eq!(None, and(Some(true), None));
198        assert_eq!(Some(false), and(Some(false), None));
199        assert_eq!(None, and(None, None));
200    }
201
202    #[test]
203    fn test_or_() {
204        assert_eq!(Some(true), or(Some(true), Some(true)));
205        assert_eq!(Some(true), or(Some(true), Some(false)));
206        assert_eq!(Some(false), or(Some(false), Some(false)));
207        assert_eq!(Some(true), or(Some(true), None));
208        assert_eq!(None, or(Some(false), None));
209        assert_eq!(None, or(None, None));
210    }
211}
risingwave_expr/expr/and_or.rs

risingwave_expr/expr/
and_or.rs
