vortex_array/arrays/primitive/array/
patch.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use vortex_error::VortexResult;
7
8use crate::ExecutionCtx;
9use crate::arrays::PrimitiveArray;
10use crate::dtype::IntegerPType;
11use crate::dtype::NativePType;
12use crate::dtype::UnsignedPType;
13use crate::match_each_integer_ptype;
14use crate::match_each_native_ptype;
15use crate::patches::PATCH_CHUNK_SIZE;
16use crate::patches::Patches;
17use crate::validity::Validity;
18use crate::vtable::ValidityHelper;
19
20impl PrimitiveArray {
21    pub fn patch(self, patches: &Patches, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
22        let patch_indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
23        let patch_values = patches.values().clone().execute::<PrimitiveArray>(ctx)?;
24
25        let patched_validity = self.validity().clone().patch(
26            self.len(),
27            patches.offset(),
28            &patch_indices.to_array(),
29            patch_values.validity(),
30            ctx,
31        )?;
32        Ok(match_each_integer_ptype!(patch_indices.ptype(), |I| {
33            match_each_native_ptype!(self.ptype(), |T| {
34                self.patch_typed::<T, I>(
35                    patch_indices,
36                    patches.offset(),
37                    patch_values,
38                    patched_validity,
39                )
40            })
41        }))
42    }
43
44    fn patch_typed<T, I>(
45        self,
46        patch_indices: PrimitiveArray,
47        patch_indices_offset: usize,
48        patch_values: PrimitiveArray,
49        patched_validity: Validity,
50    ) -> Self
51    where
52        T: NativePType,
53        I: IntegerPType,
54    {
55        let mut own_values = self.into_buffer_mut::<T>();
56
57        let patch_indices = patch_indices.as_slice::<I>();
58        let patch_values = patch_values.as_slice::<T>();
59        for (idx, value) in itertools::zip_eq(patch_indices, patch_values) {
60            own_values[idx.as_() - patch_indices_offset] = *value;
61        }
62        Self::new(own_values, patched_validity)
63    }
64}
65
66/// Computes the index range for a chunk, accounting for slice offset.
67///
68/// # Arguments
69///
70/// * `chunk_idx` - Index of the chunk
71/// * `offset` - Offset from slice
72/// * `array_len` - Length of the sliced array
73#[inline]
74pub fn chunk_range(chunk_idx: usize, offset: usize, array_len: usize) -> Range<usize> {
75    let offset_in_chunk = offset % PATCH_CHUNK_SIZE;
76    let local_start = (chunk_idx * PATCH_CHUNK_SIZE).saturating_sub(offset_in_chunk);
77    let local_end = ((chunk_idx + 1) * PATCH_CHUNK_SIZE)
78        .saturating_sub(offset_in_chunk)
79        .min(array_len);
80    local_start..local_end
81}
82
83/// Patches a chunk of decoded values.
84///
85/// # Arguments
86///
87/// * `decoded_values` - Mutable slice of decoded values to be patched
88/// * `patches_indices` - Indices indicating which positions to patch
89/// * `patches_values` - Values to apply at the patched indices
90/// * `patches_offset` - Absolute position where the slice starts
91/// * `chunk_offsets_slice` - Slice containing offsets for each chunk
92/// * `chunk_idx` - Index of the chunk to patch
93/// * `offset_within_chunk` - Number of patches to skip at the start of the first chunk
94#[inline]
95pub fn patch_chunk<T, I, C>(
96    decoded_values: &mut [T],
97    patches_indices: &[I],
98    patches_values: &[T],
99    patches_offset: usize,
100    chunk_offsets_slice: &[C],
101    chunk_idx: usize,
102    offset_within_chunk: usize,
103) where
104    T: NativePType,
105    I: UnsignedPType,
106    C: UnsignedPType,
107{
108    // Compute base_offset from the first chunk offset.
109    let base_offset: usize = chunk_offsets_slice[0].as_();
110
111    // Use the same logic as patches slice implementation for calculating patch ranges.
112    let patches_start_idx =
113        (chunk_offsets_slice[chunk_idx].as_() - base_offset).saturating_sub(offset_within_chunk);
114    let patches_end_idx = if chunk_idx + 1 < chunk_offsets_slice.len() {
115        chunk_offsets_slice[chunk_idx + 1].as_() - base_offset - offset_within_chunk
116    } else {
117        patches_indices.len()
118    };
119
120    let chunk_start = chunk_range(chunk_idx, patches_offset, /* ignore */ usize::MAX).start;
121
122    for patches_idx in patches_start_idx..patches_end_idx {
123        let chunk_relative_index =
124            (patches_indices[patches_idx].as_() - patches_offset) - chunk_start;
125        decoded_values[chunk_relative_index] = patches_values[patches_idx];
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use vortex_buffer::buffer;
132
133    use super::*;
134    use crate::ToCanonical;
135    use crate::assert_arrays_eq;
136    use crate::validity::Validity;
137
138    #[test]
139    fn patch_sliced() {
140        let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid);
141        let sliced = input.slice(2..8).unwrap();
142        assert_arrays_eq!(
143            sliced.to_primitive(),
144            PrimitiveArray::new(buffer![2u32; 6], Validity::AllValid)
145        );
146    }
147}
vortex_array/arrays/primitive/array/patch.rs

vortex_array/arrays/primitive/array/
patch.rs
