vortex_array/arrays/list/
serde.rs

1use vortex_dtype::{DType, Nullability, PType};
2use vortex_error::{VortexExpect, VortexResult, vortex_bail};
3
4use super::{ListArray, ListEncoding};
5use crate::serde::ArrayParts;
6use crate::validity::Validity;
7use crate::vtable::EncodingVTable;
8use crate::{
9    Array, ArrayChildVisitor, ArrayContext, ArrayRef, ArrayVisitorImpl, DeserializeMetadata,
10    EncodingId, RkyvMetadata,
11};
12
13impl EncodingVTable for ListEncoding {
14    fn id(&self) -> EncodingId {
15        EncodingId::new_ref("vortex.list")
16    }
17
18    fn decode(
19        &self,
20        parts: &ArrayParts,
21        ctx: &ArrayContext,
22        dtype: DType,
23        len: usize,
24    ) -> VortexResult<ArrayRef> {
25        let metadata = RkyvMetadata::<ListMetadata>::deserialize(parts.metadata())?;
26
27        let validity = if parts.nchildren() == 2 {
28            Validity::from(dtype.nullability())
29        } else if parts.nchildren() == 3 {
30            let validity = parts.child(2).decode(ctx, Validity::DTYPE, len)?;
31            Validity::Array(validity)
32        } else {
33            vortex_bail!("Expected 2 or 3 children, got {}", parts.nchildren());
34        };
35
36        let DType::List(element_dtype, _) = &dtype else {
37            vortex_bail!("Expected List dtype, got {:?}", dtype);
38        };
39        let elements =
40            parts
41                .child(0)
42                .decode(ctx, element_dtype.as_ref().clone(), metadata.elements_len)?;
43
44        let offsets = parts.child(1).decode(
45            ctx,
46            DType::Primitive(metadata.offset_ptype, Nullability::NonNullable),
47            len + 1,
48        )?;
49
50        Ok(ListArray::try_new(elements, offsets, validity)?.into_array())
51    }
52}
53
54#[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)]
55pub struct ListMetadata {
56    elements_len: usize,
57    offset_ptype: PType,
58}
59
60impl ArrayVisitorImpl<RkyvMetadata<ListMetadata>> for ListArray {
61    fn _visit_children(&self, visitor: &mut dyn ArrayChildVisitor) {
62        visitor.visit_child("elements", self.elements());
63        visitor.visit_child("offsets", self.offsets());
64        visitor.visit_validity(self.validity(), self.len());
65    }
66
67    fn _metadata(&self) -> RkyvMetadata<ListMetadata> {
68        RkyvMetadata(ListMetadata {
69            elements_len: self.elements().len(),
70            offset_ptype: PType::try_from(self.offsets().dtype())
71                .vortex_expect("Must be a valid PType"),
72        })
73    }
74}