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