vortex_array/arrays/constant/compute/
boolean.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::DType;
5use vortex_error::VortexResult;
6use vortex_error::vortex_bail;
7use vortex_error::vortex_err;
8use vortex_scalar::Scalar;
9
10use crate::Array;
11use crate::ArrayRef;
12use crate::IntoArray;
13use crate::arrays::ConstantArray;
14use crate::arrays::ConstantVTable;
15use crate::compute::BooleanKernel;
16use crate::compute::BooleanKernelAdapter;
17use crate::compute::BooleanOperator;
18use crate::register_kernel;
19
20impl BooleanKernel for ConstantVTable {
21    fn boolean(
22        &self,
23        lhs: &ConstantArray,
24        rhs: &dyn Array,
25        op: BooleanOperator,
26    ) -> VortexResult<Option<ArrayRef>> {
27        // We only implement this for constant <-> constant arrays, otherwise we allow fall back
28        // to the Arrow implementation.
29        if !rhs.is_constant() {
30            return Ok(None);
31        }
32
33        let length = lhs.len();
34        let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
35        let lhs = lhs.scalar().as_bool().value();
36        let Some(rhs) = rhs.as_constant() else {
37            vortex_bail!("Binary boolean operation requires both sides to be constant");
38        };
39        let rhs = rhs
40            .as_bool_opt()
41            .ok_or_else(|| vortex_err!("expected rhs to be boolean"))?
42            .value();
43
44        let result = match op {
45            BooleanOperator::And => and(lhs, rhs),
46            BooleanOperator::AndKleene => kleene_and(lhs, rhs),
47            BooleanOperator::Or => or(lhs, rhs),
48            BooleanOperator::OrKleene => kleene_or(lhs, rhs),
49        };
50
51        let scalar = result
52            .map(|b| Scalar::bool(b, nullable.into()))
53            .unwrap_or_else(|| Scalar::null(DType::Bool(nullable.into())));
54
55        Ok(Some(ConstantArray::new(scalar, length).into_array()))
56    }
57}
58
59register_kernel!(BooleanKernelAdapter(ConstantVTable).lift());
60
61fn and(left: Option<bool>, right: Option<bool>) -> Option<bool> {
62    left.zip(right).map(|(l, r)| l & r)
63}
64
65fn kleene_and(left: Option<bool>, right: Option<bool>) -> Option<bool> {
66    match (left, right) {
67        (Some(false), _) => Some(false),
68        (_, Some(false)) => Some(false),
69        (None, _) => None,
70        (_, None) => None,
71        (Some(l), Some(r)) => Some(l & r),
72    }
73}
74
75fn or(left: Option<bool>, right: Option<bool>) -> Option<bool> {
76    left.zip(right).map(|(l, r)| l | r)
77}
78
79fn kleene_or(left: Option<bool>, right: Option<bool>) -> Option<bool> {
80    match (left, right) {
81        (Some(true), _) => Some(true),
82        (_, Some(true)) => Some(true),
83        (None, _) => None,
84        (_, None) => None,
85        (Some(l), Some(r)) => Some(l | r),
86    }
87}
88
89#[cfg(test)]
90mod test {
91    use rstest::rstest;
92
93    use crate::Array;
94    use crate::ArrayRef;
95    use crate::IntoArray;
96    use crate::arrays::BoolArray;
97    use crate::arrays::constant::ConstantArray;
98    use crate::canonical::ToCanonical;
99    use crate::compute::and;
100    use crate::compute::or;
101
102    #[rstest]
103    #[case(ConstantArray::new(true, 4).into_array(), BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array()
104    )]
105    #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(), ConstantArray::new(true, 4).into_array()
106    )]
107    fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
108        let r = or(&lhs, &rhs).unwrap().to_bool().into_array();
109
110        let v0 = r.scalar_at(0).as_bool().value();
111        let v1 = r.scalar_at(1).as_bool().value();
112        let v2 = r.scalar_at(2).as_bool().value();
113        let v3 = r.scalar_at(3).as_bool().value();
114
115        assert!(v0.unwrap());
116        assert!(v1.unwrap());
117        assert!(v2.unwrap());
118        assert!(v3.unwrap());
119    }
120
121    #[rstest]
122    #[case(ConstantArray::new(true, 4).into_array(), BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array()
123    )]
124    #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(),
125        ConstantArray::new(true, 4).into_array())]
126    fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
127        let r = and(&lhs, &rhs).unwrap().to_bool().into_array();
128
129        let v0 = r.scalar_at(0).as_bool().value();
130        let v1 = r.scalar_at(1).as_bool().value();
131        let v2 = r.scalar_at(2).as_bool().value();
132        let v3 = r.scalar_at(3).as_bool().value();
133
134        assert!(v0.unwrap());
135        assert!(!v1.unwrap());
136        assert!(v2.unwrap());
137        assert!(!v3.unwrap());
138    }
139}
vortex_array/arrays/constant/compute/boolean.rs

vortex_array/arrays/constant/compute/
boolean.rs
