Skip to main content

vortex_array/scalar_fn/fns/binary/
boolean.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::iter::repeat_n;
5
6use vortex_buffer::BitBuffer;
7use vortex_buffer::BufferAllocatorRef;
8use vortex_buffer::BufferMut;
9use vortex_buffer::read_u64_le;
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_error::vortex_err;
13use vortex_mask::AllOr;
14use vortex_mask::Mask;
15
16use crate::ArrayRef;
17use crate::Canonical;
18use crate::ExecutionCtx;
19use crate::IntoArray;
20use crate::array::ArrayView;
21use crate::array::VTable;
22use crate::arrays::Bool;
23use crate::arrays::BoolArray;
24use crate::arrays::Constant;
25use crate::arrays::ConstantArray;
26use crate::arrays::ScalarFn;
27use crate::arrays::scalar_fn::ExactScalarFn;
28use crate::arrays::scalar_fn::ScalarFnArrayExt;
29use crate::arrays::scalar_fn::ScalarFnArrayView;
30use crate::builtins::ArrayBuiltins;
31use crate::dtype::DType;
32use crate::dtype::Nullability;
33use crate::kernel::ExecuteParentKernel;
34use crate::scalar::BoolScalar;
35use crate::scalar::Scalar;
36use crate::scalar_fn::fns::binary::Binary;
37use crate::scalar_fn::fns::operators::Operator;
38use crate::validity::Validity;
39
40/// Trait for encoding-specific boolean kernels that operate in encoded space.
41///
42/// Implementations receive the encoded array as the left operand. `rhs` may be any boolean array
43/// encoding or a constant; implementations should return `Ok(None)` when they cannot handle that
44/// operand without falling back to ordinary execution.
45///
46/// Vortex's boolean [`Operator::And`] and [`Operator::Or`] variants use Kleene semantics; there is
47/// no separate two-valued boolean operator path to dispatch here. Consequently, they are not
48/// strict: `false AND null` is `false`, and `true OR null` is `true`.
49pub trait BooleanKernel: VTable {
50    /// Execute `lhs <operator> rhs` using Kleene boolean semantics.
51    fn boolean(
52        lhs: ArrayView<'_, Self>,
53        rhs: &ArrayRef,
54        operator: Operator,
55        ctx: &mut ExecutionCtx,
56    ) -> VortexResult<Option<ArrayRef>>;
57}
58
59/// Adaptor that bridges [`BooleanKernel`] implementations to [`ExecuteParentKernel`].
60///
61/// When a `ScalarFnArray(Binary, And|Or)` wraps a child implementing [`BooleanKernel`], this
62/// adaptor extracts the other operand and delegates to the encoding-specific kernel.
63#[derive(Default, Debug)]
64pub struct BooleanExecuteAdaptor<V>(pub V);
65
66impl<V> ExecuteParentKernel<V> for BooleanExecuteAdaptor<V>
67where
68    V: BooleanKernel,
69{
70    type Parent = ExactScalarFn<Binary>;
71
72    fn execute_parent(
73        &self,
74        array: ArrayView<'_, V>,
75        parent: ScalarFnArrayView<'_, Binary>,
76        child_idx: usize,
77        ctx: &mut ExecutionCtx,
78    ) -> VortexResult<Option<ArrayRef>> {
79        let op = *parent.options;
80        if !is_boolean_operator(op) {
81            return Ok(None);
82        }
83
84        let Some(scalar_fn_array) = parent.as_opt::<ScalarFn>() else {
85            return Ok(None);
86        };
87        let other = match child_idx {
88            0 => scalar_fn_array.get_child(1),
89            1 => scalar_fn_array.get_child(0),
90            _ => return Ok(None),
91        };
92
93        if let Some(result) = constant_boolean(array.array(), other, op)? {
94            return Ok(Some(result));
95        }
96
97        V::boolean(array, other, op, ctx)
98    }
99}
100
101/// Point-wise Kleene logical _and_ between two Boolean arrays.
102#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
103pub fn and_kleene(lhs: &ArrayRef, rhs: &ArrayRef) -> VortexResult<ArrayRef> {
104    lhs.clone().binary(rhs.clone(), Operator::And)
105}
106
107/// Point-wise Kleene logical _or_ between two Boolean arrays.
108#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
109pub fn or_kleene(lhs: &ArrayRef, rhs: &ArrayRef) -> VortexResult<ArrayRef> {
110    lhs.clone().binary(rhs.clone(), Operator::Or)
111}
112
113/// Execute a Kleene boolean operation between two arrays.
114///
115/// This is the entry point for boolean operations from the binary expression.
116/// Handles constants and canonical boolean arrays directly, otherwise falls back to Arrow.
117pub(crate) fn execute_boolean(
118    lhs: ArrayRef,
119    rhs: ArrayRef,
120    op: Operator,
121    ctx: &mut ExecutionCtx,
122) -> VortexResult<ArrayRef> {
123    let nullable = boolean_nullability(&lhs, &rhs);
124
125    if lhs.is_empty() {
126        return Ok(Canonical::empty(&DType::Bool(nullable)).into_array());
127    }
128
129    if let Some(result) = constant_boolean(&lhs, &rhs, op)? {
130        return Ok(result);
131    }
132
133    let lhs = lhs.execute::<BoolArray>(ctx)?;
134    if let Some(result) = <Bool as BooleanKernel>::boolean(lhs.as_view(), &rhs, op, ctx)? {
135        return Ok(result);
136    }
137
138    let rhs = rhs.execute::<BoolArray>(ctx)?;
139    let Some(result) = <Bool as BooleanKernel>::boolean(rhs.as_view(), &lhs.into_array(), op, ctx)?
140    else {
141        vortex_bail!("No boolean kernel for two BoolArrays");
142    };
143    Ok(result)
144}
145
146/// Handles boolean operations where at least one operand is a constant array.
147fn constant_boolean(
148    lhs: &ArrayRef,
149    rhs: &ArrayRef,
150    op: Operator,
151) -> VortexResult<Option<ArrayRef>> {
152    let nullable = boolean_nullability(lhs, rhs);
153
154    match (lhs.as_opt::<Constant>(), rhs.as_opt::<Constant>()) {
155        (Some(lhs), Some(rhs)) => {
156            let result = boolean_scalar_scalar(
157                bool_scalar_value(lhs.scalar())?,
158                bool_scalar_value(rhs.scalar())?,
159                op,
160            )?;
161
162            Ok(Some(constant_bool_result(result, lhs.len(), nullable)))
163        }
164        (Some(lhs), None) => constant_array_boolean(lhs.scalar(), rhs, op, nullable),
165        (None, Some(rhs)) => constant_array_boolean(rhs.scalar(), lhs, op, nullable),
166        (None, None) => Ok(None),
167    }
168}
169
170fn constant_array_boolean(
171    constant: &Scalar,
172    array: &ArrayRef,
173    op: Operator,
174    nullability: Nullability,
175) -> VortexResult<Option<ArrayRef>> {
176    match (op, bool_scalar_value(constant)?) {
177        (Operator::And, Some(false)) => Ok(Some(constant_bool_result(
178            Some(false),
179            array.len(),
180            nullability,
181        ))),
182        (Operator::And, Some(true)) => Ok(Some(cast_bool_nullability(array, nullability)?)),
183        (Operator::Or, Some(true)) => Ok(Some(constant_bool_result(
184            Some(true),
185            array.len(),
186            nullability,
187        ))),
188        (Operator::Or, Some(false)) => Ok(Some(cast_bool_nullability(array, nullability)?)),
189        (Operator::And | Operator::Or, None) => Ok(None),
190        (other, _) => vortex_bail!("Not a boolean operator: {other}"),
191    }
192}
193
194fn boolean_scalar_scalar(
195    lhs: Option<bool>,
196    rhs: Option<bool>,
197    op: Operator,
198) -> VortexResult<Option<bool>> {
199    Ok(match op {
200        Operator::And => match (lhs, rhs) {
201            (Some(false), _) | (_, Some(false)) => Some(false),
202            (None, _) | (_, None) => None,
203            (Some(l), Some(r)) => Some(l & r),
204        },
205        Operator::Or => match (lhs, rhs) {
206            (Some(true), _) | (_, Some(true)) => Some(true),
207            (None, _) | (_, None) => None,
208            (Some(l), Some(r)) => Some(l | r),
209        },
210        other => vortex_bail!("Not a boolean operator: {other}"),
211    })
212}
213
214fn bool_scalar_value(scalar: &Scalar) -> VortexResult<Option<bool>> {
215    Ok(scalar
216        .as_bool_opt()
217        .ok_or_else(|| vortex_err!("expected boolean scalar"))?
218        .value())
219}
220
221/// Execute a Kleene boolean operation from boolean value bitmaps and validity values.
222pub fn kleene_boolean_buffers(
223    lhs_values: BitBuffer,
224    lhs_validity: Validity,
225    rhs_values: BitBuffer,
226    rhs_validity: Validity,
227    operator: Operator,
228    nullability: Nullability,
229    ctx: &mut ExecutionCtx,
230) -> VortexResult<ArrayRef> {
231    let len = lhs_values.len();
232    debug_assert_eq!(rhs_values.len(), len);
233
234    if lhs_validity.definitely_no_nulls() && rhs_validity.definitely_no_nulls() {
235        let values = match operator {
236            Operator::And => lhs_values & &rhs_values,
237            Operator::Or => lhs_values | &rhs_values,
238            other => vortex_bail!("Not a boolean operator: {other}"),
239        };
240        return Ok(BoolArray::try_new(values, Validity::from(nullability))?.into_array());
241    }
242
243    let lhs_valid = lhs_validity.execute_mask(len, ctx)?;
244    let rhs_valid = rhs_validity.execute_mask(len, ctx)?;
245    fused_boolean_buffers(
246        len,
247        &lhs_values,
248        &lhs_valid,
249        &rhs_values,
250        &rhs_valid,
251        operator,
252        nullability,
253        ctx.allocator(),
254    )
255}
256
257/// Execute a Kleene boolean operation between boolean value bits and a scalar.
258pub fn kleene_boolean_buffer_scalar(
259    values: BitBuffer,
260    validity: Validity,
261    scalar: &BoolScalar<'_>,
262    operator: Operator,
263    nullability: Nullability,
264    ctx: &mut ExecutionCtx,
265) -> VortexResult<ArrayRef> {
266    let scalar_value = scalar.value();
267    let len = values.len();
268    let result = match (operator, scalar_value) {
269        (Operator::And, Some(false)) => {
270            return Ok(constant_bool_result(Some(false), len, nullability));
271        }
272        (Operator::And, Some(true)) => {
273            return Ok(
274                BoolArray::try_new(values, validity.union_nullability(nullability))?.into_array(),
275            );
276        }
277        (Operator::Or, Some(true)) => {
278            return Ok(constant_bool_result(Some(true), len, nullability));
279        }
280        (Operator::Or, Some(false)) => {
281            return Ok(
282                BoolArray::try_new(values, validity.union_nullability(nullability))?.into_array(),
283            );
284        }
285        (Operator::And, None) => {
286            let valid = validity
287                .execute_mask(len, ctx)?
288                .bitand_not(&Mask::from_buffer(values));
289            BoolArray::try_new(
290                BitBuffer::new_unset_in(len, ctx.allocator().clone()),
291                Validity::from_mask(valid, nullability),
292            )?
293        }
294        (Operator::Or, None) => {
295            let valid = validity.execute_mask(len, ctx)? & &Mask::from_buffer(values);
296            BoolArray::try_new(
297                BitBuffer::new_set_in(len, ctx.allocator().clone()),
298                Validity::from_mask(valid, nullability),
299            )?
300        }
301        (other, _) => vortex_bail!("Not a boolean operator: {other}"),
302    };
303
304    Ok(result.into_array())
305}
306
307#[expect(
308    clippy::too_many_arguments,
309    reason = "four buffers plus result options"
310)]
311fn fused_boolean_buffers(
312    len: usize,
313    lhs_values: &BitBuffer,
314    lhs_validity: &Mask,
315    rhs_values: &BitBuffer,
316    rhs_validity: &Mask,
317    operator: Operator,
318    nullability: Nullability,
319    allocator: &BufferAllocatorRef,
320) -> VortexResult<ArrayRef> {
321    if let Some(result) = fused_boolean_buffers_aligned(
322        len,
323        lhs_values,
324        lhs_validity,
325        rhs_values,
326        rhs_validity,
327        operator,
328        nullability,
329        allocator,
330    )? {
331        return Ok(result);
332    }
333
334    let n_words = len.div_ceil(64);
335
336    macro_rules! fuse {
337        ($lhs_valid_words:expr, $rhs_valid_words:expr) => {
338            fused_boolean_words(
339                len,
340                lhs_values.chunks().iter_padded(),
341                rhs_values.chunks().iter_padded(),
342                $lhs_valid_words,
343                $rhs_valid_words,
344                operator,
345                nullability,
346                allocator,
347            )
348        };
349    }
350
351    match (lhs_validity.bit_buffer(), rhs_validity.bit_buffer()) {
352        (AllOr::All, AllOr::All) => {
353            fuse!(repeat_n(u64::MAX, n_words), repeat_n(u64::MAX, n_words))
354        }
355        (AllOr::All, AllOr::None) => {
356            fuse!(repeat_n(u64::MAX, n_words), repeat_n(0, n_words))
357        }
358        (AllOr::All, AllOr::Some(rhs_validity)) => fuse!(
359            repeat_n(u64::MAX, n_words),
360            rhs_validity.chunks().iter_padded()
361        ),
362        (AllOr::None, AllOr::All) => {
363            fuse!(repeat_n(0, n_words), repeat_n(u64::MAX, n_words))
364        }
365        (AllOr::None, AllOr::None) => {
366            fuse!(repeat_n(0, n_words), repeat_n(0, n_words))
367        }
368        (AllOr::None, AllOr::Some(rhs_validity)) => {
369            fuse!(repeat_n(0, n_words), rhs_validity.chunks().iter_padded())
370        }
371        (AllOr::Some(lhs_validity), AllOr::All) => fuse!(
372            lhs_validity.chunks().iter_padded(),
373            repeat_n(u64::MAX, n_words)
374        ),
375        (AllOr::Some(lhs_validity), AllOr::None) => {
376            fuse!(lhs_validity.chunks().iter_padded(), repeat_n(0, n_words))
377        }
378        (AllOr::Some(lhs_validity), AllOr::Some(rhs_validity)) => fuse!(
379            lhs_validity.chunks().iter_padded(),
380            rhs_validity.chunks().iter_padded()
381        ),
382    }
383}
384
385#[derive(Clone, Copy)]
386enum WordSource<'a> {
387    Fill(u64),
388    Bytes(&'a [u8]),
389}
390
391impl WordSource<'_> {
392    #[inline]
393    fn word_at(self, byte_offset: usize, len: usize) -> u64 {
394        match self {
395            Self::Fill(word) => word,
396            Self::Bytes(bytes) => read_u64_le(&bytes[byte_offset..byte_offset + len]),
397        }
398    }
399}
400
401#[expect(
402    clippy::too_many_arguments,
403    reason = "four buffers plus result options"
404)]
405fn fused_boolean_buffers_aligned(
406    len: usize,
407    lhs_values: &BitBuffer,
408    lhs_validity: &Mask,
409    rhs_values: &BitBuffer,
410    rhs_validity: &Mask,
411    operator: Operator,
412    nullability: Nullability,
413    allocator: &BufferAllocatorRef,
414) -> VortexResult<Option<ArrayRef>> {
415    let Some(lhs_values) = word_source_from_bit_buffer(lhs_values) else {
416        return Ok(None);
417    };
418    let Some(rhs_values) = word_source_from_bit_buffer(rhs_values) else {
419        return Ok(None);
420    };
421    let Some(lhs_validity) = word_source_from_mask(lhs_validity) else {
422        return Ok(None);
423    };
424    let Some(rhs_validity) = word_source_from_mask(rhs_validity) else {
425        return Ok(None);
426    };
427
428    Ok(Some(fused_boolean_word_sources(
429        len,
430        lhs_values,
431        rhs_values,
432        lhs_validity,
433        rhs_validity,
434        operator,
435        nullability,
436        allocator,
437    )?))
438}
439
440fn word_source_from_bit_buffer(buffer: &BitBuffer) -> Option<WordSource<'_>> {
441    buffer.byte_aligned_bytes().map(WordSource::Bytes)
442}
443
444fn word_source_from_mask(mask: &Mask) -> Option<WordSource<'_>> {
445    match mask.bit_buffer() {
446        AllOr::All => Some(WordSource::Fill(u64::MAX)),
447        AllOr::None => Some(WordSource::Fill(0)),
448        AllOr::Some(buffer) => word_source_from_bit_buffer(buffer),
449    }
450}
451
452#[expect(
453    clippy::too_many_arguments,
454    reason = "four word sources plus result options"
455)]
456fn fused_boolean_word_sources(
457    len: usize,
458    lhs_words: WordSource<'_>,
459    rhs_words: WordSource<'_>,
460    lhs_valid_words: WordSource<'_>,
461    rhs_valid_words: WordSource<'_>,
462    operator: Operator,
463    nullability: Nullability,
464    allocator: &BufferAllocatorRef,
465) -> VortexResult<ArrayRef> {
466    match operator {
467        Operator::And => fused_boolean_and_word_sources(
468            len,
469            lhs_words,
470            rhs_words,
471            lhs_valid_words,
472            rhs_valid_words,
473            nullability,
474            allocator,
475        ),
476        Operator::Or => fused_boolean_or_word_sources(
477            len,
478            lhs_words,
479            rhs_words,
480            lhs_valid_words,
481            rhs_valid_words,
482            nullability,
483            allocator,
484        ),
485        other => vortex_bail!("Not a boolean operator: {other}"),
486    }
487}
488
489fn fused_boolean_and_word_sources(
490    len: usize,
491    lhs_words: WordSource<'_>,
492    rhs_words: WordSource<'_>,
493    lhs_valid_words: WordSource<'_>,
494    rhs_valid_words: WordSource<'_>,
495    nullability: Nullability,
496    allocator: &BufferAllocatorRef,
497) -> VortexResult<ArrayRef> {
498    let n_bytes = len.div_ceil(8);
499    let n_words = n_bytes.div_ceil(8);
500    let full_bytes = n_bytes - n_bytes % 8;
501    let mut values = BufferMut::<u64>::with_capacity_in(n_words, allocator.clone());
502    let mut validity = BufferMut::<u64>::with_capacity_in(n_words, allocator.clone());
503
504    for byte_offset in (0..full_bytes).step_by(8) {
505        let lhs = lhs_words.word_at(byte_offset, 8);
506        let rhs = rhs_words.word_at(byte_offset, 8);
507        let lhs_valid = lhs_valid_words.word_at(byte_offset, 8);
508        let rhs_valid = rhs_valid_words.word_at(byte_offset, 8);
509
510        // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this
511        // loop plus the optional tail push emits at most `n_words` words.
512        unsafe {
513            values.push_unchecked(lhs & rhs);
514            validity
515                .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs));
516        }
517    }
518
519    if full_bytes != n_bytes {
520        let tail_len = n_bytes - full_bytes;
521        let lhs = lhs_words.word_at(full_bytes, tail_len);
522        let rhs = rhs_words.word_at(full_bytes, tail_len);
523        let lhs_valid = lhs_valid_words.word_at(full_bytes, tail_len);
524        let rhs_valid = rhs_valid_words.word_at(full_bytes, tail_len);
525
526        // SAFETY: see the loop safety comment above.
527        unsafe {
528            values.push_unchecked(lhs & rhs);
529            validity
530                .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs));
531        }
532    }
533
534    finish_fused_boolean_words(len, n_bytes, values, validity, nullability)
535}
536
537fn fused_boolean_or_word_sources(
538    len: usize,
539    lhs_words: WordSource<'_>,
540    rhs_words: WordSource<'_>,
541    lhs_valid_words: WordSource<'_>,
542    rhs_valid_words: WordSource<'_>,
543    nullability: Nullability,
544    allocator: &BufferAllocatorRef,
545) -> VortexResult<ArrayRef> {
546    let n_bytes = len.div_ceil(8);
547    let n_words = n_bytes.div_ceil(8);
548    let full_bytes = n_bytes - n_bytes % 8;
549    let mut values = BufferMut::<u64>::with_capacity_in(n_words, allocator.clone());
550    let mut validity = BufferMut::<u64>::with_capacity_in(n_words, allocator.clone());
551
552    for byte_offset in (0..full_bytes).step_by(8) {
553        let lhs = lhs_words.word_at(byte_offset, 8);
554        let rhs = rhs_words.word_at(byte_offset, 8);
555        let lhs_valid = lhs_valid_words.word_at(byte_offset, 8);
556        let rhs_valid = rhs_valid_words.word_at(byte_offset, 8);
557
558        // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this
559        // loop plus the optional tail push emits at most `n_words` words.
560        unsafe {
561            values.push_unchecked(lhs | rhs);
562            validity
563                .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs));
564        }
565    }
566
567    if full_bytes != n_bytes {
568        let tail_len = n_bytes - full_bytes;
569        let lhs = lhs_words.word_at(full_bytes, tail_len);
570        let rhs = rhs_words.word_at(full_bytes, tail_len);
571        let lhs_valid = lhs_valid_words.word_at(full_bytes, tail_len);
572        let rhs_valid = rhs_valid_words.word_at(full_bytes, tail_len);
573
574        // SAFETY: see the loop safety comment above.
575        unsafe {
576            values.push_unchecked(lhs | rhs);
577            validity
578                .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs));
579        }
580    }
581
582    finish_fused_boolean_words(len, n_bytes, values, validity, nullability)
583}
584
585fn finish_fused_boolean_words(
586    len: usize,
587    n_bytes: usize,
588    values: BufferMut<u64>,
589    validity: BufferMut<u64>,
590    nullability: Nullability,
591) -> VortexResult<ArrayRef> {
592    let mut values = values.into_byte_buffer();
593    values.truncate(n_bytes);
594    let mut validity = validity.into_byte_buffer();
595    validity.truncate(n_bytes);
596    Ok(BoolArray::try_new(
597        BitBuffer::new(values.freeze(), len),
598        Validity::from_mask(
599            Mask::from_buffer(BitBuffer::new(validity.freeze(), len)),
600            nullability,
601        ),
602    )?
603    .into_array())
604}
605
606#[expect(
607    clippy::too_many_arguments,
608    reason = "four word iterators plus result options"
609)]
610fn fused_boolean_words<L, R, LV, RV>(
611    len: usize,
612    lhs_words: L,
613    rhs_words: R,
614    lhs_valid_words: LV,
615    rhs_valid_words: RV,
616    operator: Operator,
617    nullability: Nullability,
618    allocator: &BufferAllocatorRef,
619) -> VortexResult<ArrayRef>
620where
621    L: Iterator<Item = u64>,
622    R: Iterator<Item = u64>,
623    LV: Iterator<Item = u64>,
624    RV: Iterator<Item = u64>,
625{
626    match operator {
627        Operator::And => fused_boolean_and_words(
628            len,
629            lhs_words,
630            rhs_words,
631            lhs_valid_words,
632            rhs_valid_words,
633            nullability,
634            allocator,
635        ),
636        Operator::Or => fused_boolean_or_words(
637            len,
638            lhs_words,
639            rhs_words,
640            lhs_valid_words,
641            rhs_valid_words,
642            nullability,
643            allocator,
644        ),
645        other => vortex_bail!("Not a boolean operator: {other}"),
646    }
647}
648
649fn fused_boolean_and_words<L, R, LV, RV>(
650    len: usize,
651    lhs_words: L,
652    rhs_words: R,
653    lhs_valid_words: LV,
654    rhs_valid_words: RV,
655    nullability: Nullability,
656    allocator: &BufferAllocatorRef,
657) -> VortexResult<ArrayRef>
658where
659    L: Iterator<Item = u64>,
660    R: Iterator<Item = u64>,
661    LV: Iterator<Item = u64>,
662    RV: Iterator<Item = u64>,
663{
664    let n_words = len.div_ceil(64);
665    let mut values = BufferMut::<u64>::with_capacity_in(n_words, allocator.clone());
666    let mut validity = BufferMut::<u64>::with_capacity_in(n_words, allocator.clone());
667
668    for (((lhs, rhs), lhs_valid), rhs_valid) in lhs_words
669        .zip(rhs_words)
670        .zip(lhs_valid_words)
671        .zip(rhs_valid_words)
672        .take(n_words)
673    {
674        // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this loop is
675        // capped at `n_words`.
676        unsafe {
677            values.push_unchecked(lhs & rhs);
678            validity
679                .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs));
680        }
681    }
682
683    finish_fused_boolean_words(len, len.div_ceil(8), values, validity, nullability)
684}
685
686fn fused_boolean_or_words<L, R, LV, RV>(
687    len: usize,
688    lhs_words: L,
689    rhs_words: R,
690    lhs_valid_words: LV,
691    rhs_valid_words: RV,
692    nullability: Nullability,
693    allocator: &BufferAllocatorRef,
694) -> VortexResult<ArrayRef>
695where
696    L: Iterator<Item = u64>,
697    R: Iterator<Item = u64>,
698    LV: Iterator<Item = u64>,
699    RV: Iterator<Item = u64>,
700{
701    let n_words = len.div_ceil(64);
702    let mut values = BufferMut::<u64>::with_capacity_in(n_words, allocator.clone());
703    let mut validity = BufferMut::<u64>::with_capacity_in(n_words, allocator.clone());
704
705    for (((lhs, rhs), lhs_valid), rhs_valid) in lhs_words
706        .zip(rhs_words)
707        .zip(lhs_valid_words)
708        .zip(rhs_valid_words)
709        .take(n_words)
710    {
711        // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this loop is
712        // capped at `n_words`.
713        unsafe {
714            values.push_unchecked(lhs | rhs);
715            validity
716                .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs));
717        }
718    }
719
720    finish_fused_boolean_words(len, len.div_ceil(8), values, validity, nullability)
721}
722
723fn constant_bool_result(value: Option<bool>, len: usize, nullability: Nullability) -> ArrayRef {
724    let scalar = value
725        .map(|b| Scalar::bool(b, nullability))
726        .unwrap_or_else(|| Scalar::null(DType::Bool(nullability)));
727
728    ConstantArray::new(scalar, len).into_array()
729}
730
731fn cast_bool_nullability(array: &ArrayRef, nullability: Nullability) -> VortexResult<ArrayRef> {
732    let dtype = DType::Bool(nullability);
733    if array.dtype() == &dtype {
734        Ok(array.clone())
735    } else {
736        array.cast(dtype)
737    }
738}
739
740fn boolean_nullability(lhs: &ArrayRef, rhs: &ArrayRef) -> Nullability {
741    lhs.dtype().nullability() | rhs.dtype().nullability()
742}
743
744#[inline]
745fn is_boolean_operator(operator: Operator) -> bool {
746    matches!(operator, Operator::And | Operator::Or)
747}
748
749#[cfg(test)]
750mod tests {
751    use rstest::rstest;
752    use vortex_error::VortexResult;
753
754    use crate::ArrayRef;
755    use crate::IntoArray;
756    use crate::VortexSessionExecute;
757    use crate::array_session;
758    use crate::arrays::BoolArray;
759    use crate::arrays::ConstantArray;
760    use crate::assert_arrays_eq;
761    use crate::builtins::ArrayBuiltins;
762    use crate::dtype::DType;
763    use crate::dtype::Nullability;
764    use crate::scalar::Scalar;
765    use crate::scalar_fn::fns::operators::Operator;
766
767    #[test]
768    fn test_kleene_truth_table() -> VortexResult<()> {
769        let mut ctx = array_session().create_execution_ctx();
770        let lhs = BoolArray::from_iter([
771            Some(true),
772            Some(true),
773            Some(true),
774            Some(false),
775            Some(false),
776            Some(false),
777            None,
778            None,
779            None,
780        ])
781        .into_array();
782        let rhs = BoolArray::from_iter([
783            Some(true),
784            Some(false),
785            None,
786            Some(true),
787            Some(false),
788            None,
789            Some(true),
790            Some(false),
791            None,
792        ])
793        .into_array();
794
795        assert_arrays_eq!(
796            lhs.binary(rhs.clone(), Operator::And)?,
797            BoolArray::from_iter([
798                Some(true),
799                Some(false),
800                None,
801                Some(false),
802                Some(false),
803                Some(false),
804                None,
805                Some(false),
806                None,
807            ]),
808            &mut ctx
809        );
810
811        assert_arrays_eq!(
812            lhs.binary(rhs, Operator::Or)?,
813            BoolArray::from_iter([
814                Some(true),
815                Some(true),
816                Some(true),
817                Some(true),
818                Some(false),
819                None,
820                Some(true),
821                None,
822                None,
823            ]),
824            &mut ctx
825        );
826
827        Ok(())
828    }
829
830    #[test]
831    fn test_null_constant_kleene() -> VortexResult<()> {
832        let mut ctx = array_session().create_execution_ctx();
833        let lhs = BoolArray::from_iter([Some(false), Some(true), None]).into_array();
834        let null = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), lhs.len())
835            .into_array();
836
837        assert_arrays_eq!(
838            lhs.binary(null.clone(), Operator::And)?,
839            BoolArray::from_iter([Some(false), None, None]),
840            &mut ctx
841        );
842        assert_arrays_eq!(
843            lhs.binary(null, Operator::Or)?,
844            BoolArray::from_iter([None, Some(true), None]),
845            &mut ctx
846        );
847
848        Ok(())
849    }
850
851    #[rstest]
852    #[case(
853        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(),
854        BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)]).into_array(),
855    )]
856    #[case(
857        BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)]).into_array(),
858        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(),
859    )]
860    fn test_or(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
861        let mut ctx = array_session().create_execution_ctx();
862        let r = lhs.binary(rhs, Operator::Or).unwrap();
863        let r = r.execute::<BoolArray>(&mut ctx).unwrap().into_array();
864
865        let v0 = r
866            .execute_scalar(0, &mut array_session().create_execution_ctx())
867            .unwrap()
868            .as_bool()
869            .value();
870        let v1 = r
871            .execute_scalar(1, &mut array_session().create_execution_ctx())
872            .unwrap()
873            .as_bool()
874            .value();
875        let v2 = r
876            .execute_scalar(2, &mut array_session().create_execution_ctx())
877            .unwrap()
878            .as_bool()
879            .value();
880        let v3 = r
881            .execute_scalar(3, &mut array_session().create_execution_ctx())
882            .unwrap()
883            .as_bool()
884            .value();
885
886        assert!(v0.unwrap());
887        assert!(v1.unwrap());
888        assert!(v2.unwrap());
889        assert!(!v3.unwrap());
890    }
891
892    #[rstest]
893    #[case(
894        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(),
895        BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)]).into_array(),
896    )]
897    #[case(
898        BoolArray::from_iter([Some(true), Some(false), Some(true), Some(false)]).into_array(),
899        BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(),
900    )]
901    fn test_and(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) {
902        let mut ctx = array_session().create_execution_ctx();
903        let r = lhs
904            .binary(rhs, Operator::And)
905            .unwrap()
906            .execute::<BoolArray>(&mut ctx)
907            .unwrap()
908            .into_array();
909
910        let v0 = r
911            .execute_scalar(0, &mut array_session().create_execution_ctx())
912            .unwrap()
913            .as_bool()
914            .value();
915        let v1 = r
916            .execute_scalar(1, &mut array_session().create_execution_ctx())
917            .unwrap()
918            .as_bool()
919            .value();
920        let v2 = r
921            .execute_scalar(2, &mut array_session().create_execution_ctx())
922            .unwrap()
923            .as_bool()
924            .value();
925        let v3 = r
926            .execute_scalar(3, &mut array_session().create_execution_ctx())
927            .unwrap()
928            .as_bool()
929            .value();
930
931        assert!(v0.unwrap());
932        assert!(!v1.unwrap());
933        assert!(!v2.unwrap());
934        assert!(!v3.unwrap());
935    }
936}