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::IntoArray;
10use crate::arrays::PrimitiveArray;
11use crate::dtype::IntegerPType;
12use crate::dtype::NativePType;
13use crate::dtype::UnsignedPType;
14use crate::match_each_integer_ptype;
15use crate::match_each_native_ptype;
16use crate::patches::PATCH_CHUNK_SIZE;
17use crate::patches::Patches;
18use crate::validity::Validity;
19use crate::vtable::ValidityHelper;
20
21impl PrimitiveArray {
22    pub fn patch(self, patches: &Patches, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
23        let patch_indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
24        let patch_values = patches.values().clone().execute::<PrimitiveArray>(ctx)?;
25
26        let patched_validity = self.validity().clone().patch(
27            self.len(),
28            patches.offset(),
29            &patch_indices.clone().into_array(),
30            patch_values.validity(),
31            ctx,
32        )?;
33        Ok(match_each_integer_ptype!(patch_indices.ptype(), |I| {
34            match_each_native_ptype!(self.ptype(), |T| {
35                self.patch_typed::<T, I>(
36                    patch_indices,
37                    patches.offset(),
38                    patch_values,
39                    patched_validity,
40                )
41            })
42        }))
43    }
44
45    fn patch_typed<T, I>(
46        self,
47        patch_indices: PrimitiveArray,
48        patch_indices_offset: usize,
49        patch_values: PrimitiveArray,
50        patched_validity: Validity,
51    ) -> Self
52    where
53        T: NativePType,
54        I: IntegerPType,
55    {
56        let mut own_values = self.into_buffer_mut::<T>();
57
58        let patch_indices = patch_indices.as_slice::<I>();
59        let patch_values = patch_values.as_slice::<T>();
60        for (idx, value) in itertools::zip_eq(patch_indices, patch_values) {
61            own_values[idx.as_() - patch_indices_offset] = *value;
62        }
63        Self::new(own_values, patched_validity)
64    }
65}
66
67/// Computes the index range for a chunk, accounting for slice offset.
68///
69/// # Arguments
70///
71/// * `chunk_idx` - Index of the chunk
72/// * `offset` - Offset from slice
73/// * `array_len` - Length of the sliced array
74#[inline]
75pub fn chunk_range(chunk_idx: usize, offset: usize, array_len: usize) -> Range<usize> {
76    let offset_in_chunk = offset % PATCH_CHUNK_SIZE;
77    let local_start = (chunk_idx * PATCH_CHUNK_SIZE).saturating_sub(offset_in_chunk);
78    let local_end = ((chunk_idx + 1) * PATCH_CHUNK_SIZE)
79        .saturating_sub(offset_in_chunk)
80        .min(array_len);
81    local_start..local_end
82}
83
84/// Patches a chunk of decoded values.
85///
86/// # Arguments
87///
88/// * `decoded_values` - Mutable slice of decoded values to be patched
89/// * `patches_indices` - Indices indicating which positions to patch
90/// * `patches_values` - Values to apply at the patched indices
91/// * `patches_offset` - Absolute position where the slice starts
92/// * `chunk_offsets_slice` - Slice containing offsets for each chunk
93/// * `chunk_idx` - Index of the chunk to patch
94/// * `offset_within_chunk` - Number of patches to skip at the start of the first chunk
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
