vortex_array/arrays/listview/vtable/
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::{VortexResult, vortex_bail, vortex_ensure};
7
8use crate::arrays::{ListViewArray, ListViewEncoding, ListViewVTable};
9use crate::serde::ArrayChildren;
10use crate::validity::Validity;
11use crate::vtable::SerdeVTable;
12use crate::{Array, ProstMetadata};
13
14#[derive(Clone, prost::Message)]
15pub struct ListViewMetadata {
16    #[prost(uint64, tag = "1")]
17    elements_len: u64,
18    #[prost(enumeration = "PType", tag = "2")]
19    offset_ptype: i32,
20    #[prost(enumeration = "PType", tag = "3")]
21    size_ptype: i32,
22}
23
24impl SerdeVTable<ListViewVTable> for ListViewVTable {
25    type Metadata = ProstMetadata<ListViewMetadata>;
26
27    fn metadata(array: &ListViewArray) -> VortexResult<Option<Self::Metadata>> {
28        Ok(Some(ProstMetadata(ListViewMetadata {
29            elements_len: array.elements().len() as u64,
30            offset_ptype: PType::try_from(array.offsets().dtype())? as i32,
31            size_ptype: PType::try_from(array.sizes().dtype())? as i32,
32        })))
33    }
34
35    fn build(
36        _encoding: &ListViewEncoding,
37        dtype: &DType,
38        len: usize,
39        metadata: &ListViewMetadata,
40        buffers: &[ByteBuffer],
41        children: &dyn ArrayChildren,
42    ) -> VortexResult<ListViewArray> {
43        vortex_ensure!(
44            buffers.is_empty(),
45            "`ListViewArray::build` expects no buffers"
46        );
47
48        let DType::List(element_dtype, _) = dtype else {
49            vortex_bail!("Expected List dtype, got {:?}", dtype);
50        };
51
52        let validity = if children.len() == 3 {
53            Validity::from(dtype.nullability())
54        } else if children.len() == 4 {
55            let validity = children.get(3, &Validity::DTYPE, len)?;
56            Validity::Array(validity)
57        } else {
58            vortex_bail!(
59                "`ListViewArray::build` expects 3 or 4 children, got {}",
60                children.len()
61            );
62        };
63
64        // Get elements with the correct length from metadata.
65        let elements = children.get(
66            0,
67            element_dtype.as_ref(),
68            usize::try_from(metadata.elements_len)?,
69        )?;
70
71        // Get offsets with proper type from metadata.
72        let offsets = children.get(
73            1,
74            &DType::Primitive(metadata.offset_ptype(), Nullability::NonNullable),
75            len,
76        )?;
77
78        // Get sizes with proper type from metadata.
79        let sizes = children.get(
80            2,
81            &DType::Primitive(metadata.size_ptype(), Nullability::NonNullable),
82            len,
83        )?;
84
85        ListViewArray::try_new(elements, offsets, sizes, validity)
86    }
87}