vortex_array/arrays/list/
serde.rs

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

vortex_array/arrays/list/
serde.rs
