vortex_array/arrays/list/
serde.rs

1use vortex_buffer::ByteBuffer;
2use vortex_dtype::{DType, Nullability, PType};
3use vortex_error::{VortexExpect, VortexResult, vortex_bail};
4
5use super::{ListArray, ListVTable};
6use crate::arrays::ListEncoding;
7use crate::serde::ArrayChildren;
8use crate::validity::Validity;
9use crate::vtable::{SerdeVTable, ValidityHelper, VisitorVTable};
10use crate::{Array, ArrayBufferVisitor, ArrayChildVisitor, ProstMetadata};
11
12#[derive(Clone, prost::Message)]
13pub struct ListMetadata {
14    #[prost(uint64, tag = "1")]
15    elements_len: u64,
16    #[prost(enumeration = "PType", tag = "2")]
17    offset_ptype: i32,
18}
19
20impl SerdeVTable<ListVTable> for ListVTable {
21    type Metadata = ProstMetadata<ListMetadata>;
22
23    fn metadata(array: &ListArray) -> VortexResult<Option<Self::Metadata>> {
24        Ok(Some(ProstMetadata(ListMetadata {
25            elements_len: array.elements().len() as u64,
26            offset_ptype: PType::try_from(array.offsets().dtype())
27                .vortex_expect("Must be a valid PType") as i32,
28        })))
29    }
30
31    fn build(
32        _encoding: &ListEncoding,
33        dtype: &DType,
34        len: usize,
35        metadata: &ListMetadata,
36        _buffers: &[ByteBuffer],
37        children: &dyn ArrayChildren,
38    ) -> VortexResult<ListArray> {
39        let validity = if children.len() == 2 {
40            Validity::from(dtype.nullability())
41        } else if children.len() == 3 {
42            let validity = children.get(2, &Validity::DTYPE, len)?;
43            Validity::Array(validity)
44        } else {
45            vortex_bail!("Expected 2 or 3 children, got {}", children.len());
46        };
47
48        let DType::List(element_dtype, _) = &dtype else {
49            vortex_bail!("Expected List dtype, got {:?}", dtype);
50        };
51        let elements = children.get(
52            0,
53            element_dtype.as_ref(),
54            usize::try_from(metadata.elements_len).vortex_expect("Too many elements"),
55        )?;
56
57        let offsets = children.get(
58            1,
59            &DType::Primitive(metadata.offset_ptype(), Nullability::NonNullable),
60            len + 1,
61        )?;
62
63        ListArray::try_new(elements, offsets, validity)
64    }
65}
66
67impl VisitorVTable<ListVTable> for ListVTable {
68    fn visit_buffers(_array: &ListArray, _visitor: &mut dyn ArrayBufferVisitor) {}
69
70    fn visit_children(array: &ListArray, visitor: &mut dyn ArrayChildVisitor) {
71        visitor.visit_child("elements", array.elements());
72        visitor.visit_child("offsets", array.offsets());
73        visitor.visit_validity(array.validity(), array.len());
74    }
75}