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    use crate::scalar::Scalar;
59
60    fn chunked_array() -> ChunkedArray {
61        ChunkedArray::try_new(
62            vec![
63                buffer![1u64, 2, 3].into_array(),
64                buffer![4u64, 5, 6].into_array(),
65                buffer![7u64, 8, 9].into_array(),
66            ],
67            DType::Primitive(PType::U64, Nullability::NonNullable),
68        )
69        .unwrap()
70    }
71
72    #[rstest]
73    #[case::middle(2..5, &[3u64, 4, 5])]
74    #[case::begin(1..3, &[2u64, 3])]
75    #[case::aligned(3..6, &[4u64, 5, 6])]
76    #[case::many_aligned(0..6, &[1u64, 2, 3, 4, 5, 6])]
77    #[case::end(7..8, &[8u64])]
78    #[case::exactly_end(6..9, &[7u64, 8, 9])]
79    fn slice(#[case] range: Range<usize>, #[case] expected: &[u64]) {
80        let mut ctx = array_session().create_execution_ctx();
81        assert_arrays_eq!(
82            chunked_array().slice(range).unwrap(),
83            PrimitiveArray::from_iter(expected.iter().copied()),
84            &mut ctx
85        );
86    }
87
88    #[test]
89    fn slice_empty() {
90        let chunked = ChunkedArray::try_new(vec![], PType::U32.into()).unwrap();
91        let sliced = chunked.slice(0..0).unwrap();
92
93        assert_arrays_eq!(
94            sliced,
95            PrimitiveArray::empty::<u32>(Nullability::NonNullable),
96            &mut array_session().create_execution_ctx()
97        );
98    }
99
100    #[test]
101    fn scalar_at_empty_children_both_sides() {
102        let mut ctx = array_session().create_execution_ctx();
103        let array = ChunkedArray::try_new(
104            vec![
105                Buffer::<u64>::empty().into_array(),
106                Buffer::<u64>::empty().into_array(),
107                buffer![1u64, 2].into_array(),
108                Buffer::<u64>::empty().into_array(),
109                Buffer::<u64>::empty().into_array(),
110            ],
111            DType::Primitive(PType::U64, Nullability::NonNullable),
112        )
113        .unwrap();
114        for (index, value) in [1u64, 2].into_iter().enumerate() {
115            assert_eq!(
116                array.execute_scalar(index, &mut ctx).unwrap(),
117                Scalar::from(value)
118            );
119        }
120        assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2]), &mut ctx);
121    }
122
123    #[test]
124    fn scalar_at_empty_children_trailing() {
125        let mut ctx = array_session().create_execution_ctx();
126        let array = ChunkedArray::try_new(
127            vec![
128                buffer![1u64, 2].into_array(),
129                Buffer::<u64>::empty().into_array(),
130                Buffer::<u64>::empty().into_array(),
131                buffer![3u64, 4].into_array(),
132            ],
133            DType::Primitive(PType::U64, Nullability::NonNullable),
134        )
135        .unwrap();
136        for (index, value) in [1u64, 2, 3, 4].into_iter().enumerate() {
137            assert_eq!(
138                array.execute_scalar(index, &mut ctx).unwrap(),
139                Scalar::from(value)
140            );
141        }
142        assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]), &mut ctx);
143    }
144
145    #[test]
146    fn scalar_at_empty_children_leading() {
147        let mut ctx = array_session().create_execution_ctx();
148        let array = ChunkedArray::try_new(
149            vec![
150                Buffer::<u64>::empty().into_array(),
151                Buffer::<u64>::empty().into_array(),
152                buffer![1u64, 2].into_array(),
153                buffer![3u64, 4].into_array(),
154            ],
155            DType::Primitive(PType::U64, Nullability::NonNullable),
156        )
157        .unwrap();
158        for (index, value) in [1u64, 2, 3, 4].into_iter().enumerate() {
159            assert_eq!(
160                array.execute_scalar(index, &mut ctx).unwrap(),
161                Scalar::from(value)
162            );
163        }
164        assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]), &mut ctx);
165    }
166}
vortex_array/arrays/chunked/vtable/operations.rs

vortex_array/arrays/chunked/vtable/
operations.rs
