vortex_array/arrays/chunked/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::Array;
7use crate::arrays::ChunkedArray;
8use crate::arrays::ChunkedVTable;
9use crate::scalar::Scalar;
10use crate::vtable::OperationsVTable;
11
12impl OperationsVTable<ChunkedVTable> for ChunkedVTable {
13    fn scalar_at(array: &ChunkedArray, index: usize) -> VortexResult<Scalar> {
14        let (chunk_index, chunk_offset) = array.find_chunk_idx(index)?;
15        array.chunk(chunk_index).scalar_at(chunk_offset)
16    }
17}
18
19#[cfg(test)]
20mod tests {
21    use std::ops::Range;
22
23    use rstest::rstest;
24    use vortex_buffer::Buffer;
25    use vortex_buffer::buffer;
26    use vortex_dtype::DType;
27    use vortex_dtype::Nullability;
28    use vortex_dtype::PType;
29
30    use crate::IntoArray;
31    use crate::arrays::ChunkedArray;
32    use crate::arrays::PrimitiveArray;
33    use crate::assert_arrays_eq;
34
35    fn chunked_array() -> ChunkedArray {
36        ChunkedArray::try_new(
37            vec![
38                buffer![1u64, 2, 3].into_array(),
39                buffer![4u64, 5, 6].into_array(),
40                buffer![7u64, 8, 9].into_array(),
41            ],
42            DType::Primitive(PType::U64, Nullability::NonNullable),
43        )
44        .unwrap()
45    }
46
47    #[rstest]
48    #[case::middle(2..5, &[3u64, 4, 5])]
49    #[case::begin(1..3, &[2u64, 3])]
50    #[case::aligned(3..6, &[4u64, 5, 6])]
51    #[case::many_aligned(0..6, &[1u64, 2, 3, 4, 5, 6])]
52    #[case::end(7..8, &[8u64])]
53    #[case::exactly_end(6..9, &[7u64, 8, 9])]
54    fn slice(#[case] range: Range<usize>, #[case] expected: &[u64]) {
55        assert_arrays_eq!(
56            chunked_array().slice(range).unwrap(),
57            PrimitiveArray::from_iter(expected.iter().copied())
58        );
59    }
60
61    #[test]
62    fn slice_empty() {
63        let chunked = ChunkedArray::try_new(vec![], PType::U32.into()).unwrap();
64        let sliced = chunked.slice(0..0).unwrap();
65
66        assert!(sliced.is_empty());
67    }
68
69    #[test]
70    fn scalar_at_empty_children_both_sides() {
71        let array = ChunkedArray::try_new(
72            vec![
73                Buffer::<u64>::empty().into_array(),
74                Buffer::<u64>::empty().into_array(),
75                buffer![1u64, 2].into_array(),
76                Buffer::<u64>::empty().into_array(),
77                Buffer::<u64>::empty().into_array(),
78            ],
79            DType::Primitive(PType::U64, Nullability::NonNullable),
80        )
81        .unwrap();
82        assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2]));
83    }
84
85    #[test]
86    fn scalar_at_empty_children_trailing() {
87        let array = ChunkedArray::try_new(
88            vec![
89                buffer![1u64, 2].into_array(),
90                Buffer::<u64>::empty().into_array(),
91                Buffer::<u64>::empty().into_array(),
92                buffer![3u64, 4].into_array(),
93            ],
94            DType::Primitive(PType::U64, Nullability::NonNullable),
95        )
96        .unwrap();
97        assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]));
98    }
99
100    #[test]
101    fn scalar_at_empty_children_leading() {
102        let array = ChunkedArray::try_new(
103            vec![
104                Buffer::<u64>::empty().into_array(),
105                Buffer::<u64>::empty().into_array(),
106                buffer![1u64, 2].into_array(),
107                buffer![3u64, 4].into_array(),
108            ],
109            DType::Primitive(PType::U64, Nullability::NonNullable),
110        )
111        .unwrap();
112        assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]));
113    }
114}
vortex_array/arrays/chunked/vtable/operations.rs

vortex_array/arrays/chunked/vtable/
operations.rs
