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, ExtensionArrayTrait, ListArrayTrait, NullArrayTrait,
9    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_utf8_typed(&self) -> Option<&dyn Utf8ArrayTrait> {
28        Some(self)
29    }
30
31    fn _as_binary_typed(&self) -> Option<&dyn BinaryArrayTrait> {
32        Some(self)
33    }
34
35    fn _as_struct_typed(&self) -> Option<&dyn StructArrayTrait> {
36        Some(self)
37    }
38
39    fn _as_list_typed(&self) -> Option<&dyn ListArrayTrait> {
40        Some(self)
41    }
42
43    fn _as_extension_typed(&self) -> Option<&dyn ExtensionArrayTrait> {
44        Some(self)
45    }
46}
47
48impl NullArrayTrait for ChunkedArray {}
49
50impl BoolArrayTrait for ChunkedArray {}
51
52impl PrimitiveArrayTrait for ChunkedArray {}
53
54impl Utf8ArrayTrait for ChunkedArray {}
55
56impl BinaryArrayTrait for ChunkedArray {}
57
58impl StructArrayTrait for ChunkedArray {
59    fn maybe_null_field_by_idx(&self, idx: usize) -> VortexResult<ArrayRef> {
60        let mut chunks = Vec::with_capacity(self.nchunks());
61        for chunk in self.chunks() {
62            chunks.push(
63                chunk
64                    .as_struct_typed()
65                    .ok_or_else(|| vortex_err!("Chunk was not a StructArray"))?
66                    .maybe_null_field_by_idx(idx)?,
67            );
68        }
69
70        let projected_dtype = self
71            .dtype()
72            .as_struct()
73            .ok_or_else(|| vortex_err!("Not a struct dtype"))?
74            .field_by_index(idx)?;
75        let chunked = ChunkedArray::try_new(chunks, projected_dtype.clone())
76            .unwrap_or_else(|err| {
77                vortex_panic!(
78                    err,
79                    "Failed to create new chunked array with dtype {}",
80                    projected_dtype
81                )
82            })
83            .into_array();
84        Ok(chunked)
85    }
86
87    fn project(&self, projection: &[FieldName]) -> VortexResult<ArrayRef> {
88        let mut chunks = Vec::with_capacity(self.nchunks());
89        for chunk in self.chunks() {
90            chunks.push(
91                chunk
92                    .as_struct_typed()
93                    .ok_or_else(|| vortex_err!("Chunk was not a StructArray"))?
94                    .project(projection)?,
95            );
96        }
97
98        let projected_dtype = self
99            .dtype()
100            .as_struct()
101            .ok_or_else(|| vortex_err!("Not a struct dtype"))?
102            .project(projection)?;
103        Ok(ChunkedArray::new_unchecked(
104            chunks,
105            DType::Struct(Arc::new(projected_dtype), self.dtype().nullability()),
106        )
107        .into_array())
108    }
109}
110
111impl ListArrayTrait for ChunkedArray {}
112
113impl ExtensionArrayTrait for ChunkedArray {
114    fn storage_data(&self) -> ArrayRef {
115        ChunkedArray::new_unchecked(
116            self.chunks()
117                .iter()
118                .map(|chunk| {
119                    chunk
120                        .as_extension_typed()
121                        .vortex_expect("Expected extension array")
122                        .storage_data()
123                })
124                .collect(),
125            self.ext_dtype().storage_dtype().clone(),
126        )
127        .into_array()
128    }
129}
vortex_array/arrays/chunked/variants.rs

vortex_array/arrays/chunked/
variants.rs
