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

vortex_array/arrays/chunked/vtable/
operations.rs
