vortex_bytebool/
compute.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use num_traits::AsPrimitive;
5use vortex_array::compute::{
6    CastKernel, CastKernelAdapter, MaskKernel, MaskKernelAdapter, TakeKernel, TakeKernelAdapter,
7};
8use vortex_array::vtable::ValidityHelper;
9use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical, register_kernel};
10use vortex_dtype::{DType, match_each_integer_ptype};
11use vortex_error::VortexResult;
12use vortex_mask::Mask;
13
14use super::{ByteBoolArray, ByteBoolVTable};
15
16impl CastKernel for ByteBoolVTable {
17    fn cast(&self, array: &ByteBoolArray, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
18        // ByteBool is essentially a bool array stored as bytes
19        // The main difference from BoolArray is the storage format
20        // For casting, we can decode to canonical (BoolArray) and let it handle the cast
21
22        // If just changing nullability, we can optimize
23        if array.dtype().eq_ignore_nullability(dtype) {
24            let new_validity = array
25                .validity()
26                .clone()
27                .cast_nullability(dtype.nullability(), array.len())?;
28
29            return Ok(Some(
30                ByteBoolArray::new(array.buffer().clone(), new_validity).into_array(),
31            ));
32        }
33
34        // For other casts, decode to canonical and let BoolArray handle it
35        Ok(None)
36    }
37}
38
39register_kernel!(CastKernelAdapter(ByteBoolVTable).lift());
40
41impl MaskKernel for ByteBoolVTable {
42    fn mask(&self, array: &ByteBoolArray, mask: &Mask) -> VortexResult<ArrayRef> {
43        Ok(ByteBoolArray::new(array.buffer().clone(), array.validity().mask(mask)).into_array())
44    }
45}
46
47register_kernel!(MaskKernelAdapter(ByteBoolVTable).lift());
48
49impl TakeKernel for ByteBoolVTable {
50    fn take(&self, array: &ByteBoolArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
51        let indices = indices.to_primitive();
52        let bools = array.as_slice();
53
54        // This handles combining validity from both source array and nullable indices
55        let validity = array.validity().take(indices.as_ref())?;
56
57        let taken_bools = match_each_integer_ptype!(indices.ptype(), |I| {
58            indices
59                .as_slice::<I>()
60                .iter()
61                .map(|&idx| {
62                    let idx: usize = idx.as_();
63                    bools[idx]
64                })
65                .collect::<Vec<bool>>()
66        });
67
68        Ok(ByteBoolArray::from_vec(taken_bools, validity).into_array())
69    }
70}
71
72register_kernel!(TakeKernelAdapter(ByteBoolVTable).lift());
73
74#[cfg(test)]
75mod tests {
76    use rstest::rstest;
77    use vortex_array::assert_arrays_eq;
78    use vortex_array::compute::conformance::filter::test_filter_conformance;
79    use vortex_array::compute::conformance::mask::test_mask_conformance;
80    use vortex_array::compute::conformance::take::test_take_conformance;
81    use vortex_array::compute::{Operator, compare};
82
83    use super::*;
84
85    #[test]
86    fn test_slice() {
87        let original = vec![Some(true), Some(true), None, Some(false), None];
88        let vortex_arr = ByteBoolArray::from(original);
89
90        let sliced_arr = vortex_arr.slice(1..4);
91
92        let expected = ByteBoolArray::from(vec![Some(true), None, Some(false)]);
93        assert_arrays_eq!(sliced_arr, expected.to_array());
94    }
95
96    #[test]
97    fn test_compare_all_equal() {
98        let lhs = ByteBoolArray::from(vec![true; 5]);
99        let rhs = ByteBoolArray::from(vec![true; 5]);
100
101        let arr = compare(lhs.as_ref(), rhs.as_ref(), Operator::Eq).unwrap();
102
103        let expected = ByteBoolArray::from(vec![true; 5]);
104        assert_arrays_eq!(arr, expected.to_array());
105    }
106
107    #[test]
108    fn test_compare_all_different() {
109        let lhs = ByteBoolArray::from(vec![false; 5]);
110        let rhs = ByteBoolArray::from(vec![true; 5]);
111
112        let arr = compare(lhs.as_ref(), rhs.as_ref(), Operator::Eq).unwrap();
113
114        let expected = ByteBoolArray::from(vec![false; 5]);
115        assert_arrays_eq!(arr, expected.to_array());
116    }
117
118    #[test]
119    fn test_compare_with_nulls() {
120        let lhs = ByteBoolArray::from(vec![true; 5]);
121        let rhs = ByteBoolArray::from(vec![Some(true), Some(true), Some(true), Some(false), None]);
122
123        let arr = compare(lhs.as_ref(), rhs.as_ref(), Operator::Eq).unwrap();
124
125        let expected =
126            ByteBoolArray::from(vec![Some(true), Some(true), Some(true), Some(false), None]);
127        assert_arrays_eq!(arr, expected.to_array());
128    }
129
130    #[test]
131    fn test_mask_byte_bool() {
132        test_mask_conformance(ByteBoolArray::from(vec![true, false, true, true, false]).as_ref());
133        test_mask_conformance(
134            ByteBoolArray::from(vec![Some(true), Some(true), None, Some(false), None]).as_ref(),
135        );
136    }
137
138    #[test]
139    fn test_filter_byte_bool() {
140        test_filter_conformance(ByteBoolArray::from(vec![true, false, true, true, false]).as_ref());
141        test_filter_conformance(
142            ByteBoolArray::from(vec![Some(true), Some(true), None, Some(false), None]).as_ref(),
143        );
144    }
145
146    #[rstest]
147    #[case(ByteBoolArray::from(vec![true, false, true, true, false]))]
148    #[case(ByteBoolArray::from(vec![Some(true), Some(true), None, Some(false), None]))]
149    #[case(ByteBoolArray::from(vec![true, false]))]
150    #[case(ByteBoolArray::from(vec![true]))]
151    fn test_take_byte_bool_conformance(#[case] array: ByteBoolArray) {
152        test_take_conformance(array.as_ref());
153    }
154
155    use vortex_array::compute::cast;
156    use vortex_array::compute::conformance::cast::test_cast_conformance;
157    use vortex_array::compute::conformance::consistency::test_array_consistency;
158    use vortex_dtype::{DType, Nullability};
159
160    #[test]
161    fn test_cast_bytebool_to_nullable() {
162        let array = ByteBoolArray::from(vec![true, false, true, false]);
163        let casted = cast(array.as_ref(), &DType::Bool(Nullability::Nullable)).unwrap();
164        assert_eq!(casted.dtype(), &DType::Bool(Nullability::Nullable));
165        assert_eq!(casted.len(), 4);
166    }
167
168    #[rstest]
169    #[case(ByteBoolArray::from(vec![true, false, true, true, false]))]
170    #[case(ByteBoolArray::from(vec![Some(true), Some(false), None, Some(true), None]))]
171    #[case(ByteBoolArray::from(vec![false]))]
172    #[case(ByteBoolArray::from(vec![true]))]
173    #[case(ByteBoolArray::from(vec![Some(true), None]))]
174    fn test_cast_bytebool_conformance(#[case] array: ByteBoolArray) {
175        test_cast_conformance(array.as_ref());
176    }
177
178    #[rstest]
179    #[case::non_nullable(ByteBoolArray::from(vec![true, false, true, true, false]))]
180    #[case::nullable(ByteBoolArray::from(vec![Some(true), Some(false), None, Some(true), None]))]
181    #[case::all_true(ByteBoolArray::from(vec![true, true, true, true]))]
182    #[case::all_false(ByteBoolArray::from(vec![false, false, false, false]))]
183    #[case::single_true(ByteBoolArray::from(vec![true]))]
184    #[case::single_false(ByteBoolArray::from(vec![false]))]
185    #[case::single_null(ByteBoolArray::from(vec![None]))]
186    #[case::mixed_with_nulls(ByteBoolArray::from(vec![Some(true), None, Some(false), None, Some(true)]))]
187    fn test_bytebool_consistency(#[case] array: ByteBoolArray) {
188        test_array_consistency(array.as_ref());
189    }
190}
vortex_bytebool/compute.rs

vortex_bytebool/
compute.rs
