vortex_array/arrays/struct_/
serde.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use vortex_buffer::ByteBuffer;
6use vortex_dtype::DType;
7use vortex_error::{VortexExpect, VortexResult, vortex_bail};
8
9use super::StructEncoding;
10use crate::arrays::{StructArray, StructVTable};
11use crate::serde::ArrayChildren;
12use crate::validity::Validity;
13use crate::vtable::{SerdeVTable, ValidityHelper, VisitorVTable};
14use crate::{ArrayBufferVisitor, ArrayChildVisitor, EmptyMetadata};
15
16impl SerdeVTable<StructVTable> for StructVTable {
17    type Metadata = EmptyMetadata;
18
19    fn metadata(_array: &StructArray) -> VortexResult<Option<Self::Metadata>> {
20        Ok(Some(EmptyMetadata))
21    }
22
23    fn build(
24        _encoding: &StructEncoding,
25        dtype: &DType,
26        len: usize,
27        _metadata: &Self::Metadata,
28        _buffers: &[ByteBuffer],
29        children: &dyn ArrayChildren,
30    ) -> VortexResult<StructArray> {
31        let DType::Struct(struct_dtype, nullability) = dtype else {
32            vortex_bail!("Expected struct dtype, found {:?}", dtype)
33        };
34
35        let (validity, non_data_children) = if children.len() == struct_dtype.nfields() {
36            (Validity::from(*nullability), 0_usize)
37        } else if children.len() == struct_dtype.nfields() + 1 {
38            // Validity is the first child if it exists.
39            let validity = children.get(0, &Validity::DTYPE, len)?;
40            (Validity::Array(validity), 1_usize)
41        } else {
42            vortex_bail!(
43                "Expected {} or {} children, found {}",
44                struct_dtype.nfields(),
45                struct_dtype.nfields() + 1,
46                children.len()
47            );
48        };
49
50        let children = (0..struct_dtype.nfields())
51            .map(|i| {
52                let child_dtype = struct_dtype
53                    .field_by_index(i)
54                    .vortex_expect("no out of bounds");
55                children.get(non_data_children + i, &child_dtype, len)
56            })
57            .try_collect()?;
58
59        StructArray::try_new_with_dtype(children, struct_dtype.clone(), len, validity)
60    }
61}
62
63impl VisitorVTable<StructVTable> for StructVTable {
64    fn visit_buffers(_array: &StructArray, _visitor: &mut dyn ArrayBufferVisitor) {}
65
66    fn visit_children(array: &StructArray, visitor: &mut dyn ArrayChildVisitor) {
67        visitor.visit_validity(array.validity(), array.len());
68        for (idx, name) in array.names().iter().enumerate() {
69            visitor.visit_child(name.as_ref(), &array.fields()[idx]);
70        }
71    }
72}