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                #[allow(clippy::cast_possible_truncation)]
30                own_values.set_bit(*idx as usize - offset + bit_offset, value);
31            }
32        });
33
34        Ok(Self::new(
35            own_values.finish().slice(bit_offset, len),
36            patched_validity,
37        ))
38    }
39}
40
41#[cfg(test)]
42mod tests {
43    use arrow_buffer::BooleanBuffer;
44
45    use crate::ToCanonical;
46    use crate::arrays::BoolArray;
47
48    #[test]
49    fn patch_sliced_bools() {
50        let arr = BoolArray::from(BooleanBuffer::new_set(12));
51        let sliced = arr.slice(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 = arr.slice(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}