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
15pub fn and(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
22 boolean(lhs, rhs, BooleanOperator::And)
23}
24
25pub fn and_kleene(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
29 boolean(lhs, rhs, BooleanOperator::AndKleene)
30}
31
32pub fn or(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
39 boolean(lhs, rhs, BooleanOperator::Or)
40}
41
42pub fn or_kleene(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
46 boolean(lhs, rhs, BooleanOperator::OrKleene)
47}
48
49pub 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() && !rhs_is_constant {
115 return Ok(boolean(rhs, lhs, operator)?.into());
116 }
117
118 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 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 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
199pub enum BooleanOperator {
200 And,
209 AndKleene,
218 Or,
227 OrKleene,
236 }
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
280pub(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}