vortex_array/arrays/struct_/
serde.rs1use itertools::Itertools;
2use vortex_dtype::DType;
3use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
4
5use crate::arrays::{StructArray, StructEncoding};
6use crate::serde::ArrayParts;
7use crate::validity::Validity;
8use crate::variants::StructArrayTrait;
9use crate::vtable::SerdeVTable;
10use crate::{Array, ArrayChildVisitor, ArrayContext, ArrayRef, ArrayVisitorImpl, EmptyMetadata};
11
12impl ArrayVisitorImpl for StructArray {
13 fn _children(&self, visitor: &mut dyn ArrayChildVisitor) {
14 visitor.visit_validity(self.validity(), self.len());
15 for (idx, name) in self.names().iter().enumerate() {
16 let child = self
17 .maybe_null_field_by_idx(idx)
18 .vortex_expect("no out of bounds");
19 visitor.visit_child(name.as_ref(), &child);
20 }
21 }
22
23 fn _metadata(&self) -> EmptyMetadata {
24 EmptyMetadata
25 }
26}
27
28impl SerdeVTable<&StructArray> for StructEncoding {
29 fn decode(
30 &self,
31 parts: &ArrayParts,
32 ctx: &ArrayContext,
33 dtype: DType,
34 len: usize,
35 ) -> VortexResult<ArrayRef> {
36 let struct_dtype = dtype
37 .as_struct()
38 .ok_or_else(|| vortex_err!("Expected struct dtype, found {:?}", dtype))?;
39
40 let validity = if parts.nchildren() == struct_dtype.nfields() {
41 Validity::from(dtype.nullability())
42 } else if parts.nchildren() == struct_dtype.nfields() + 1 {
43 let validity = parts.child(0).decode(ctx, Validity::DTYPE, len)?;
45 Validity::Array(validity)
46 } else {
47 vortex_bail!(
48 "Expected {} or {} children, found {}",
49 struct_dtype.nfields(),
50 struct_dtype.nfields() + 1,
51 parts.nchildren()
52 );
53 };
54
55 let children = (0..parts.nchildren())
56 .map(|i| {
57 let child_parts = parts.child(i);
58 let child_dtype = struct_dtype
59 .field_by_index(i)
60 .vortex_expect("no out of bounds");
61 child_parts.decode(ctx, child_dtype, len)
62 })
63 .try_collect()?;
64
65 Ok(
66 StructArray::try_new(struct_dtype.names().clone(), children, len, validity)?
67 .into_array(),
68 )
69 }
70}