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

vortex_array/arrays/chunked/vtable/
operations.rs
