vortex_array/arrays/listview/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use vortex_buffer::BufferHandle;
7use vortex_dtype::DType;
8use vortex_dtype::Nullability;
9use vortex_dtype::PType;
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_error::vortex_ensure;
13use vortex_vector::Vector;
14use vortex_vector::listview::ListViewVector;
15
16use crate::DeserializeMetadata;
17use crate::ProstMetadata;
18use crate::SerializeMetadata;
19use crate::arrays::ListViewArray;
20use crate::execution::ExecutionCtx;
21use crate::serde::ArrayChildren;
22use crate::validity::Validity;
23use crate::vtable;
24use crate::vtable::ArrayId;
25use crate::vtable::ArrayVTable;
26use crate::vtable::ArrayVTableExt;
27use crate::vtable::NotSupported;
28use crate::vtable::VTable;
29use crate::vtable::ValidityVTableFromValidityHelper;
30
31mod array;
32mod canonical;
33mod operations;
34mod validity;
35mod visitor;
36
37vtable!(ListView);
38
39#[derive(Debug)]
40pub struct ListViewVTable;
41
42#[derive(Clone, prost::Message)]
43pub struct ListViewMetadata {
44    #[prost(uint64, tag = "1")]
45    elements_len: u64,
46    #[prost(enumeration = "PType", tag = "2")]
47    offset_ptype: i32,
48    #[prost(enumeration = "PType", tag = "3")]
49    size_ptype: i32,
50}
51
52impl VTable for ListViewVTable {
53    type Array = ListViewArray;
54
55    type Metadata = ProstMetadata<ListViewMetadata>;
56
57    type ArrayVTable = Self;
58    type CanonicalVTable = Self;
59    type OperationsVTable = Self;
60    type ValidityVTable = ValidityVTableFromValidityHelper;
61    type VisitorVTable = Self;
62    type ComputeVTable = NotSupported;
63    type EncodeVTable = NotSupported;
64
65    fn id(&self) -> ArrayId {
66        ArrayId::new_ref("vortex.listview")
67    }
68
69    fn encoding(_array: &Self::Array) -> ArrayVTable {
70        ListViewVTable.as_vtable()
71    }
72
73    fn metadata(array: &ListViewArray) -> VortexResult<Self::Metadata> {
74        Ok(ProstMetadata(ListViewMetadata {
75            elements_len: array.elements().len() as u64,
76            offset_ptype: PType::try_from(array.offsets().dtype())? as i32,
77            size_ptype: PType::try_from(array.sizes().dtype())? as i32,
78        }))
79    }
80
81    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
82        Ok(Some(metadata.serialize()))
83    }
84
85    fn deserialize(bytes: &[u8]) -> VortexResult<Self::Metadata> {
86        let metadata = <Self::Metadata as DeserializeMetadata>::deserialize(bytes)?;
87        Ok(ProstMetadata(metadata))
88    }
89
90    fn build(
91        &self,
92        dtype: &DType,
93        len: usize,
94        metadata: &Self::Metadata,
95        buffers: &[BufferHandle],
96        children: &dyn ArrayChildren,
97    ) -> VortexResult<ListViewArray> {
98        vortex_ensure!(
99            buffers.is_empty(),
100            "`ListViewArray::build` expects no buffers"
101        );
102
103        let DType::List(element_dtype, _) = dtype else {
104            vortex_bail!("Expected List dtype, got {:?}", dtype);
105        };
106
107        let validity = if children.len() == 3 {
108            Validity::from(dtype.nullability())
109        } else if children.len() == 4 {
110            let validity = children.get(3, &Validity::DTYPE, len)?;
111            Validity::Array(validity)
112        } else {
113            vortex_bail!(
114                "`ListViewArray::build` expects 3 or 4 children, got {}",
115                children.len()
116            );
117        };
118
119        // Get elements with the correct length from metadata.
120        let elements = children.get(
121            0,
122            element_dtype.as_ref(),
123            usize::try_from(metadata.0.elements_len)?,
124        )?;
125
126        // Get offsets with proper type from metadata.
127        let offsets = children.get(
128            1,
129            &DType::Primitive(metadata.0.offset_ptype(), Nullability::NonNullable),
130            len,
131        )?;
132
133        // Get sizes with proper type from metadata.
134        let sizes = children.get(
135            2,
136            &DType::Primitive(metadata.0.size_ptype(), Nullability::NonNullable),
137            len,
138        )?;
139
140        ListViewArray::try_new(elements, offsets, sizes, validity)
141    }
142
143    fn batch_execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
144        Ok(unsafe {
145            ListViewVector::new_unchecked(
146                Arc::new(array.elements().batch_execute(ctx)?),
147                array.offsets().batch_execute(ctx)?.into_primitive(),
148                array.sizes().batch_execute(ctx)?.into_primitive(),
149                array.validity_mask(),
150            )
151        }
152        .into())
153    }
154}
vortex_array/arrays/listview/vtable/mod.rs

vortex_array/arrays/listview/vtable/
mod.rs
