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::DynArray;
11use crate::IntoArray;
12use crate::arrays::ConstantArray;
13use crate::arrays::ConstantVTable;
14use crate::arrow::FromArrowArray;
15use crate::arrow::IntoArrowArray;
16use crate::builtins::ArrayBuiltins;
17use crate::dtype::DType;
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.to_array().binary(rhs.to_array(), 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.to_array().binary(rhs.to_array(), 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) -> VortexResult<ArrayRef> {
42    if let Some(result) = constant_boolean(lhs, rhs, op)? {
43        return Ok(result);
44    }
45    arrow_execute_boolean(lhs.to_array(), rhs.to_array(), op)
46}
47
48/// Arrow implementation for Kleene boolean operations using [`Operator`].
49fn arrow_execute_boolean(lhs: ArrayRef, rhs: ArrayRef, op: Operator) -> VortexResult<ArrayRef> {
50    let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
51
52    let lhs = lhs.into_arrow(&DataType::Boolean)?.as_boolean().clone();
53    let rhs = rhs.into_arrow(&DataType::Boolean)?.as_boolean().clone();
54
55    let array = match op {
56        Operator::And => arrow_arith::boolean::and_kleene(&lhs, &rhs)?,
57        Operator::Or => arrow_arith::boolean::or_kleene(&lhs, &rhs)?,
58        other => return Err(vortex_err!("Not a boolean operator: {other}")),
59    };
60
61    ArrayRef::from_arrow(&array, nullable)
62}
63
64/// Constant-folds a boolean operation between two constant arrays.
65fn constant_boolean(
66    lhs: &ArrayRef,
67    rhs: &ArrayRef,
68    op: Operator,
69) -> VortexResult<Option<ArrayRef>> {
70    let (Some(lhs), Some(rhs)) = (
71        lhs.as_opt::<ConstantVTable>(),
72        rhs.as_opt::<ConstantVTable>(),
73    ) else {
74        return Ok(None);
75    };
76
77    let length = lhs.len();
78    let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
79    let lhs_val = lhs.scalar().as_bool().value();
80    let rhs_val = rhs
81        .scalar()
82        .as_bool_opt()
83        .ok_or_else(|| vortex_err!("expected rhs to be boolean"))?
84        .value();
85
86    let result = match op {
87        Operator::And => match (lhs_val, rhs_val) {
88            (Some(false), _) | (_, Some(false)) => Some(false),
89            (None, _) | (_, None) => None,
90            (Some(l), Some(r)) => Some(l & r),
91        },
92        Operator::Or => match (lhs_val, rhs_val) {
93            (Some(true), _) | (_, Some(true)) => Some(true),
94            (None, _) | (_, None) => None,
95            (Some(l), Some(r)) => Some(l | r),
96        },
97        other => return Err(vortex_err!("Not a boolean operator: {other}")),
98    };
99
100    let scalar = result
101        .map(|b| Scalar::bool(b, nullable.into()))
102        .unwrap_or_else(|| Scalar::null(DType::Bool(nullable.into())));
103
104    Ok(Some(ConstantArray::new(scalar, length).into_array()))
105}
106
107#[cfg(test)]
108mod tests {
109    use rstest::rstest;
110
111    use crate::ArrayRef;
112    use crate::IntoArray;
113    use crate::arrays::BoolArray;
114    use crate::builtins::ArrayBuiltins;
115    use crate::canonical::ToCanonical;
116    use crate::scalar_fn::fns::operators::Operator;
117
118    #[rstest]
119    #[case(
120        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(),
121        BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)]).into_array(),
122    )]
123    #[case(
124        BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)]).into_array(),
125        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(),
126    )]
127    fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
128        let r = lhs.binary(rhs, Operator::Or).unwrap();
129        let r = r.to_bool().into_array();
130
131        let v0 = r.scalar_at(0).unwrap().as_bool().value();
132        let v1 = r.scalar_at(1).unwrap().as_bool().value();
133        let v2 = r.scalar_at(2).unwrap().as_bool().value();
134        let v3 = r.scalar_at(3).unwrap().as_bool().value();
135
136        assert!(v0.unwrap());
137        assert!(v1.unwrap());
138        assert!(v2.unwrap());
139        assert!(!v3.unwrap());
140    }
141
142    #[rstest]
143    #[case(
144        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(),
145        BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)]).into_array(),
146    )]
147    #[case(
148        BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)]).into_array(),
149        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(),
150    )]
151    fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
152        let r = lhs
153            .binary(rhs, Operator::And)
154            .unwrap()
155            .to_bool()
156            .into_array();
157
158        let v0 = r.scalar_at(0).unwrap().as_bool().value();
159        let v1 = r.scalar_at(1).unwrap().as_bool().value();
160        let v2 = r.scalar_at(2).unwrap().as_bool().value();
161        let v3 = r.scalar_at(3).unwrap().as_bool().value();
162
163        assert!(v0.unwrap());
164        assert!(!v1.unwrap());
165        assert!(!v2.unwrap());
166        assert!(!v3.unwrap());
167    }
168}
vortex_array/scalar_fn/fns/binary/boolean.rs

vortex_array/scalar_fn/fns/binary/
boolean.rs
