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::{
5    IntegerPType, NativePType, UnsignedPType, match_each_integer_ptype, match_each_native_ptype,
6};
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) -> 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    ) -> Self
45    where
46        T: NativePType,
47        I: IntegerPType,
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        Self::new(own_values, patched_validity)
57    }
58}
59
60/// Patches a chunk of decoded values.
61///
62/// # Arguments
63///
64/// * `decoded_values` - Mutable slice of decoded values to be patched
65/// * `patches_indices` - Indices indicating which positions to patch
66/// * `patches_values` - Values to apply at the patched indices
67/// * `patches_offset` - Offset to subtract from patch indices
68/// * `chunk_offsets_slice` - Slice containing offsets for each chunk
69/// * `chunk_idx` - Index of the chunk to patch
70/// * `base_offset` - Base offset from the first chunk
71/// * `offset_within_chunk` - Offset within chunk for sliced patches
72#[inline]
73#[allow(clippy::too_many_arguments)]
74pub fn patch_chunk<T, I, C>(
75    decoded_values: &mut [T],
76    patches_indices: &[I],
77    patches_values: &[T],
78    patches_offset: usize,
79    chunk_offsets_slice: &[C],
80    chunk_idx: usize,
81    base_offset: usize,
82    offset_within_chunk: usize,
83) where
84    T: NativePType,
85    I: UnsignedPType,
86    C: UnsignedPType,
87{
88    // Use the same logic as patches slice implementation for calculating patch ranges.
89    let patches_start_idx =
90        (chunk_offsets_slice[chunk_idx].as_() - base_offset).saturating_sub(offset_within_chunk);
91    let patches_end_idx = if chunk_idx + 1 < chunk_offsets_slice.len() {
92        chunk_offsets_slice[chunk_idx + 1].as_() - base_offset - offset_within_chunk
93    } else {
94        patches_indices.len()
95    };
96
97    let chunk_start = chunk_idx * 1024;
98    for patches_idx in patches_start_idx..patches_end_idx {
99        let patched_value = patches_values[patches_idx];
100        let absolute_index: usize = patches_indices[patches_idx].as_() - patches_offset;
101        let chunk_relative_index = absolute_index - chunk_start;
102        decoded_values[chunk_relative_index] = patched_value;
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use vortex_buffer::buffer;
109
110    use super::*;
111    use crate::ToCanonical;
112    use crate::validity::Validity;
113
114    #[test]
115    fn patch_sliced() {
116        let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid);
117        let sliced = input.slice(2..8);
118        assert_eq!(sliced.to_primitive().as_slice::<u32>(), &[2u32; 6]);
119    }
120}
vortex_array/arrays/primitive/array/patch.rs

vortex_array/arrays/primitive/array/
patch.rs
