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

vortex_array/arrays/constant/compute/
boolean.rs
