vortex_array/arrays/patched/vtable/
slice.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::ArrayRef;
9use crate::IntoArray;
10use crate::array::ArrayView;
11use crate::arrays::Patched;
12use crate::arrays::patched::PatchedArrayExt;
13use crate::arrays::patched::PatchedArraySlotsExt;
14use crate::arrays::patched::PatchedSlots;
15use crate::arrays::slice::SliceReduce;
16
17impl SliceReduce for Patched {
18    fn slice(array: ArrayView<'_, Self>, range: Range<usize>) -> VortexResult<Option<ArrayRef>> {
19        // We **always** slice the patches at 1024-element chunk boundaries. We keep the offset + len
20        // around so that when we execute we know how much to chop off.
21        let new_offset = (range.start + array.offset()) % 1024;
22        let chunk_start = (range.start + array.offset()) / 1024;
23        let chunk_stop = (range.end + array.offset()).div_ceil(1024);
24        let sliced_lane_offsets = array
25            .lane_offsets()
26            .slice((chunk_start * array.n_lanes())..(chunk_stop * array.n_lanes()) + 1)?;
27
28        // Unlike the patches, we slice the inner to the exact range. This is handled
29        // at execution time by making sure to skip patch indices that are < offset
30        // or >= len.
31        let inner = array.inner().slice(range.start..range.end)?;
32        let len = inner.len();
33
34        let slots = PatchedSlots {
35            inner,
36            lane_offsets: sliced_lane_offsets,
37            patch_indices: array.patch_indices().clone(),
38            patch_values: array.patch_values().clone(),
39        }
40        .into_slots();
41
42        Ok(Some(
43            unsafe {
44                Patched::new_unchecked(
45                    array.dtype().clone(),
46                    len,
47                    slots,
48                    array.n_lanes(),
49                    new_offset,
50                )
51            }
52            .into_array(),
53        ))
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use std::ops::Range;
60
61    use rstest::rstest;
62    use vortex_buffer::Buffer;
63    use vortex_buffer::BufferMut;
64    use vortex_buffer::buffer;
65    use vortex_error::VortexResult;
66
67    use crate::Canonical;
68    use crate::IntoArray;
69    use crate::VortexSessionExecute;
70    use crate::arrays::Patched;
71    use crate::arrays::PrimitiveArray;
72    use crate::assert_arrays_eq;
73    use crate::dtype::NativePType;
74    use crate::patches::Patches;
75
76    #[test]
77    fn test_reduce() -> VortexResult<()> {
78        let values = buffer![0u16; 512].into_array();
79        let patch_indices = buffer![1u32, 8, 30].into_array();
80        let patch_values = buffer![u16::MAX; 3].into_array();
81        let patches = Patches::new(512, 0, patch_indices, patch_values, None)?;
82
83        let mut ctx = crate::array_session().create_execution_ctx();
84
85        let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx)?;
86
87        let sliced = patched_array.into_array().slice(1..10)?;
88
89        insta::assert_snapshot!(
90            sliced.display_tree_encodings_only(),
91            @r#"
92            root: vortex.patched(u16, len=9)
93              inner: vortex.primitive(u16, len=9)
94              lane_offsets: vortex.primitive(u32, len=33)
95              patch_indices: vortex.primitive(u16, len=3)
96              patch_values: vortex.primitive(u16, len=3)
97            "#);
98
99        let executed = sliced.execute::<Canonical>(&mut ctx)?.into_primitive();
100
101        assert_eq!(
102            &[u16::MAX, 0, 0, 0, 0, 0, 0, u16::MAX, 0],
103            executed.as_slice::<u16>()
104        );
105
106        Ok(())
107    }
108
109    #[rstest]
110    #[case::trivial(buffer![1u64; 2], buffer![1u32], buffer![u64::MAX], 1..2)]
111    #[case::one_chunk(buffer![0u64; 1024], buffer![1u32, 8, 30], buffer![u64::MAX; 3], 1..10)]
112    #[case::multichunk(buffer![1u64; 10_000], buffer![0u32, 1, 2, 3, 4, 16, 17, 18, 19, 1024, 2048, 2049], buffer![u64::MAX; 12], 1024..5000)]
113    fn test_cases<T: NativePType>(
114        #[case] inner: Buffer<T>,
115        #[case] patch_indices: Buffer<u32>,
116        #[case] patch_values: Buffer<T>,
117        #[case] range: Range<usize>,
118    ) {
119        // Create patched array.
120        let patches = Patches::new(
121            inner.len(),
122            0,
123            patch_indices.into_array(),
124            patch_values.into_array(),
125            None,
126        )
127        .unwrap();
128
129        let mut ctx = crate::array_session().create_execution_ctx();
130
131        let patched_array = Patched::from_array_and_patches(inner.into_array(), &patches, &mut ctx)
132            .unwrap()
133            .into_array();
134
135        // Verify that applying slice first yields same result as applying slice at end.
136        let slice_first = patched_array
137            .slice(range.clone())
138            .unwrap()
139            .execute::<Canonical>(&mut ctx)
140            .unwrap()
141            .into_array();
142
143        let slice_last = patched_array
144            .execute::<Canonical>(&mut ctx)
145            .unwrap()
146            .into_primitive()
147            .slice(range)
148            .unwrap();
149
150        assert_arrays_eq!(slice_first, slice_last, &mut ctx);
151    }
152
153    #[test]
154    fn test_stacked_slices() {
155        let values = PrimitiveArray::from_iter(0u64..10_000).into_array();
156
157        let patched_indices = buffer![1u32, 2, 1024, 2048, 3072, 3088].into_array();
158        let patched_values = buffer![0u64, 1, 2, 3, 4, 5].into_array();
159
160        let patches = Patches::new(10_000, 0, patched_indices, patched_values, None).unwrap();
161        let mut ctx = crate::array_session().create_execution_ctx();
162
163        let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx)
164            .unwrap()
165            .into_array();
166
167        let sliced = patched_array
168            .slice(1024..5000)
169            .unwrap()
170            .slice(1..2065)
171            .unwrap()
172            .execute::<Canonical>(&mut ctx)
173            .unwrap()
174            .into_array();
175
176        let mut expected = BufferMut::from_iter(1025u64..=3088);
177        expected[1023] = 3;
178        expected[2047] = 4;
179        expected[2063] = 5;
180
181        let expected = expected.into_array();
182
183        assert_arrays_eq!(expected, sliced, &mut ctx);
184    }
185}
vortex_array/arrays/patched/vtable/slice.rs

vortex_array/arrays/patched/vtable/
slice.rs
