vortex_array/arrays/bool/
patch.rs

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