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

vortex_array/scalar_fn/fns/binary/
boolean.rs
