vortex_array/arrays/chunked/compute/
scalar_at.rs

1use vortex_error::VortexResult;
2use vortex_scalar::Scalar;
3
4use crate::arrays::{ChunkedArray, ChunkedEncoding};
5use crate::compute::{ScalarAtFn, scalar_at};
6
7impl ScalarAtFn<&ChunkedArray> for ChunkedEncoding {
8    fn scalar_at(&self, array: &ChunkedArray, index: usize) -> VortexResult<Scalar> {
9        let (chunk_index, chunk_offset) = array.find_chunk_idx(index);
10        scalar_at(array.chunk(chunk_index)?, chunk_offset)
11    }
12}
13
14#[cfg(test)]
15mod tests {
16    use vortex_buffer::Buffer;
17    use vortex_dtype::{DType, Nullability, PType};
18
19    use crate::IntoArray;
20    use crate::array::Array;
21    use crate::arrays::{ChunkedArray, PrimitiveArray};
22    use crate::compute::scalar_at;
23
24    #[test]
25    fn empty_children_both_sides() {
26        let array = ChunkedArray::try_new(
27            vec![
28                Buffer::<u64>::empty().into_array(),
29                Buffer::<u64>::empty().into_array(),
30                PrimitiveArray::from_iter([1u64, 2]).into_array(),
31                Buffer::<u64>::empty().into_array(),
32                Buffer::<u64>::empty().into_array(),
33            ],
34            DType::Primitive(PType::U64, Nullability::NonNullable),
35        )
36        .unwrap();
37        assert_eq!(scalar_at(&array, 0).unwrap(), 1u64.into());
38        assert_eq!(scalar_at(&array, 1).unwrap(), 2u64.into());
39    }
40
41    #[test]
42    fn empty_children_trailing() {
43        let array = ChunkedArray::try_new(
44            vec![
45                PrimitiveArray::from_iter([1u64, 2]).into_array(),
46                Buffer::<u64>::empty().into_array(),
47                Buffer::<u64>::empty().into_array(),
48                PrimitiveArray::from_iter([3u64, 4]).into_array(),
49            ],
50            DType::Primitive(PType::U64, Nullability::NonNullable),
51        )
52        .unwrap();
53        assert_eq!(scalar_at(&array, 0).unwrap(), 1u64.into());
54        assert_eq!(scalar_at(&array, 1).unwrap(), 2u64.into());
55        assert_eq!(scalar_at(&array, 2).unwrap(), 3u64.into());
56        assert_eq!(scalar_at(&array, 3).unwrap(), 4u64.into());
57    }
58
59    #[test]
60    fn empty_children_leading() {
61        let array = ChunkedArray::try_new(
62            vec![
63                Buffer::<u64>::empty().into_array(),
64                Buffer::<u64>::empty().into_array(),
65                PrimitiveArray::from_iter([1u64, 2]).into_array(),
66                PrimitiveArray::from_iter([3u64, 4]).into_array(),
67            ],
68            DType::Primitive(PType::U64, Nullability::NonNullable),
69        )
70        .unwrap();
71        assert_eq!(scalar_at(&array, 0).unwrap(), 1u64.into());
72        assert_eq!(scalar_at(&array, 1).unwrap(), 2u64.into());
73        assert_eq!(scalar_at(&array, 2).unwrap(), 3u64.into());
74        assert_eq!(scalar_at(&array, 3).unwrap(), 4u64.into());
75    }
76}
vortex_array/arrays/chunked/compute/scalar_at.rs

vortex_array/arrays/chunked/compute/
scalar_at.rs
