vortex_array/compute/
boolean.rs

1use std::any::Any;
2use std::sync::LazyLock;
3
4use arcref::ArcRef;
5use arrow_array::cast::AsArray;
6use arrow_schema::DataType;
7use vortex_dtype::DType;
8use vortex_error::{VortexError, VortexExpect, VortexResult, vortex_bail, vortex_err};
9
10use crate::arrow::{FromArrowArray, IntoArrowArray};
11use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Options, Output};
12use crate::vtable::VTable;
13use crate::{Array, ArrayRef};
14
15/// Point-wise logical _and_ between two Boolean arrays.
16///
17/// This method uses Arrow-style null propagation rather than the Kleene logic semantics. This
18/// semantics is also known as "Bochvar logic" and "weak Kleene logic".
19///
20/// See also [BooleanOperator::And]
21pub fn and(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
22    boolean(lhs, rhs, BooleanOperator::And)
23}
24
25/// Point-wise Kleene logical _and_ between two Boolean arrays.
26///
27/// See also [BooleanOperator::AndKleene]
28pub fn and_kleene(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
29    boolean(lhs, rhs, BooleanOperator::AndKleene)
30}
31
32/// Point-wise logical _or_ between two Boolean arrays.
33///
34/// This method uses Arrow-style null propagation rather than the Kleene logic semantics. This
35/// semantics is also known as "Bochvar logic" and "weak Kleene logic".
36///
37/// See also [BooleanOperator::Or]
38pub fn or(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
39    boolean(lhs, rhs, BooleanOperator::Or)
40}
41
42/// Point-wise Kleene logical _or_ between two Boolean arrays.
43///
44/// See also [BooleanOperator::OrKleene]
45pub fn or_kleene(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
46    boolean(lhs, rhs, BooleanOperator::OrKleene)
47}
48
49/// Point-wise logical operator between two Boolean arrays.
50///
51/// This method uses Arrow-style null propagation rather than the Kleene logic semantics. This
52/// semantics is also known as "Bochvar logic" and "weak Kleene logic".
53pub fn boolean(lhs: &dyn Array, rhs: &dyn Array, op: BooleanOperator) -> VortexResult<ArrayRef> {
54    BOOLEAN_FN
55        .invoke(&InvocationArgs {
56            inputs: &[lhs.into(), rhs.into()],
57            options: &op,
58        })?
59        .unwrap_array()
60}
61
62pub struct BooleanKernelRef(ArcRef<dyn Kernel>);
63inventory::collect!(BooleanKernelRef);
64
65pub trait BooleanKernel: VTable {
66    fn boolean(
67        &self,
68        array: &Self::Array,
69        other: &dyn Array,
70        op: BooleanOperator,
71    ) -> VortexResult<Option<ArrayRef>>;
72}
73
74#[derive(Debug)]
75pub struct BooleanKernelAdapter<V: VTable>(pub V);
76
77impl<V: VTable + BooleanKernel> BooleanKernelAdapter<V> {
78    pub const fn lift(&'static self) -> BooleanKernelRef {
79        BooleanKernelRef(ArcRef::new_ref(self))
80    }
81}
82
83impl<V: VTable + BooleanKernel> Kernel for BooleanKernelAdapter<V> {
84    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
85        let inputs = BooleanArgs::try_from(args)?;
86        let Some(array) = inputs.lhs.as_opt::<V>() else {
87            return Ok(None);
88        };
89        Ok(V::boolean(&self.0, array, inputs.rhs, inputs.operator)?.map(|array| array.into()))
90    }
91}
92
93pub static BOOLEAN_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
94    let compute = ComputeFn::new("boolean".into(), ArcRef::new_ref(&Boolean));
95    for kernel in inventory::iter::<BooleanKernelRef> {
96        compute.register_kernel(kernel.0.clone());
97    }
98    compute
99});
100
101struct Boolean;
102
103impl ComputeFnVTable for Boolean {
104    fn invoke(
105        &self,
106        args: &InvocationArgs,
107        kernels: &[ArcRef<dyn Kernel>],
108    ) -> VortexResult<Output> {
109        let BooleanArgs { lhs, rhs, operator } = BooleanArgs::try_from(args)?;
110
111        let rhs_is_constant = rhs.is_constant();
112
113        // If LHS is constant, then we make sure it's on the RHS.
114        if lhs.is_constant() && !rhs_is_constant {
115            return Ok(boolean(rhs, lhs, operator)?.into());
116        }
117
118        // If the RHS is constant and the LHS is Arrow, we can't do any better than arrow_compare.
119        if lhs.is_arrow() && (rhs.is_arrow() || rhs_is_constant) {
120            return Ok(arrow_boolean(lhs.to_array(), rhs.to_array(), operator)?.into());
121        }
122
123        // Check if either LHS or RHS supports the operation directly.
124        for kernel in kernels {
125            if let Some(output) = kernel.invoke(args)? {
126                return Ok(output);
127            }
128        }
129        if let Some(output) = lhs.invoke(&BOOLEAN_FN, args)? {
130            return Ok(output);
131        }
132
133        let inverse_args = InvocationArgs {
134            inputs: &[rhs.into(), lhs.into()],
135            options: &operator,
136        };
137        for kernel in kernels {
138            if let Some(output) = kernel.invoke(&inverse_args)? {
139                return Ok(output);
140            }
141        }
142        if let Some(output) = rhs.invoke(&BOOLEAN_FN, &inverse_args)? {
143            return Ok(output);
144        }
145
146        log::debug!(
147            "No boolean implementation found for LHS {}, RHS {}, and operator {:?} (or inverse)",
148            rhs.encoding_id(),
149            lhs.encoding_id(),
150            operator,
151        );
152
153        // If neither side implements the trait, then we delegate to Arrow compute.
154        Ok(arrow_boolean(lhs.to_array(), rhs.to_array(), operator)?.into())
155    }
156
157    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
158        let BooleanArgs { lhs, rhs, .. } = BooleanArgs::try_from(args)?;
159
160        if !lhs.dtype().is_boolean()
161            || !rhs.dtype().is_boolean()
162            || !lhs.dtype().eq_ignore_nullability(rhs.dtype())
163        {
164            vortex_bail!(
165                "Boolean operations are only supported on boolean arrays: {} and {}",
166                lhs.dtype(),
167                rhs.dtype()
168            )
169        }
170
171        Ok(DType::Bool(
172            (lhs.dtype().is_nullable() || rhs.dtype().is_nullable()).into(),
173        ))
174    }
175
176    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
177        let BooleanArgs { lhs, rhs, .. } = BooleanArgs::try_from(args)?;
178
179        if lhs.len() != rhs.len() {
180            vortex_bail!(
181                "Boolean operations aren't supported on arrays of different lengths: {} and {}",
182                lhs.len(),
183                rhs.len()
184            )
185        }
186
187        Ok(lhs.len())
188    }
189
190    fn is_elementwise(&self) -> bool {
191        true
192    }
193}
194
195/// Operations over the nullable Boolean values.
196///
197/// All three operators accept and produce values from the set {true, false, and null}.
198#[derive(Debug, Clone, Copy, PartialEq, Eq)]
199pub enum BooleanOperator {
200    /// Logical and, unless either value is null, in which case the result is null.
201    ///
202    /// | A ∧ B |       | **B** |       |       |
203    /// |:-----:|:-----:|:-----:|:-----:|:-----:|
204    /// |       |       | **F** | **U** | **T** |
205    /// | **A** | **F** | F     | U     | F     |
206    /// |       | **U** | U     | U     | U     |
207    /// |       | **T** | F     | U     | T     |
208    And,
209    /// [Kleene (three-valued) logical and](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics).
210    ///
211    /// | A ∧ B |       | **B** |       |       |
212    /// |:-----:|:-----:|:-----:|:-----:|:-----:|
213    /// |       |       | **F** | **U** | **T** |
214    /// | **A** | **F** | F     | F     | F     |
215    /// |       | **U** | F     | U     | U     |
216    /// |       | **T** | F     | U     | T     |
217    AndKleene,
218    /// Logical or, unless either value is null, in which case the result is null.
219    ///
220    /// | A ∨ B |       | **B** |       |       |
221    /// |:-----:|:-----:|:-----:|:-----:|:-----:|
222    /// |       |       | **F** | **U** | **T** |
223    /// | **A** | **F** | F     | U     | T     |
224    /// |       | **U** | U     | U     | U     |
225    /// |       | **T** | T     | U     | T     |
226    Or,
227    /// [Kleene (three-valued) logical or](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics).
228    ///
229    /// | A ∨ B |       | **B** |       |       |
230    /// |:-----:|:-----:|:-----:|:-----:|:-----:|
231    /// |       |       | **F** | **U** | **T** |
232    /// | **A** | **F** | F     | U     | T     |
233    /// |       | **U** | U     | U     | T     |
234    /// |       | **T** | T     | T     | T     |
235    OrKleene,
236    // AndNot,
237    // AndNotKleene,
238    // Xor,
239}
240
241impl Options for BooleanOperator {
242    fn as_any(&self) -> &dyn Any {
243        self
244    }
245}
246
247struct BooleanArgs<'a> {
248    lhs: &'a dyn Array,
249    rhs: &'a dyn Array,
250    operator: BooleanOperator,
251}
252
253impl<'a> TryFrom<&InvocationArgs<'a>> for BooleanArgs<'a> {
254    type Error = VortexError;
255
256    fn try_from(value: &InvocationArgs<'a>) -> VortexResult<Self> {
257        if value.inputs.len() != 2 {
258            vortex_bail!("Expected 2 inputs, found {}", value.inputs.len());
259        }
260        let lhs = value.inputs[0]
261            .array()
262            .ok_or_else(|| vortex_err!("Expected input 0 to be an array"))?;
263        let rhs = value.inputs[1]
264            .array()
265            .ok_or_else(|| vortex_err!("Expected input 1 to be an array"))?;
266        let operator = value
267            .options
268            .as_any()
269            .downcast_ref::<BooleanOperator>()
270            .vortex_expect("Expected options to be an operator");
271
272        Ok(BooleanArgs {
273            lhs,
274            rhs,
275            operator: *operator,
276        })
277    }
278}
279
280/// Implementation of `BinaryBooleanFn` using the Arrow crate.
281///
282/// Note that other encodings should handle a constant RHS value, so we can assume here that
283/// the RHS is not constant and expand to a full array.
284pub(crate) fn arrow_boolean(
285    lhs: ArrayRef,
286    rhs: ArrayRef,
287    operator: BooleanOperator,
288) -> VortexResult<ArrayRef> {
289    let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
290
291    let lhs = lhs.into_arrow(&DataType::Boolean)?.as_boolean().clone();
292    let rhs = rhs.into_arrow(&DataType::Boolean)?.as_boolean().clone();
293
294    let array = match operator {
295        BooleanOperator::And => arrow_arith::boolean::and(&lhs, &rhs)?,
296        BooleanOperator::AndKleene => arrow_arith::boolean::and_kleene(&lhs, &rhs)?,
297        BooleanOperator::Or => arrow_arith::boolean::or(&lhs, &rhs)?,
298        BooleanOperator::OrKleene => arrow_arith::boolean::or_kleene(&lhs, &rhs)?,
299    };
300
301    Ok(ArrayRef::from_arrow(&array, nullable))
302}
303
304#[cfg(test)]
305mod tests {
306    use rstest::rstest;
307
308    use super::*;
309    use crate::IntoArray;
310    use crate::arrays::BoolArray;
311    use crate::canonical::ToCanonical;
312    #[rstest]
313    #[case(BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter())
314    .into_array(), BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter())
315    .into_array())]
316    #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(),
317        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter()).into_array())]
318    fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
319        let r = or(&lhs, &rhs).unwrap();
320
321        let r = r.to_bool().unwrap().into_array();
322
323        let v0 = r.scalar_at(0).unwrap().as_bool().value();
324        let v1 = r.scalar_at(1).unwrap().as_bool().value();
325        let v2 = r.scalar_at(2).unwrap().as_bool().value();
326        let v3 = r.scalar_at(3).unwrap().as_bool().value();
327
328        assert!(v0.unwrap());
329        assert!(v1.unwrap());
330        assert!(v2.unwrap());
331        assert!(!v3.unwrap());
332    }
333
334    #[rstest]
335    #[case(BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter())
336    .into_array(), BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter())
337    .into_array())]
338    #[case(BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)].into_iter()).into_array(),
339        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)].into_iter()).into_array())]
340    fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
341        let r = and(&lhs, &rhs).unwrap().to_bool().unwrap().into_array();
342
343        let v0 = r.scalar_at(0).unwrap().as_bool().value();
344        let v1 = r.scalar_at(1).unwrap().as_bool().value();
345        let v2 = r.scalar_at(2).unwrap().as_bool().value();
346        let v3 = r.scalar_at(3).unwrap().as_bool().value();
347
348        assert!(v0.unwrap());
349        assert!(!v1.unwrap());
350        assert!(!v2.unwrap());
351        assert!(!v3.unwrap());
352    }
353}