vortex_array/arrays/chunked/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use itertools::Itertools;
7use vortex_scalar::Scalar;
8
9use crate::Array;
10use crate::ArrayRef;
11use crate::IntoArray;
12use crate::arrays::ChunkedArray;
13use crate::arrays::ChunkedVTable;
14use crate::vtable::OperationsVTable;
15
16impl OperationsVTable<ChunkedVTable> for ChunkedVTable {
17    fn slice(array: &ChunkedArray, range: Range<usize>) -> ArrayRef {
18        assert!(
19            !array.is_empty() || (range.start > 0 && range.end > 0),
20            "Empty chunked array can't be sliced from {} to {}",
21            range.start,
22            range.end
23        );
24
25        if array.is_empty() {
26            // SAFETY: empty chunked array trivially satisfies all validations
27            unsafe {
28                return ChunkedArray::new_unchecked(vec![], array.dtype().clone()).into_array();
29            }
30        }
31
32        let (offset_chunk, offset_in_first_chunk) = array.find_chunk_idx(range.start);
33        let (length_chunk, length_in_last_chunk) = array.find_chunk_idx(range.end);
34
35        if length_chunk == offset_chunk {
36            let chunk = array.chunk(offset_chunk);
37            return chunk.slice(offset_in_first_chunk..length_in_last_chunk);
38        }
39
40        let mut chunks = (offset_chunk..length_chunk + 1)
41            .map(|i| array.chunk(i).clone())
42            .collect_vec();
43        if let Some(c) = chunks.first_mut() {
44            *c = c.slice(offset_in_first_chunk..c.len());
45        }
46
47        if length_in_last_chunk == 0 {
48            chunks.pop();
49        } else if let Some(c) = chunks.last_mut() {
50            *c = c.slice(0..length_in_last_chunk);
51        }
52
53        // SAFETY: chunks are slices of the original valid chunks, preserving their dtype.
54        // All chunks maintain the same dtype as the original array.
55        unsafe { ChunkedArray::new_unchecked(chunks, array.dtype().clone()).into_array() }
56    }
57
58    fn scalar_at(array: &ChunkedArray, index: usize) -> Scalar {
59        let (chunk_index, chunk_offset) = array.find_chunk_idx(index);
60        array.chunk(chunk_index).scalar_at(chunk_offset)
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use vortex_buffer::Buffer;
67    use vortex_buffer::buffer;
68    use vortex_dtype::DType;
69    use vortex_dtype::NativePType;
70    use vortex_dtype::Nullability;
71    use vortex_dtype::PType;
72
73    use crate::IntoArray;
74    use crate::array::Array;
75    use crate::arrays::ChunkedArray;
76    use crate::arrays::ChunkedVTable;
77    use crate::canonical::ToCanonical;
78
79    fn chunked_array() -> ChunkedArray {
80        ChunkedArray::try_new(
81            vec![
82                buffer![1u64, 2, 3].into_array(),
83                buffer![4u64, 5, 6].into_array(),
84                buffer![7u64, 8, 9].into_array(),
85            ],
86            DType::Primitive(PType::U64, Nullability::NonNullable),
87        )
88        .unwrap()
89    }
90
91    fn assert_equal_slices<T: NativePType>(arr: &dyn Array, slice: &[T]) {
92        let mut values = Vec::with_capacity(arr.len());
93        if let Some(arr) = arr.as_opt::<ChunkedVTable>() {
94            arr.chunks()
95                .iter()
96                .map(|a| a.to_primitive())
97                .for_each(|a| values.extend_from_slice(a.as_slice::<T>()));
98        } else {
99            values.extend_from_slice(arr.to_primitive().as_slice::<T>());
100        }
101        assert_eq!(values, slice);
102    }
103
104    #[test]
105    fn slice_middle() {
106        assert_equal_slices(&chunked_array().slice(2..5), &[3u64, 4, 5])
107    }
108
109    #[test]
110    fn slice_begin() {
111        assert_equal_slices(&chunked_array().slice(1..3), &[2u64, 3]);
112    }
113
114    #[test]
115    fn slice_aligned() {
116        assert_equal_slices(&chunked_array().slice(3..6), &[4u64, 5, 6]);
117    }
118
119    #[test]
120    fn slice_many_aligned() {
121        assert_equal_slices(&chunked_array().slice(0..6), &[1u64, 2, 3, 4, 5, 6]);
122    }
123
124    #[test]
125    fn slice_end() {
126        assert_equal_slices(&chunked_array().slice(7..8), &[8u64]);
127    }
128
129    #[test]
130    fn slice_exactly_end() {
131        assert_equal_slices(&chunked_array().slice(6..9), &[7u64, 8, 9]);
132    }
133
134    #[test]
135    fn slice_empty() {
136        let chunked = ChunkedArray::try_new(vec![], PType::U32.into()).unwrap();
137        let sliced = chunked.slice(0..0);
138
139        assert!(sliced.is_empty());
140    }
141
142    #[test]
143    fn scalar_at_empty_children_both_sides() {
144        let array = ChunkedArray::try_new(
145            vec![
146                Buffer::<u64>::empty().into_array(),
147                Buffer::<u64>::empty().into_array(),
148                buffer![1u64, 2].into_array(),
149                Buffer::<u64>::empty().into_array(),
150                Buffer::<u64>::empty().into_array(),
151            ],
152            DType::Primitive(PType::U64, Nullability::NonNullable),
153        )
154        .unwrap();
155        assert_eq!(array.scalar_at(0), 1u64.into());
156        assert_eq!(array.scalar_at(1), 2u64.into());
157    }
158
159    #[test]
160    fn scalar_at_empty_children_trailing() {
161        let array = ChunkedArray::try_new(
162            vec![
163                buffer![1u64, 2].into_array(),
164                Buffer::<u64>::empty().into_array(),
165                Buffer::<u64>::empty().into_array(),
166                buffer![3u64, 4].into_array(),
167            ],
168            DType::Primitive(PType::U64, Nullability::NonNullable),
169        )
170        .unwrap();
171        assert_eq!(array.scalar_at(0), 1u64.into());
172        assert_eq!(array.scalar_at(1), 2u64.into());
173        assert_eq!(array.scalar_at(2), 3u64.into());
174        assert_eq!(array.scalar_at(3), 4u64.into());
175    }
176
177    #[test]
178    fn scalar_at_empty_children_leading() {
179        let array = ChunkedArray::try_new(
180            vec![
181                Buffer::<u64>::empty().into_array(),
182                Buffer::<u64>::empty().into_array(),
183                buffer![1u64, 2].into_array(),
184                buffer![3u64, 4].into_array(),
185            ],
186            DType::Primitive(PType::U64, Nullability::NonNullable),
187        )
188        .unwrap();
189        assert_eq!(array.scalar_at(0), 1u64.into());
190        assert_eq!(array.scalar_at(1), 2u64.into());
191        assert_eq!(array.scalar_at(2), 3u64.into());
192        assert_eq!(array.scalar_at(3), 4u64.into());
193    }
194}
vortex_array/arrays/chunked/vtable/operations.rs

vortex_array/arrays/chunked/vtable/
operations.rs
