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

vortex_array/scalar_fn/fns/binary/
boolean.rs
