1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

use vortex_dtype::DType;
use vortex_error::VortexResult;

use crate::array::{ChunkedArray, ChunkedEncoding};
use crate::compute::{binary_boolean, slice, BinaryBooleanFn, BinaryOperator};
use crate::{ArrayDType, ArrayData, IntoArrayData};

impl BinaryBooleanFn<ChunkedArray> for ChunkedEncoding {
    fn binary_boolean(
        &self,
        lhs: &ChunkedArray,
        rhs: &ArrayData,
        op: BinaryOperator,
    ) -> VortexResult<Option<ArrayData>> {
        let mut idx = 0;
        let mut chunks = Vec::with_capacity(lhs.nchunks());

        for chunk in lhs.chunks() {
            let sliced = slice(rhs, idx, idx + chunk.len())?;
            let result = binary_boolean(&chunk, &sliced, op)?;
            chunks.push(result);
            idx += chunk.len();
        }

        let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
        let dtype = DType::Bool(nullable.into());
        Ok(Some(ChunkedArray::try_new(chunks, dtype)?.into_array()))
    }
}

#[cfg(test)]
mod tests {
    use vortex_dtype::{DType, Nullability};

    use crate::array::{BoolArray, ChunkedArray};
    use crate::compute::{binary_boolean, BinaryOperator};
    use crate::{IntoArrayData, IntoArrayVariant};

    #[test]
    fn test_bin_bool_chunked() {
        let arr0 = BoolArray::from_iter(vec![true, false]).into_array();
        let arr1 = BoolArray::from_iter(vec![false, false, true]).into_array();
        let chunked1 =
            ChunkedArray::try_new(vec![arr0, arr1], DType::Bool(Nullability::NonNullable)).unwrap();

        let arr2 = BoolArray::from_iter(vec![Some(false), Some(true)]).into_array();
        let arr3 = BoolArray::from_iter(vec![Some(false), None, Some(false)]).into_array();
        let chunked2 =
            ChunkedArray::try_new(vec![arr2, arr3], DType::Bool(Nullability::Nullable)).unwrap();

        assert_eq!(
            binary_boolean(
                &chunked1.into_array(),
                &chunked2.into_array(),
                BinaryOperator::Or
            )
            .unwrap()
            .into_bool()
            .unwrap()
            .boolean_buffer(),
            vec![true, true, false, false, true].into()
        );
    }
}