vortex_array/arrays/primitive/
patch.rs

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

vortex_array/arrays/primitive/
patch.rs
