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

vortex_array/arrays/chunked/compute/
mod.rs
