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 (_, offset, patch_indices, patch_values) = patches.clone().into_parts();
15        let patch_indices = patch_indices.to_primitive()?;
16        let patch_values = patch_values.to_primitive()?;
17
18        let patched_validity = self.validity().clone().patch(
19            self.len(),
20            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>(patch_indices, offset, patch_values, patched_validity)
27            })
28        })
29    }
30
31    fn patch_typed<T, I>(
32        self,
33        patch_indices: PrimitiveArray,
34        patch_indices_offset: usize,
35        patch_values: PrimitiveArray,
36        patched_validity: Validity,
37    ) -> VortexResult<Self>
38    where
39        T: NativePType + ArrowNativeType,
40        I: NativePType + ArrowNativeType,
41    {
42        let mut own_values = self.into_buffer_mut::<T>();
43
44        let patch_indices = patch_indices.as_slice::<I>();
45        let patch_values = patch_values.as_slice::<T>();
46        for (idx, value) in itertools::zip_eq(patch_indices, patch_values) {
47            own_values[idx.as_usize() - patch_indices_offset] = *value;
48        }
49        Ok(Self::new(own_values, patched_validity))
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use vortex_buffer::buffer;
56
57    use super::*;
58    use crate::ToCanonical;
59    use crate::validity::Validity;
60
61    #[test]
62    fn patch_sliced() {
63        let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid);
64        let sliced = input.slice(2, 8).unwrap();
65        assert_eq!(sliced.to_primitive().unwrap().as_slice::<u32>(), &[2u32; 6]);
66    }
67}
vortex_array/arrays/primitive/patch.rs

vortex_array/arrays/primitive/
patch.rs
