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::IntegerPType;
5use vortex_dtype::NativePType;
6use vortex_dtype::UnsignedPType;
7use vortex_dtype::match_each_integer_ptype;
8use vortex_dtype::match_each_native_ptype;
9
10use crate::ToCanonical;
11use crate::arrays::PrimitiveArray;
12use crate::patches::Patches;
13use crate::validity::Validity;
14use crate::vtable::ValidityHelper;
15
16impl PrimitiveArray {
17    pub fn patch(self, patches: &Patches) -> Self {
18        let patch_indices = patches.indices().to_primitive();
19        let patch_values = patches.values().to_primitive();
20
21        let patched_validity = self.validity().clone().patch(
22            self.len(),
23            patches.offset(),
24            patch_indices.as_ref(),
25            patch_values.validity(),
26        );
27        match_each_integer_ptype!(patch_indices.ptype(), |I| {
28            match_each_native_ptype!(self.ptype(), |T| {
29                self.patch_typed::<T, I>(
30                    patch_indices,
31                    patches.offset(),
32                    patch_values,
33                    patched_validity,
34                )
35            })
36        })
37    }
38
39    fn patch_typed<T, I>(
40        self,
41        patch_indices: PrimitiveArray,
42        patch_indices_offset: usize,
43        patch_values: PrimitiveArray,
44        patched_validity: Validity,
45    ) -> Self
46    where
47        T: NativePType,
48        I: IntegerPType,
49    {
50        let mut own_values = self.into_buffer_mut::<T>();
51
52        let patch_indices = patch_indices.as_slice::<I>();
53        let patch_values = patch_values.as_slice::<T>();
54        for (idx, value) in itertools::zip_eq(patch_indices, patch_values) {
55            own_values[idx.as_() - patch_indices_offset] = *value;
56        }
57        Self::new(own_values, patched_validity)
58    }
59}
60
61/// Patches a chunk of decoded values.
62///
63/// # Arguments
64///
65/// * `decoded_values` - Mutable slice of decoded values to be patched
66/// * `patches_indices` - Indices indicating which positions to patch
67/// * `patches_values` - Values to apply at the patched indices
68/// * `patches_offset` - Offset to subtract from patch indices
69/// * `chunk_offsets_slice` - Slice containing offsets for each chunk
70/// * `chunk_idx` - Index of the chunk to patch
71/// * `base_offset` - Base offset from the first chunk
72/// * `offset_within_chunk` - Offset within chunk for sliced patches
73#[inline]
74#[expect(
75    clippy::too_many_arguments,
76    reason = "all arguments are needed for the patching operation"
77)]
78pub fn patch_chunk<T, I, C>(
79    decoded_values: &mut [T],
80    patches_indices: &[I],
81    patches_values: &[T],
82    patches_offset: usize,
83    chunk_offsets_slice: &[C],
84    chunk_idx: usize,
85    base_offset: usize,
86    offset_within_chunk: usize,
87) where
88    T: NativePType,
89    I: UnsignedPType,
90    C: UnsignedPType,
91{
92    // Use the same logic as patches slice implementation for calculating patch ranges.
93    let patches_start_idx =
94        (chunk_offsets_slice[chunk_idx].as_() - base_offset).saturating_sub(offset_within_chunk);
95    let patches_end_idx = if chunk_idx + 1 < chunk_offsets_slice.len() {
96        chunk_offsets_slice[chunk_idx + 1].as_() - base_offset - offset_within_chunk
97    } else {
98        patches_indices.len()
99    };
100
101    let chunk_start = chunk_idx * 1024;
102    for patches_idx in patches_start_idx..patches_end_idx {
103        let patched_value = patches_values[patches_idx];
104        let absolute_index: usize = patches_indices[patches_idx].as_() - patches_offset;
105        let chunk_relative_index = absolute_index - chunk_start;
106        decoded_values[chunk_relative_index] = patched_value;
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use vortex_buffer::buffer;
113
114    use super::*;
115    use crate::ToCanonical;
116    use crate::validity::Validity;
117
118    #[test]
119    fn patch_sliced() {
120        let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid);
121        let sliced = input.slice(2..8);
122        assert_eq!(sliced.to_primitive().as_slice::<u32>(), &[2u32; 6]);
123    }
124}
vortex_array/arrays/primitive/array/patch.rs

vortex_array/arrays/primitive/array/
patch.rs
