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

vortex_array/arrays/chunked/vtable/
operations.rs
