1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72

use itertools::Itertools;
use vortex_dtype::match_each_integer_ptype;
use vortex_error::VortexResult;

use crate::array::BoolArray;
use crate::patches::Patches;
use crate::variants::PrimitiveArrayTrait;
use crate::{ArrayLen, IntoArrayVariant, ToArrayData};

impl BoolArray {
    pub fn patch(self, patches: Patches) -> VortexResult<Self> {
        let length = self.len();
        let indices = patches.indices().clone().into_primitive()?;
        let values = patches.values().clone().into_bool()?;

        let patched_validity =
            self.validity()
                .patch(self.len(), &indices.to_array(), values.validity())?;

        let (mut own_values, bit_offset) = self.into_boolean_builder();
        match_each_integer_ptype!(indices.ptype(), |$I| {
            for (idx, value) in indices
                .into_maybe_null_slice::<$I>()
                .into_iter()
                .zip_eq(values.boolean_buffer().iter())
            {
                own_values.set_bit(idx as usize + bit_offset, value);
            }
        });

        Self::try_new(
            own_values.finish().slice(bit_offset, length),
            patched_validity,
        )
    }
}

#[cfg(test)]
mod tests {
    use arrow_buffer::BooleanBuffer;

    use crate::array::BoolArray;
    use crate::compute::slice;
    use crate::IntoArrayVariant;

    #[test]
    fn patch_sliced_bools() {
        let arr = BoolArray::from(BooleanBuffer::new_set(12));
        let sliced = slice(arr, 4, 12).unwrap();
        let (values, offset) = sliced.into_bool().unwrap().into_boolean_builder();
        assert_eq!(offset, 4);
        assert_eq!(values.as_slice(), &[255, 15]);
    }

    #[test]
    fn patch_sliced_bools_offset() {
        let arr = BoolArray::from(BooleanBuffer::new_set(15));
        let sliced = slice(arr, 4, 15).unwrap();
        let (values, offset) = sliced.into_bool().unwrap().into_boolean_builder();
        assert_eq!(offset, 4);
        assert_eq!(values.as_slice(), &[255, 127]);
    }

    #[test]
    fn patch_sliced_bools_even() {
        let arr = BoolArray::from(BooleanBuffer::new_set(31));
        let sliced = slice(arr, 8, 24).unwrap();
        let (values, offset) = sliced.into_bool().unwrap().into_boolean_builder();
        assert_eq!(offset, 0);
        assert_eq!(values.as_slice(), &[255, 255]);
    }
}