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
95#[inline]
96pub fn patch_chunk<T, I, C>(
97    decoded_values: &mut [T],
98    patches_indices: &[I],
99    patches_values: &[T],
100    patches_offset: usize,
101    chunk_offsets_slice: &[C],
102    chunk_idx: usize,
103    offset_within_chunk: usize,
104) where
105    T: NativePType,
106    I: UnsignedPType,
107    C: UnsignedPType,
108{
109    // Compute base_offset from the first chunk offset.
110    let base_offset: usize = chunk_offsets_slice[0].as_();
111
112    // Use the same logic as patches slice implementation for calculating patch ranges.
113    let patches_start_idx =
114        (chunk_offsets_slice[chunk_idx].as_() - base_offset).saturating_sub(offset_within_chunk);
115    let patches_end_idx = if chunk_idx + 1 < chunk_offsets_slice.len() {
116        chunk_offsets_slice[chunk_idx + 1].as_() - base_offset - offset_within_chunk
117    } else {
118        patches_indices.len()
119    };
120
121    let chunk_start = chunk_range(chunk_idx, patches_offset, /* ignore */ usize::MAX).start;
122
123    for patches_idx in patches_start_idx..patches_end_idx {
124        let chunk_relative_index =
125            (patches_indices[patches_idx].as_() - patches_offset) - chunk_start;
126        decoded_values[chunk_relative_index] = patches_values[patches_idx];
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use vortex_buffer::buffer;
133
134    use super::*;
135    use crate::ToCanonical;
136    use crate::assert_arrays_eq;
137    use crate::validity::Validity;
138
139    #[test]
140    fn patch_sliced() {
141        let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid);
142        let sliced = input.slice(2..8).unwrap();
143        assert_arrays_eq!(
144            sliced.to_primitive(),
145            PrimitiveArray::new(buffer![2u32; 6], Validity::AllValid)
146        );
147    }
148}
vortex_array/arrays/primitive/array/patch.rs

vortex_array/arrays/primitive/array/
patch.rs
