vortex_array/arrays/constant/compute/
boolean.rs

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

vortex_array/arrays/constant/compute/
boolean.rs
