vortex_array/arrays/struct_/
serde.rs

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