vortex_array/arrays/chunked/compute/
mod.rs

1use vortex_dtype::DType;
2use vortex_error::VortexResult;
3
4use crate::arrays::ChunkedEncoding;
5use crate::arrays::chunked::ChunkedArray;
6use crate::compute::{
7    BinaryBooleanFn, BinaryNumericFn, CastFn, CompareFn, FillNullFn, FilterFn, InvertFn,
8    IsConstantFn, MaskFn, MinMaxFn, ScalarAtFn, SliceFn, TakeFn, UncompressedSizeFn, try_cast,
9};
10use crate::vtable::ComputeVTable;
11use crate::{Array, ArrayRef};
12
13mod binary_numeric;
14mod boolean;
15mod compare;
16mod fill_null;
17mod filter;
18mod invert;
19mod is_constant;
20mod mask;
21mod min_max;
22mod scalar_at;
23mod slice;
24mod sum;
25mod take;
26mod uncompressed_size;
27
28impl ComputeVTable for ChunkedEncoding {
29    fn binary_boolean_fn(&self) -> Option<&dyn BinaryBooleanFn<&dyn Array>> {
30        Some(self)
31    }
32
33    fn binary_numeric_fn(&self) -> Option<&dyn BinaryNumericFn<&dyn Array>> {
34        Some(self)
35    }
36
37    fn cast_fn(&self) -> Option<&dyn CastFn<&dyn Array>> {
38        Some(self)
39    }
40
41    fn compare_fn(&self) -> Option<&dyn CompareFn<&dyn Array>> {
42        Some(self)
43    }
44
45    fn fill_null_fn(&self) -> Option<&dyn FillNullFn<&dyn Array>> {
46        Some(self)
47    }
48
49    fn filter_fn(&self) -> Option<&dyn FilterFn<&dyn Array>> {
50        Some(self)
51    }
52
53    fn invert_fn(&self) -> Option<&dyn InvertFn<&dyn Array>> {
54        Some(self)
55    }
56
57    fn is_constant_fn(&self) -> Option<&dyn IsConstantFn<&dyn Array>> {
58        Some(self)
59    }
60
61    fn mask_fn(&self) -> Option<&dyn MaskFn<&dyn Array>> {
62        Some(self)
63    }
64
65    fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
66        Some(self)
67    }
68
69    fn slice_fn(&self) -> Option<&dyn SliceFn<&dyn Array>> {
70        Some(self)
71    }
72
73    fn take_fn(&self) -> Option<&dyn TakeFn<&dyn Array>> {
74        Some(self)
75    }
76
77    fn min_max_fn(&self) -> Option<&dyn MinMaxFn<&dyn Array>> {
78        Some(self)
79    }
80
81    fn uncompressed_size_fn(&self) -> Option<&dyn UncompressedSizeFn<&dyn Array>> {
82        Some(self)
83    }
84}
85
86impl CastFn<&ChunkedArray> for ChunkedEncoding {
87    fn cast(&self, array: &ChunkedArray, dtype: &DType) -> VortexResult<ArrayRef> {
88        let mut cast_chunks = Vec::new();
89        for chunk in array.chunks() {
90            cast_chunks.push(try_cast(chunk, dtype)?);
91        }
92
93        Ok(ChunkedArray::new_unchecked(cast_chunks, dtype.clone()).into_array())
94    }
95}
96
97#[cfg(test)]
98mod test {
99    use vortex_buffer::buffer;
100    use vortex_dtype::{DType, Nullability, PType};
101
102    use crate::IntoArray;
103    use crate::array::Array;
104    use crate::arrays::chunked::ChunkedArray;
105    use crate::canonical::ToCanonical;
106    use crate::compute::try_cast;
107
108    #[test]
109    fn test_cast_chunked() {
110        let arr0 = buffer![0u32, 1].into_array();
111        let arr1 = buffer![2u32, 3].into_array();
112
113        let chunked = ChunkedArray::try_new(
114            vec![arr0, arr1],
115            DType::Primitive(PType::U32, Nullability::NonNullable),
116        )
117        .unwrap()
118        .into_array();
119
120        // Two levels of chunking, just to be fancy.
121        let root = ChunkedArray::try_new(
122            vec![chunked],
123            DType::Primitive(PType::U32, Nullability::NonNullable),
124        )
125        .unwrap()
126        .into_array();
127
128        assert_eq!(
129            try_cast(
130                &root,
131                &DType::Primitive(PType::U64, Nullability::NonNullable)
132            )
133            .unwrap()
134            .to_primitive()
135            .unwrap()
136            .as_slice::<u64>(),
137            &[0u64, 1, 2, 3],
138        );
139    }
140}