vortex_array/arrays/chunked/compute/
boolean.rs

1use vortex_dtype::DType;
2use vortex_error::VortexResult;
3
4use crate::arrays::{ChunkedArray, ChunkedEncoding};
5use crate::compute::{BinaryBooleanFn, BinaryOperator, binary_boolean, slice};
6use crate::{Array, ArrayRef};
7
8impl BinaryBooleanFn<&ChunkedArray> for ChunkedEncoding {
9    fn binary_boolean(
10        &self,
11        lhs: &ChunkedArray,
12        rhs: &dyn Array,
13        op: BinaryOperator,
14    ) -> VortexResult<Option<ArrayRef>> {
15        let mut idx = 0;
16        let mut chunks = Vec::with_capacity(lhs.nchunks());
17
18        for chunk in lhs.non_empty_chunks() {
19            let sliced = slice(rhs, idx, idx + chunk.len())?;
20            let result = binary_boolean(chunk, &sliced, op)?;
21            chunks.push(result);
22            idx += chunk.len();
23        }
24
25        let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
26        let dtype = DType::Bool(nullable.into());
27        Ok(Some(ChunkedArray::try_new(chunks, dtype)?.into_array()))
28    }
29}
30
31#[cfg(test)]
32mod tests {
33    use vortex_dtype::{DType, Nullability};
34
35    use crate::array::Array;
36    use crate::arrays::{BoolArray, BooleanBuffer, ChunkedArray};
37    use crate::canonical::ToCanonical;
38    use crate::compute::{BinaryOperator, binary_boolean};
39
40    #[test]
41    fn test_bin_bool_chunked() {
42        let arr0 = BoolArray::from_iter(vec![true, false]).to_array();
43        let arr1 = BoolArray::from_iter(vec![false, false, true]).to_array();
44        let chunked1 =
45            ChunkedArray::try_new(vec![arr0, arr1], DType::Bool(Nullability::NonNullable)).unwrap();
46
47        let arr2 = BoolArray::from_iter(vec![Some(false), Some(true)]).to_array();
48        let arr3 = BoolArray::from_iter(vec![Some(false), None, Some(false)]).to_array();
49        let chunked2 =
50            ChunkedArray::try_new(vec![arr2, arr3], DType::Bool(Nullability::Nullable)).unwrap();
51
52        let result = binary_boolean(&chunked1, &chunked2, BinaryOperator::Or)
53            .unwrap()
54            .to_bool()
55            .unwrap();
56        assert_eq!(
57            result.boolean_buffer(),
58            &BooleanBuffer::from_iter([true, true, false, false, true])
59        );
60    }
61}
vortex_array/arrays/chunked/compute/boolean.rs

vortex_array/arrays/chunked/compute/
boolean.rs
