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::ExecutionCtx;
69    use crate::IntoArray;
70    use crate::LEGACY_SESSION;
71    use crate::arrays::Patched;
72    use crate::arrays::PrimitiveArray;
73    use crate::assert_arrays_eq;
74    use crate::dtype::NativePType;
75    use crate::patches::Patches;
76
77    #[test]
78    fn test_reduce() -> VortexResult<()> {
79        let values = buffer![0u16; 512].into_array();
80        let patch_indices = buffer![1u32, 8, 30].into_array();
81        let patch_values = buffer![u16::MAX; 3].into_array();
82        let patches = Patches::new(512, 0, patch_indices, patch_values, None)?;
83
84        let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone());
85
86        let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx)?;
87
88        let sliced = patched_array.into_array().slice(1..10)?;
89
90        insta::assert_snapshot!(
91            sliced.display_tree_encodings_only(),
92            @r#"
93            root: vortex.patched(u16, len=9)
94              inner: vortex.primitive(u16, len=9)
95              lane_offsets: vortex.primitive(u32, len=33)
96              patch_indices: vortex.primitive(u16, len=3)
97              patch_values: vortex.primitive(u16, len=3)
98            "#);
99
100        let executed = sliced.execute::<Canonical>(&mut ctx)?.into_primitive();
101
102        assert_eq!(
103            &[u16::MAX, 0, 0, 0, 0, 0, 0, u16::MAX, 0],
104            executed.as_slice::<u16>()
105        );
106
107        Ok(())
108    }
109
110    #[rstest]
111    #[case::trivial(buffer![1u64; 2], buffer![1u32], buffer![u64::MAX], 1..2)]
112    #[case::one_chunk(buffer![0u64; 1024], buffer![1u32, 8, 30], buffer![u64::MAX; 3], 1..10)]
113    #[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)]
114    fn test_cases<T: NativePType>(
115        #[case] inner: Buffer<T>,
116        #[case] patch_indices: Buffer<u32>,
117        #[case] patch_values: Buffer<T>,
118        #[case] range: Range<usize>,
119    ) {
120        // Create patched array.
121        let patches = Patches::new(
122            inner.len(),
123            0,
124            patch_indices.into_array(),
125            patch_values.into_array(),
126            None,
127        )
128        .unwrap();
129
130        let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone());
131
132        let patched_array = Patched::from_array_and_patches(inner.into_array(), &patches, &mut ctx)
133            .unwrap()
134            .into_array();
135
136        // Verify that applying slice first yields same result as applying slice at end.
137        let slice_first = patched_array
138            .slice(range.clone())
139            .unwrap()
140            .execute::<Canonical>(&mut ctx)
141            .unwrap()
142            .into_array();
143
144        let slice_last = patched_array
145            .execute::<Canonical>(&mut ctx)
146            .unwrap()
147            .into_primitive()
148            .slice(range)
149            .unwrap();
150
151        assert_arrays_eq!(slice_first, slice_last);
152    }
153
154    #[test]
155    fn test_stacked_slices() {
156        let values = PrimitiveArray::from_iter(0u64..10_000).into_array();
157
158        let patched_indices = buffer![1u32, 2, 1024, 2048, 3072, 3088].into_array();
159        let patched_values = buffer![0u64, 1, 2, 3, 4, 5].into_array();
160
161        let patches = Patches::new(10_000, 0, patched_indices, patched_values, None).unwrap();
162        let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone());
163
164        let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx)
165            .unwrap()
166            .into_array();
167
168        let sliced = patched_array
169            .slice(1024..5000)
170            .unwrap()
171            .slice(1..2065)
172            .unwrap()
173            .execute::<Canonical>(&mut ctx)
174            .unwrap()
175            .into_array();
176
177        let mut expected = BufferMut::from_iter(1025u64..=3088);
178        expected[1023] = 3;
179        expected[2047] = 4;
180        expected[2063] = 5;
181
182        let expected = expected.into_array();
183
184        assert_arrays_eq!(expected, sliced);
185    }
186}
vortex_array/arrays/patched/vtable/slice.rs

vortex_array/arrays/patched/vtable/
slice.rs
