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

vortex_array/arrays/primitive/
patch.rs
