vortex_array/arrays/list/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::BufferHandle;
5use vortex_dtype::DType;
6use vortex_dtype::Nullability;
7use vortex_dtype::PType;
8use vortex_error::VortexResult;
9use vortex_error::vortex_bail;
10
11use crate::ProstMetadata;
12use crate::arrays::ListArray;
13use crate::metadata::DeserializeMetadata;
14use crate::metadata::SerializeMetadata;
15use crate::serde::ArrayChildren;
16use crate::validity::Validity;
17use crate::vtable;
18use crate::vtable::ArrayId;
19use crate::vtable::ArrayVTable;
20use crate::vtable::ArrayVTableExt;
21use crate::vtable::NotSupported;
22use crate::vtable::VTable;
23use crate::vtable::ValidityVTableFromValidityHelper;
24
25mod array;
26mod canonical;
27mod operations;
28mod validity;
29mod visitor;
30
31vtable!(List);
32
33#[derive(Clone, prost::Message)]
34pub struct ListMetadata {
35    #[prost(uint64, tag = "1")]
36    elements_len: u64,
37    #[prost(enumeration = "PType", tag = "2")]
38    offset_ptype: i32,
39}
40
41impl VTable for ListVTable {
42    type Array = ListArray;
43
44    type Metadata = ProstMetadata<ListMetadata>;
45
46    type ArrayVTable = Self;
47    type CanonicalVTable = Self;
48    type OperationsVTable = Self;
49    type ValidityVTable = ValidityVTableFromValidityHelper;
50    type VisitorVTable = Self;
51    type ComputeVTable = NotSupported;
52    type EncodeVTable = NotSupported;
53
54    fn id(&self) -> ArrayId {
55        ArrayId::new_ref("vortex.list")
56    }
57
58    fn encoding(_array: &Self::Array) -> ArrayVTable {
59        ListVTable.as_vtable()
60    }
61
62    fn metadata(array: &ListArray) -> VortexResult<Self::Metadata> {
63        Ok(ProstMetadata(ListMetadata {
64            elements_len: array.elements().len() as u64,
65            offset_ptype: PType::try_from(array.offsets().dtype())? as i32,
66        }))
67    }
68
69    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
70        Ok(Some(SerializeMetadata::serialize(metadata)))
71    }
72
73    fn deserialize(bytes: &[u8]) -> VortexResult<Self::Metadata> {
74        Ok(ProstMetadata(
75            <ProstMetadata<ListMetadata> as DeserializeMetadata>::deserialize(bytes)?,
76        ))
77    }
78
79    fn build(
80        &self,
81        dtype: &DType,
82        len: usize,
83        metadata: &Self::Metadata,
84        _buffers: &[BufferHandle],
85        children: &dyn ArrayChildren,
86    ) -> VortexResult<ListArray> {
87        let validity = if children.len() == 2 {
88            Validity::from(dtype.nullability())
89        } else if children.len() == 3 {
90            let validity = children.get(2, &Validity::DTYPE, len)?;
91            Validity::Array(validity)
92        } else {
93            vortex_bail!("Expected 2 or 3 children, got {}", children.len());
94        };
95
96        let DType::List(element_dtype, _) = &dtype else {
97            vortex_bail!("Expected List dtype, got {:?}", dtype);
98        };
99        let elements = children.get(
100            0,
101            element_dtype.as_ref(),
102            usize::try_from(metadata.0.elements_len)?,
103        )?;
104
105        let offsets = children.get(
106            1,
107            &DType::Primitive(metadata.0.offset_ptype(), Nullability::NonNullable),
108            len + 1,
109        )?;
110
111        ListArray::try_new(elements, offsets, validity)
112    }
113}
114
115#[derive(Debug)]
116pub struct ListVTable;
vortex_array/arrays/list/vtable/mod.rs

vortex_array/arrays/list/vtable/
mod.rs
