vortex_array/arrays/chunked/
ops.rs

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

vortex_array/arrays/chunked/
ops.rs
