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

vortex_array/scalar_fn/fns/binary/
boolean.rs
