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;
6
7use crate::ToCanonical;
8use crate::arrays::BoolArray;
9use crate::patches::Patches;
10use crate::vtable::ValidityHelper;
11
12impl BoolArray {
13    pub fn patch(self, patches: &Patches) -> Self {
14        let len = self.len();
15        let offset = patches.offset();
16        let indices = patches.indices().to_primitive();
17        let values = patches.values().to_bool();
18
19        let patched_validity =
20            self.validity()
21                .clone()
22                .patch(len, offset, indices.as_ref(), values.validity());
23
24        let (mut own_values, bit_offset) = self.into_boolean_builder();
25        match_each_integer_ptype!(indices.ptype(), |I| {
26            for (idx, value) in indices
27                .as_slice::<I>()
28                .iter()
29                .zip_eq(values.boolean_buffer().iter())
30            {
31                #[allow(clippy::cast_possible_truncation)]
32                own_values.set_bit(*idx as usize - offset + bit_offset, value);
33            }
34        });
35
36        Self::new(own_values.finish().slice(bit_offset, len), patched_validity)
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);
51        let (values, offset) = sliced.to_bool().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);
61        let (values, offset) = sliced.to_bool().into_boolean_builder();
62        assert_eq!(offset, 4);
63        assert_eq!(values.as_slice(), &[255, 127]);
64    }
65}