vortex_array/arrays/bool/
patch.rs

1use itertools::Itertools;
2use vortex_dtype::match_each_integer_ptype;
3use vortex_error::VortexResult;
4
5use crate::ToCanonical;
6use crate::arrays::BoolArray;
7use crate::patches::Patches;
8use crate::vtable::ValidityHelper;
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.as_ref(), 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
47    #[test]
48    fn patch_sliced_bools() {
49        let arr = BoolArray::from(BooleanBuffer::new_set(12));
50        let sliced = arr.slice(4, 12).unwrap();
51        let (values, offset) = sliced.to_bool().unwrap().into_boolean_builder();
52        assert_eq!(offset, 4);
53        assert_eq!(values.len(), 12);
54        assert_eq!(values.as_slice(), &[255, 15]);
55    }
56
57    #[test]
58    fn patch_sliced_bools_offset() {
59        let arr = BoolArray::from(BooleanBuffer::new_set(15));
60        let sliced = arr.slice(4, 15).unwrap();
61        let (values, offset) = sliced.to_bool().unwrap().into_boolean_builder();
62        assert_eq!(offset, 4);
63        assert_eq!(values.as_slice(), &[255, 127]);
64    }
65}