Skip to main content

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_buffer::BitBufferMut;
6use vortex_error::VortexResult;
7
8use crate::ExecutionCtx;
9use crate::arrays::BoolArray;
10use crate::arrays::PrimitiveArray;
11use crate::arrays::bool::BoolArrayExt;
12use crate::match_each_unsigned_integer_ptype;
13use crate::patches::Patches;
14
15impl BoolArray {
16    pub fn patch(self, patches: &Patches, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
17        let len = self.len();
18        let offset = patches.offset();
19        let indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
20        let values = patches.values().clone().execute::<BoolArray>(ctx)?;
21
22        let patched_validity =
23            self.validity()?
24                .patch(len, offset, patches.indices(), &values.validity()?, ctx)?;
25
26        let bit_buffer = self.into_bit_buffer();
27        let mut own_values = bit_buffer
28            .try_into_mut()
29            .unwrap_or_else(|bb| BitBufferMut::copy_from(&bb));
30        match_each_unsigned_integer_ptype!(indices.ptype(), |I| {
31            for (idx, value) in indices
32                .as_slice::<I>()
33                .iter()
34                .zip_eq(values.to_bit_buffer().iter())
35            {
36                #[allow(clippy::cast_possible_truncation)]
37                own_values.set_to(*idx as usize - offset, value);
38            }
39        });
40
41        Ok(Self::new(own_values.freeze(), patched_validity))
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use vortex_buffer::BitBuffer;
48
49    use crate::IntoArray;
50    use crate::arrays::BoolArray;
51    use crate::assert_arrays_eq;
52
53    #[test]
54    fn patch_sliced_bools() {
55        let arr = BoolArray::from(BitBuffer::new_set(12));
56        let sliced = arr.into_array().slice(4..12).unwrap();
57        let expected = BoolArray::from_iter([true; 8]);
58        assert_arrays_eq!(sliced, expected);
59    }
60
61    #[test]
62    fn patch_sliced_bools_offset() {
63        let arr = BoolArray::from(BitBuffer::new_set(15));
64        let sliced = arr.into_array().slice(4..15).unwrap();
65        let expected = BoolArray::from_iter([true; 11]);
66        assert_arrays_eq!(sliced, expected);
67    }
68}