vortex_array/arrays/chunked/
variants.rs

1use std::sync::Arc;
2
3use vortex_dtype::{DType, FieldName};
4use vortex_error::{VortexExpect, VortexResult, vortex_err, vortex_panic};
5
6use crate::arrays::chunked::ChunkedArray;
7use crate::variants::{
8    BinaryArrayTrait, BoolArrayTrait, DecimalArrayTrait, ExtensionArrayTrait, ListArrayTrait,
9    NullArrayTrait, PrimitiveArrayTrait, StructArrayTrait, Utf8ArrayTrait,
10};
11use crate::{Array, ArrayRef, ArrayVariantsImpl};
12
13/// Chunked arrays support all DTypes
14impl ArrayVariantsImpl for ChunkedArray {
15    fn _as_null_typed(&self) -> Option<&dyn NullArrayTrait> {
16        Some(self)
17    }
18
19    fn _as_bool_typed(&self) -> Option<&dyn BoolArrayTrait> {
20        Some(self)
21    }
22
23    fn _as_primitive_typed(&self) -> Option<&dyn PrimitiveArrayTrait> {
24        Some(self)
25    }
26
27    fn _as_decimal_typed(&self) -> Option<&dyn DecimalArrayTrait> {
28        Some(self)
29    }
30
31    fn _as_utf8_typed(&self) -> Option<&dyn Utf8ArrayTrait> {
32        Some(self)
33    }
34
35    fn _as_binary_typed(&self) -> Option<&dyn BinaryArrayTrait> {
36        Some(self)
37    }
38
39    fn _as_struct_typed(&self) -> Option<&dyn StructArrayTrait> {
40        Some(self)
41    }
42
43    fn _as_list_typed(&self) -> Option<&dyn ListArrayTrait> {
44        Some(self)
45    }
46
47    fn _as_extension_typed(&self) -> Option<&dyn ExtensionArrayTrait> {
48        Some(self)
49    }
50}
51
52impl NullArrayTrait for ChunkedArray {}
53
54impl BoolArrayTrait for ChunkedArray {}
55
56impl PrimitiveArrayTrait for ChunkedArray {}
57
58impl DecimalArrayTrait for ChunkedArray {}
59
60impl Utf8ArrayTrait for ChunkedArray {}
61
62impl BinaryArrayTrait for ChunkedArray {}
63
64impl StructArrayTrait for ChunkedArray {
65    fn maybe_null_field_by_idx(&self, idx: usize) -> VortexResult<ArrayRef> {
66        let mut chunks = Vec::with_capacity(self.nchunks());
67        for chunk in self.chunks() {
68            chunks.push(
69                chunk
70                    .as_struct_typed()
71                    .ok_or_else(|| vortex_err!("Chunk was not a StructArray"))?
72                    .maybe_null_field_by_idx(idx)?,
73            );
74        }
75
76        let projected_dtype = self
77            .dtype()
78            .as_struct()
79            .ok_or_else(|| vortex_err!("Not a struct dtype"))?
80            .field_by_index(idx)?;
81        let chunked = ChunkedArray::try_new(chunks, projected_dtype.clone())
82            .unwrap_or_else(|err| {
83                vortex_panic!(
84                    err,
85                    "Failed to create new chunked array with dtype {}",
86                    projected_dtype
87                )
88            })
89            .into_array();
90        Ok(chunked)
91    }
92
93    fn project(&self, projection: &[FieldName]) -> VortexResult<ArrayRef> {
94        let mut chunks = Vec::with_capacity(self.nchunks());
95        for chunk in self.chunks() {
96            chunks.push(
97                chunk
98                    .as_struct_typed()
99                    .ok_or_else(|| vortex_err!("Chunk was not a StructArray"))?
100                    .project(projection)?,
101            );
102        }
103
104        let projected_dtype = self
105            .dtype()
106            .as_struct()
107            .ok_or_else(|| vortex_err!("Not a struct dtype"))?
108            .project(projection)?;
109        Ok(ChunkedArray::new_unchecked(
110            chunks,
111            DType::Struct(Arc::new(projected_dtype), self.dtype().nullability()),
112        )
113        .into_array())
114    }
115}
116
117impl ListArrayTrait for ChunkedArray {}
118
119impl ExtensionArrayTrait for ChunkedArray {
120    fn storage_data(&self) -> ArrayRef {
121        ChunkedArray::new_unchecked(
122            self.chunks()
123                .iter()
124                .map(|chunk| {
125                    chunk
126                        .as_extension_typed()
127                        .vortex_expect("Expected extension array")
128                        .storage_data()
129                })
130                .collect(),
131            self.ext_dtype().storage_dtype().clone(),
132        )
133        .into_array()
134    }
135}
vortex_array/arrays/chunked/variants.rs

vortex_array/arrays/chunked/
variants.rs
