vortex_array/arrays/bool/
patch.rs

1use itertools::Itertools;
2use vortex_dtype::match_each_integer_ptype;
3use vortex_error::VortexResult;
4
5use crate::arrays::BoolArray;
6use crate::patches::Patches;
7use crate::variants::PrimitiveArrayTrait;
8use crate::{Array, ToCanonical};
9
10impl BoolArray {
11    pub fn patch(self, patches: &Patches) -> VortexResult<Self> {
12        let len = self.len();
13        let (_, offset, indices, values) = patches.clone().into_parts();
14        let indices = indices.to_primitive()?;
15        let values = values.to_bool()?;
16
17        let patched_validity =
18            self.validity()
19                .clone()
20                .patch(len, offset, &indices, values.validity())?;
21
22        let (mut own_values, bit_offset) = self.into_boolean_builder();
23        match_each_integer_ptype!(indices.ptype(), |$I| {
24            for (idx, value) in indices
25                .as_slice::<$I>()
26                .iter()
27                .zip_eq(values.boolean_buffer().iter())
28            {
29                own_values.set_bit(*idx as usize - offset + bit_offset, value);
30            }
31        });
32
33        Ok(Self::new(
34            own_values.finish().slice(bit_offset, len),
35            patched_validity,
36        ))
37    }
38}
39
40#[cfg(test)]
41mod tests {
42    use arrow_buffer::BooleanBuffer;
43
44    use crate::ToCanonical;
45    use crate::arrays::BoolArray;
46    use crate::compute::slice;
47
48    #[test]
49    fn patch_sliced_bools() {
50        let arr = BoolArray::from(BooleanBuffer::new_set(12));
51        let sliced = slice(&arr, 4, 12).unwrap();
52        let (values, offset) = sliced.to_bool().unwrap().into_boolean_builder();
53        assert_eq!(offset, 4);
54        assert_eq!(values.len(), 12);
55        assert_eq!(values.as_slice(), &[255, 15]);
56    }
57
58    #[test]
59    fn patch_sliced_bools_offset() {
60        let arr = BoolArray::from(BooleanBuffer::new_set(15));
61        let sliced = slice(&arr, 4, 15).unwrap();
62        let (values, offset) = sliced.to_bool().unwrap().into_boolean_builder();
63        assert_eq!(offset, 4);
64        assert_eq!(values.as_slice(), &[255, 127]);
65    }
66}