vortex_array/arrays/primitive/array/
patch.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::{IntegerPType, NativePType, match_each_integer_ptype, match_each_native_ptype};
5
6use crate::ToCanonical;
7use crate::arrays::PrimitiveArray;
8use crate::patches::Patches;
9use crate::validity::Validity;
10use crate::vtable::ValidityHelper;
11
12impl PrimitiveArray {
13    #[allow(clippy::cognitive_complexity)]
14    pub fn patch(self, patches: &Patches) -> Self {
15        let patch_indices = patches.indices().to_primitive();
16        let patch_values = patches.values().to_primitive();
17
18        let patched_validity = self.validity().clone().patch(
19            self.len(),
20            patches.offset(),
21            patch_indices.as_ref(),
22            patch_values.validity(),
23        );
24        match_each_integer_ptype!(patch_indices.ptype(), |I| {
25            match_each_native_ptype!(self.ptype(), |T| {
26                self.patch_typed::<T, I>(
27                    patch_indices,
28                    patches.offset(),
29                    patch_values,
30                    patched_validity,
31                )
32            })
33        })
34    }
35
36    fn patch_typed<T, I>(
37        self,
38        patch_indices: PrimitiveArray,
39        patch_indices_offset: usize,
40        patch_values: PrimitiveArray,
41        patched_validity: Validity,
42    ) -> Self
43    where
44        T: NativePType,
45        I: IntegerPType,
46    {
47        let mut own_values = self.into_buffer_mut::<T>();
48
49        let patch_indices = patch_indices.as_slice::<I>();
50        let patch_values = patch_values.as_slice::<T>();
51        for (idx, value) in itertools::zip_eq(patch_indices, patch_values) {
52            own_values[idx.as_() - patch_indices_offset] = *value;
53        }
54        Self::new(own_values, patched_validity)
55    }
56}
57
58#[cfg(test)]
59mod tests {
60    use vortex_buffer::buffer;
61
62    use super::*;
63    use crate::ToCanonical;
64    use crate::validity::Validity;
65
66    #[test]
67    fn patch_sliced() {
68        let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid);
69        let sliced = input.slice(2..8);
70        assert_eq!(sliced.to_primitive().as_slice::<u32>(), &[2u32; 6]);
71    }
72}
vortex_array/arrays/primitive/array/patch.rs

vortex_array/arrays/primitive/array/
patch.rs
