vortex_array/arrays/primitive/
patch.rs

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

vortex_array/arrays/primitive/
patch.rs
