vortex_array/arrays/chunked/
ops.rs

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

vortex_array/arrays/chunked/
ops.rs
