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::ExecutionCtx;
7use crate::array::ArrayView;
8use crate::array::OperationsVTable;
9use crate::arrays::Chunked;
10use crate::arrays::chunked::ChunkedArrayExt;
11use crate::scalar::Scalar;
12
13impl OperationsVTable<Chunked> for Chunked {
14    fn scalar_at(
15        array: ArrayView<'_, Chunked>,
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).execute_scalar(chunk_offset, ctx)
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::VortexSessionExecute;
34    use crate::array_session;
35    use crate::arrays::ChunkedArray;
36    use crate::arrays::PrimitiveArray;
37    use crate::assert_arrays_eq;
38    use crate::dtype::DType;
39    use crate::dtype::Nullability;
40    use crate::dtype::PType;
41
42    fn chunked_array() -> ChunkedArray {
43        ChunkedArray::try_new(
44            vec![
45                buffer![1u64, 2, 3].into_array(),
46                buffer![4u64, 5, 6].into_array(),
47                buffer![7u64, 8, 9].into_array(),
48            ],
49            DType::Primitive(PType::U64, Nullability::NonNullable),
50        )
51        .unwrap()
52    }
53
54    #[rstest]
55    #[case::middle(2..5, &[3u64, 4, 5])]
56    #[case::begin(1..3, &[2u64, 3])]
57    #[case::aligned(3..6, &[4u64, 5, 6])]
58    #[case::many_aligned(0..6, &[1u64, 2, 3, 4, 5, 6])]
59    #[case::end(7..8, &[8u64])]
60    #[case::exactly_end(6..9, &[7u64, 8, 9])]
61    fn slice(#[case] range: Range<usize>, #[case] expected: &[u64]) {
62        let mut ctx = array_session().create_execution_ctx();
63        assert_arrays_eq!(
64            chunked_array().slice(range).unwrap(),
65            PrimitiveArray::from_iter(expected.iter().copied()),
66            &mut ctx
67        );
68    }
69
70    #[test]
71    fn slice_empty() {
72        let chunked = ChunkedArray::try_new(vec![], PType::U32.into()).unwrap();
73        let sliced = chunked.slice(0..0).unwrap();
74
75        assert!(sliced.is_empty());
76    }
77
78    #[test]
79    fn scalar_at_empty_children_both_sides() {
80        let mut ctx = array_session().create_execution_ctx();
81        let array = ChunkedArray::try_new(
82            vec![
83                Buffer::<u64>::empty().into_array(),
84                Buffer::<u64>::empty().into_array(),
85                buffer![1u64, 2].into_array(),
86                Buffer::<u64>::empty().into_array(),
87                Buffer::<u64>::empty().into_array(),
88            ],
89            DType::Primitive(PType::U64, Nullability::NonNullable),
90        )
91        .unwrap();
92        assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2]), &mut ctx);
93    }
94
95    #[test]
96    fn scalar_at_empty_children_trailing() {
97        let mut ctx = array_session().create_execution_ctx();
98        let array = ChunkedArray::try_new(
99            vec![
100                buffer![1u64, 2].into_array(),
101                Buffer::<u64>::empty().into_array(),
102                Buffer::<u64>::empty().into_array(),
103                buffer![3u64, 4].into_array(),
104            ],
105            DType::Primitive(PType::U64, Nullability::NonNullable),
106        )
107        .unwrap();
108        assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]), &mut ctx);
109    }
110
111    #[test]
112    fn scalar_at_empty_children_leading() {
113        let mut ctx = array_session().create_execution_ctx();
114        let array = ChunkedArray::try_new(
115            vec![
116                Buffer::<u64>::empty().into_array(),
117                Buffer::<u64>::empty().into_array(),
118                buffer![1u64, 2].into_array(),
119                buffer![3u64, 4].into_array(),
120            ],
121            DType::Primitive(PType::U64, Nullability::NonNullable),
122        )
123        .unwrap();
124        assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]), &mut ctx);
125    }
126}
vortex_array/arrays/chunked/vtable/operations.rs

vortex_array/arrays/chunked/vtable/
operations.rs
