vortex_array/arrays/bool/
mod.rs

1use arrow_array::BooleanArray;
2use arrow_buffer::MutableBuffer;
3
4use crate::validity::Validity;
5
6mod array;
7pub mod compute;
8mod patch;
9mod serde;
10mod stats;
11
12pub use array::*;
13// Re-export the BooleanBuffer type on our API surface.
14pub use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
15
16impl BoolArray {
17    /// Create a new BoolArray from a set of indices and a length.
18    /// All indices must be less than the length.
19    pub fn from_indices<I: IntoIterator<Item = usize>>(length: usize, indices: I) -> Self {
20        let mut buffer = MutableBuffer::new_null(length);
21        indices
22            .into_iter()
23            .for_each(|idx| arrow_buffer::bit_util::set_bit(&mut buffer, idx));
24        Self::new(
25            BooleanBufferBuilder::new_from_buffer(buffer, length).finish(),
26            Validity::NonNullable,
27        )
28    }
29}
30
31impl From<BooleanBuffer> for BoolArray {
32    fn from(value: BooleanBuffer) -> Self {
33        Self::new(value, Validity::NonNullable)
34    }
35}
36
37impl FromIterator<bool> for BoolArray {
38    fn from_iter<T: IntoIterator<Item = bool>>(iter: T) -> Self {
39        Self::new(BooleanBuffer::from_iter(iter), Validity::NonNullable)
40    }
41}
42
43impl FromIterator<Option<bool>> for BoolArray {
44    fn from_iter<I: IntoIterator<Item = Option<bool>>>(iter: I) -> Self {
45        let (buffer, nulls) = BooleanArray::from_iter(iter).into_parts();
46        Self::new(
47            buffer,
48            nulls.map(Validity::from).unwrap_or(Validity::AllValid),
49        )
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
56    use vortex_buffer::buffer;
57
58    use crate::ToCanonical;
59    use crate::array::Array;
60    use crate::arrays::{BoolArray, PrimitiveArray};
61    use crate::compute::test_harness::test_mask;
62    use crate::compute::{scalar_at, slice};
63    use crate::patches::Patches;
64    use crate::validity::Validity;
65
66    #[test]
67    fn bool_array() {
68        let arr = BoolArray::from_iter([true, false, true]);
69        let scalar = bool::try_from(&scalar_at(&arr, 0).unwrap()).unwrap();
70        assert!(scalar);
71    }
72
73    #[test]
74    fn test_all_some_iter() {
75        let arr = BoolArray::from_iter([Some(true), Some(false)]);
76
77        assert!(matches!(arr.validity(), Validity::AllValid));
78
79        let scalar = bool::try_from(&scalar_at(&arr, 0).unwrap()).unwrap();
80        assert!(scalar);
81        let scalar = bool::try_from(&scalar_at(&arr, 1).unwrap()).unwrap();
82        assert!(!scalar);
83    }
84
85    #[test]
86    fn test_bool_from_iter() {
87        let arr = BoolArray::from_iter([Some(true), Some(true), None, Some(false), None]);
88
89        let scalar = bool::try_from(&scalar_at(&arr, 0).unwrap()).unwrap();
90        assert!(scalar);
91
92        let scalar = bool::try_from(&scalar_at(&arr, 1).unwrap()).unwrap();
93        assert!(scalar);
94
95        let scalar = scalar_at(&arr, 2).unwrap();
96        assert!(scalar.is_null());
97
98        let scalar = bool::try_from(&scalar_at(&arr, 3).unwrap()).unwrap();
99        assert!(!scalar);
100
101        let scalar = scalar_at(&arr, 4).unwrap();
102        assert!(scalar.is_null());
103    }
104
105    #[test]
106    fn patch_sliced_bools() {
107        let arr = {
108            let mut builder = BooleanBufferBuilder::new(12);
109            builder.append(false);
110            builder.append_n(11, true);
111            BoolArray::from(builder.finish())
112        };
113        let sliced = slice(&arr, 4, 12).unwrap();
114        let sliced_len = sliced.len();
115        let (values, offset) = sliced.to_bool().unwrap().into_boolean_builder();
116        assert_eq!(offset, 4);
117        assert_eq!(values.as_slice(), &[254, 15]);
118
119        // patch the underlying array
120        let patches = Patches::new(
121            arr.len(),
122            0,
123            PrimitiveArray::new(buffer![4u32], Validity::AllValid).into_array(),
124            BoolArray::from(BooleanBuffer::new_unset(1)).into_array(),
125        );
126        let arr = arr.patch(&patches).unwrap();
127        let arr_len = arr.len();
128        let (values, offset) = arr.to_bool().unwrap().into_boolean_builder();
129        assert_eq!(offset, 0);
130        assert_eq!(values.len(), arr_len + offset);
131        assert_eq!(values.as_slice(), &[238, 15]);
132
133        // the slice should be unchanged
134        let (values, offset) = sliced.to_bool().unwrap().into_boolean_builder();
135        assert_eq!(offset, 4);
136        assert_eq!(values.len(), sliced_len + offset);
137        assert_eq!(values.as_slice(), &[254, 15]); // unchanged
138    }
139
140    #[test]
141    fn slice_array_in_middle() {
142        let arr = BoolArray::from(BooleanBuffer::new_set(16));
143        let sliced = slice(&arr, 4, 12).unwrap();
144        let sliced_len = sliced.len();
145        let (values, offset) = sliced.to_bool().unwrap().into_boolean_builder();
146        assert_eq!(offset, 4);
147        assert_eq!(values.len(), sliced_len + offset);
148        assert_eq!(values.as_slice(), &[255, 15]);
149    }
150
151    #[test]
152    #[should_panic]
153    fn patch_bools_owned() {
154        let buffer = buffer![255u8; 2];
155        let buf = BooleanBuffer::new(buffer.into_arrow_buffer(), 0, 15);
156        let arr = BoolArray::new(buf, Validity::NonNullable);
157        let buf_ptr = arr.boolean_buffer().sliced().as_ptr();
158
159        let patches = Patches::new(
160            arr.len(),
161            0,
162            PrimitiveArray::new(buffer![0u32], Validity::AllValid).into_array(),
163            BoolArray::from(BooleanBuffer::new_unset(1)).into_array(),
164        );
165        let arr = arr.patch(&patches).unwrap();
166        assert_eq!(arr.boolean_buffer().sliced().as_ptr(), buf_ptr);
167
168        let (values, _byte_bit_offset) = arr.to_bool().unwrap().into_boolean_builder();
169        assert_eq!(values.as_slice(), &[254, 127]);
170    }
171
172    #[test]
173    fn test_mask_primitive_array() {
174        test_mask(&BoolArray::from_iter([true, false, true, true, false]));
175    }
176}
vortex_array/arrays/bool/mod.rs

vortex_array/arrays/bool/
mod.rs
