vortex_array/scalar_fn/fns/binary/
boolean.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use arrow_array::cast::AsArray;
5use arrow_schema::DataType;
6use vortex_error::VortexResult;
7use vortex_error::vortex_err;
8
9use crate::ArrayRef;
10use crate::IntoArray;
11use crate::arrays::Constant;
12use crate::arrays::ConstantArray;
13use crate::arrow::FromArrowArray;
14use crate::arrow::IntoArrowArray;
15use crate::builtins::ArrayBuiltins;
16use crate::dtype::DType;
17use crate::scalar::Scalar;
18use crate::scalar_fn::fns::operators::Operator;
19
20/// Point-wise Kleene logical _and_ between two Boolean arrays.
21#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
22pub fn and_kleene(lhs: &ArrayRef, rhs: &ArrayRef) -> VortexResult<ArrayRef> {
23    lhs.clone().binary(rhs.clone(), Operator::And)
24}
25
26/// Point-wise Kleene logical _or_ between two Boolean arrays.
27#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
28pub fn or_kleene(lhs: &ArrayRef, rhs: &ArrayRef) -> VortexResult<ArrayRef> {
29    lhs.clone().binary(rhs.clone(), Operator::Or)
30}
31
32/// Execute a Kleene boolean operation between two arrays.
33///
34/// This is the entry point for boolean operations from the binary expression.
35/// Handles constant-constant directly, otherwise falls back to Arrow.
36pub(crate) fn execute_boolean(
37    lhs: &ArrayRef,
38    rhs: &ArrayRef,
39    op: Operator,
40) -> VortexResult<ArrayRef> {
41    if let Some(result) = constant_boolean(lhs, rhs, op)? {
42        return Ok(result);
43    }
44    arrow_execute_boolean(lhs.clone(), rhs.clone(), op)
45}
46
47/// Arrow implementation for Kleene boolean operations using [`Operator`].
48fn arrow_execute_boolean(lhs: ArrayRef, rhs: ArrayRef, op: Operator) -> VortexResult<ArrayRef> {
49    let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
50
51    let lhs = lhs.into_arrow(&DataType::Boolean)?.as_boolean().clone();
52    let rhs = rhs.into_arrow(&DataType::Boolean)?.as_boolean().clone();
53
54    let array = match op {
55        Operator::And => arrow_arith::boolean::and_kleene(&lhs, &rhs)?,
56        Operator::Or => arrow_arith::boolean::or_kleene(&lhs, &rhs)?,
57        other => return Err(vortex_err!("Not a boolean operator: {other}")),
58    };
59
60    ArrayRef::from_arrow(&array, nullable)
61}
62
63/// Constant-folds a boolean operation between two constant arrays.
64fn constant_boolean(
65    lhs: &ArrayRef,
66    rhs: &ArrayRef,
67    op: Operator,
68) -> VortexResult<Option<ArrayRef>> {
69    let (Some(lhs), Some(rhs)) = (lhs.as_opt::<Constant>(), rhs.as_opt::<Constant>()) else {
70        return Ok(None);
71    };
72
73    let length = lhs.len();
74    let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
75    let lhs_val = lhs.scalar().as_bool().value();
76    let rhs_val = rhs
77        .scalar()
78        .as_bool_opt()
79        .ok_or_else(|| vortex_err!("expected rhs to be boolean"))?
80        .value();
81
82    let result = match op {
83        Operator::And => match (lhs_val, rhs_val) {
84            (Some(false), _) | (_, Some(false)) => Some(false),
85            (None, _) | (_, None) => None,
86            (Some(l), Some(r)) => Some(l & r),
87        },
88        Operator::Or => match (lhs_val, rhs_val) {
89            (Some(true), _) | (_, Some(true)) => Some(true),
90            (None, _) | (_, None) => None,
91            (Some(l), Some(r)) => Some(l | r),
92        },
93        other => return Err(vortex_err!("Not a boolean operator: {other}")),
94    };
95
96    let scalar = result
97        .map(|b| Scalar::bool(b, nullable.into()))
98        .unwrap_or_else(|| Scalar::null(DType::Bool(nullable.into())));
99
100    Ok(Some(ConstantArray::new(scalar, length).into_array()))
101}
102
103#[cfg(test)]
104mod tests {
105    use rstest::rstest;
106
107    use crate::ArrayRef;
108    use crate::IntoArray;
109    use crate::LEGACY_SESSION;
110    use crate::VortexSessionExecute;
111    use crate::arrays::BoolArray;
112    use crate::builtins::ArrayBuiltins;
113    use crate::canonical::ToCanonical;
114    use crate::scalar_fn::fns::operators::Operator;
115
116    #[rstest]
117    #[case(
118        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(),
119        BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)]).into_array(),
120    )]
121    #[case(
122        BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)]).into_array(),
123        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(),
124    )]
125    fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
126        let r = lhs.binary(rhs, Operator::Or).unwrap();
127        let r = r.to_bool().into_array();
128
129        let v0 = r
130            .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())
131            .unwrap()
132            .as_bool()
133            .value();
134        let v1 = r
135            .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())
136            .unwrap()
137            .as_bool()
138            .value();
139        let v2 = r
140            .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())
141            .unwrap()
142            .as_bool()
143            .value();
144        let v3 = r
145            .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx())
146            .unwrap()
147            .as_bool()
148            .value();
149
150        assert!(v0.unwrap());
151        assert!(v1.unwrap());
152        assert!(v2.unwrap());
153        assert!(!v3.unwrap());
154    }
155
156    #[rstest]
157    #[case(
158        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(),
159        BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)]).into_array(),
160    )]
161    #[case(
162        BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)]).into_array(),
163        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(),
164    )]
165    fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
166        let r = lhs
167            .binary(rhs, Operator::And)
168            .unwrap()
169            .to_bool()
170            .into_array();
171
172        let v0 = r
173            .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())
174            .unwrap()
175            .as_bool()
176            .value();
177        let v1 = r
178            .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())
179            .unwrap()
180            .as_bool()
181            .value();
182        let v2 = r
183            .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())
184            .unwrap()
185            .as_bool()
186            .value();
187        let v3 = r
188            .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx())
189            .unwrap()
190            .as_bool()
191            .value();
192
193        assert!(v0.unwrap());
194        assert!(!v1.unwrap());
195        assert!(!v2.unwrap());
196        assert!(!v3.unwrap());
197    }
198}
vortex_array/scalar_fn/fns/binary/boolean.rs

vortex_array/scalar_fn/fns/binary/
boolean.rs
