vortex_array/arrays/struct_/
serde.rs

1use 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            // Validity is the first child if it exists.
44            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}