vortex_array/arrays/chunked/compute/
mod.rs1use 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 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}