Skip to main content

vortex_array/compute/
boolean.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5
6use arrow_array::cast::AsArray;
7use arrow_schema::DataType;
8use vortex_error::VortexResult;
9
10use crate::Array;
11use crate::ArrayRef;
12use crate::IntoArray;
13use crate::arrays::ScalarFnArray;
14use crate::arrow::FromArrowArray;
15use crate::arrow::IntoArrowArray;
16use crate::compute::Options;
17use crate::expr::Binary;
18use crate::expr::ScalarFn;
19use crate::expr::operators::Operator;
20
21/// Point-wise logical _and_ between two Boolean arrays.
22///
23/// This method uses Arrow-style null propagation rather than the Kleene logic semantics. This
24/// semantics is also known as "Bochvar logic" and "weak Kleene logic".
25///
26/// See also [BooleanOperator::And]
27#[deprecated(note = "Use and_kleene instead. Non-Kleene boolean ops cannot be lazily evaluated.")]
28pub fn and(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
29    boolean(lhs, rhs, BooleanOperator::And)
30}
31
32/// Point-wise Kleene logical _and_ between two Boolean arrays.
33///
34/// See also [BooleanOperator::AndKleene]
35pub fn and_kleene(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
36    boolean(lhs, rhs, BooleanOperator::AndKleene)
37}
38
39/// Point-wise logical _or_ between two Boolean arrays.
40///
41/// This method uses Arrow-style null propagation rather than the Kleene logic semantics. This
42/// semantics is also known as "Bochvar logic" and "weak Kleene logic".
43///
44/// See also [BooleanOperator::Or]
45#[deprecated(note = "Use or_kleene instead. Non-Kleene boolean ops cannot be lazily evaluated.")]
46pub fn or(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
47    boolean(lhs, rhs, BooleanOperator::Or)
48}
49
50/// Point-wise Kleene logical _or_ between two Boolean arrays.
51///
52/// See also [BooleanOperator::OrKleene]
53pub fn or_kleene(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
54    boolean(lhs, rhs, BooleanOperator::OrKleene)
55}
56
57/// Point-wise logical operator between two Boolean arrays.
58pub fn boolean(lhs: &dyn Array, rhs: &dyn Array, op: BooleanOperator) -> VortexResult<ArrayRef> {
59    match Operator::try_from(op) {
60        Ok(expr_op) => Ok(ScalarFnArray::try_new(
61            ScalarFn::new(Binary, expr_op),
62            vec![lhs.to_array(), rhs.to_array()],
63            lhs.len(),
64        )?
65        .into_array()),
66        Err(_) => {
67            tracing::trace!(
68                "non-Kleene boolean op {op:?} cannot be lazily evaluated, falling back to eager Arrow evaluation"
69            );
70            arrow_boolean(lhs.to_array(), rhs.to_array(), op)
71        }
72    }
73}
74
75/// Operations over the nullable Boolean values.
76///
77/// All three operators accept and produce values from the set {true, false, and null}.
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub enum BooleanOperator {
80    /// Logical and, unless either value is null, in which case the result is null.
81    ///
82    /// | A ∧ B |       | **B** |       |       |
83    /// |:-----:|:-----:|:-----:|:-----:|:-----:|
84    /// |       |       | **F** | **U** | **T** |
85    /// | **A** | **F** | F     | U     | F     |
86    /// |       | **U** | U     | U     | U     |
87    /// |       | **T** | F     | U     | T     |
88    And,
89    /// [Kleene (three-valued) logical and](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics).
90    ///
91    /// | A ∧ B |       | **B** |       |       |
92    /// |:-----:|:-----:|:-----:|:-----:|:-----:|
93    /// |       |       | **F** | **U** | **T** |
94    /// | **A** | **F** | F     | F     | F     |
95    /// |       | **U** | F     | U     | U     |
96    /// |       | **T** | F     | U     | T     |
97    AndKleene,
98    /// Logical or, unless either value is null, in which case the result is null.
99    ///
100    /// | A ∨ B |       | **B** |       |       |
101    /// |:-----:|:-----:|:-----:|:-----:|:-----:|
102    /// |       |       | **F** | **U** | **T** |
103    /// | **A** | **F** | F     | U     | T     |
104    /// |       | **U** | U     | U     | U     |
105    /// |       | **T** | T     | U     | T     |
106    Or,
107    /// [Kleene (three-valued) logical or](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics).
108    ///
109    /// | A ∨ B |       | **B** |       |       |
110    /// |:-----:|:-----:|:-----:|:-----:|:-----:|
111    /// |       |       | **F** | **U** | **T** |
112    /// | **A** | **F** | F     | U     | T     |
113    /// |       | **U** | U     | U     | T     |
114    /// |       | **T** | T     | T     | T     |
115    OrKleene,
116    // AndNot,
117    // AndNotKleene,
118    // Xor,
119}
120
121impl Options for BooleanOperator {
122    fn as_any(&self) -> &dyn Any {
123        self
124    }
125}
126
127/// Implementation of `BinaryBooleanFn` using the Arrow crate.
128///
129/// Note that other encodings should handle a constant RHS value, so we can assume here that
130/// the RHS is not constant and expand to a full array.
131pub(crate) fn arrow_boolean(
132    lhs: ArrayRef,
133    rhs: ArrayRef,
134    operator: BooleanOperator,
135) -> VortexResult<ArrayRef> {
136    let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
137
138    let lhs = lhs.into_arrow(&DataType::Boolean)?.as_boolean().clone();
139    let rhs = rhs.into_arrow(&DataType::Boolean)?.as_boolean().clone();
140
141    let array = match operator {
142        BooleanOperator::And => arrow_arith::boolean::and(&lhs, &rhs)?,
143        BooleanOperator::AndKleene => arrow_arith::boolean::and_kleene(&lhs, &rhs)?,
144        BooleanOperator::Or => arrow_arith::boolean::or(&lhs, &rhs)?,
145        BooleanOperator::OrKleene => arrow_arith::boolean::or_kleene(&lhs, &rhs)?,
146    };
147
148    ArrayRef::from_arrow(&array, nullable)
149}
150
151#[cfg(test)]
152mod tests {
153    use rstest::rstest;
154
155    use super::*;
156    use crate::IntoArray;
157    use crate::arrays::BoolArray;
158    use crate::canonical::ToCanonical;
159    #[rstest]
160    #[case(BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter())
161    .into_array(), BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter())
162    .into_array())]
163    #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(),
164        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter()).into_array())]
165    fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
166        let r = or_kleene(&lhs, &rhs).unwrap();
167
168        let r = r.to_bool().into_array();
169
170        let v0 = r.scalar_at(0).unwrap().as_bool().value();
171        let v1 = r.scalar_at(1).unwrap().as_bool().value();
172        let v2 = r.scalar_at(2).unwrap().as_bool().value();
173        let v3 = r.scalar_at(3).unwrap().as_bool().value();
174
175        assert!(v0.unwrap());
176        assert!(v1.unwrap());
177        assert!(v2.unwrap());
178        assert!(!v3.unwrap());
179    }
180
181    #[rstest]
182    #[case(BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter())
183    .into_array(), BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter())
184    .into_array())]
185    #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(),
186        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter()).into_array())]
187    fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
188        let r = and_kleene(&lhs, &rhs).unwrap().to_bool().into_array();
189
190        let v0 = r.scalar_at(0).unwrap().as_bool().value();
191        let v1 = r.scalar_at(1).unwrap().as_bool().value();
192        let v2 = r.scalar_at(2).unwrap().as_bool().value();
193        let v3 = r.scalar_at(3).unwrap().as_bool().value();
194
195        assert!(v0.unwrap());
196        assert!(!v1.unwrap());
197        assert!(!v2.unwrap());
198        assert!(!v3.unwrap());
199    }
200}