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

vortex_bytebool/
compute.rs
