vortex_array/arrays/listview/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use kernel::PARENT_KERNELS;
5use vortex_dtype::DType;
6use vortex_dtype::Nullability;
7use vortex_dtype::PType;
8use vortex_error::VortexExpect;
9use vortex_error::VortexResult;
10use vortex_error::vortex_bail;
11use vortex_error::vortex_ensure;
12use vortex_session::VortexSession;
13
14use crate::ArrayRef;
15use crate::DeserializeMetadata;
16use crate::ExecutionCtx;
17use crate::ProstMetadata;
18use crate::SerializeMetadata;
19use crate::arrays::ListViewArray;
20use crate::arrays::listview::compute::rules::PARENT_RULES;
21use crate::buffer::BufferHandle;
22use crate::serde::ArrayChildren;
23use crate::validity::Validity;
24use crate::vtable;
25use crate::vtable::ArrayId;
26use crate::vtable::VTable;
27use crate::vtable::ValidityVTableFromValidityHelper;
28
29mod array;
30mod kernel;
31mod operations;
32mod validity;
33mod visitor;
34
35vtable!(ListView);
36
37#[derive(Debug)]
38pub struct ListViewVTable;
39
40impl ListViewVTable {
41    pub const ID: ArrayId = ArrayId::new_ref("vortex.listview");
42}
43
44#[derive(Clone, prost::Message)]
45pub struct ListViewMetadata {
46    #[prost(uint64, tag = "1")]
47    elements_len: u64,
48    #[prost(enumeration = "PType", tag = "2")]
49    offset_ptype: i32,
50    #[prost(enumeration = "PType", tag = "3")]
51    size_ptype: i32,
52}
53
54impl VTable for ListViewVTable {
55    type Array = ListViewArray;
56
57    type Metadata = ProstMetadata<ListViewMetadata>;
58
59    type ArrayVTable = Self;
60    type OperationsVTable = Self;
61    type ValidityVTable = ValidityVTableFromValidityHelper;
62    type VisitorVTable = Self;
63
64    fn id(_array: &Self::Array) -> ArrayId {
65        Self::ID
66    }
67
68    fn metadata(array: &ListViewArray) -> VortexResult<Self::Metadata> {
69        Ok(ProstMetadata(ListViewMetadata {
70            elements_len: array.elements().len() as u64,
71            offset_ptype: PType::try_from(array.offsets().dtype())? as i32,
72            size_ptype: PType::try_from(array.sizes().dtype())? as i32,
73        }))
74    }
75
76    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
77        Ok(Some(metadata.serialize()))
78    }
79
80    fn deserialize(
81        bytes: &[u8],
82        _dtype: &DType,
83        _len: usize,
84        _buffers: &[BufferHandle],
85        _session: &VortexSession,
86    ) -> VortexResult<Self::Metadata> {
87        let metadata = <Self::Metadata as DeserializeMetadata>::deserialize(bytes)?;
88        Ok(ProstMetadata(metadata))
89    }
90
91    fn build(
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 with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
144        vortex_ensure!(
145            children.len() == 3 || children.len() == 4,
146            "ListViewArray expects 3 or 4 children, got {}",
147            children.len()
148        );
149
150        let mut iter = children.into_iter();
151        let elements = iter
152            .next()
153            .vortex_expect("children length already validated");
154        let offsets = iter
155            .next()
156            .vortex_expect("children length already validated");
157        let sizes = iter
158            .next()
159            .vortex_expect("children length already validated");
160        let validity = if let Some(validity_array) = iter.next() {
161            Validity::Array(validity_array)
162        } else {
163            Validity::from(array.dtype.nullability())
164        };
165
166        let new_array = ListViewArray::try_new(elements, offsets, sizes, validity)?;
167        *array = new_array;
168        Ok(())
169    }
170
171    fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
172        Ok(array.to_array())
173    }
174
175    fn reduce_parent(
176        array: &Self::Array,
177        parent: &ArrayRef,
178        child_idx: usize,
179    ) -> VortexResult<Option<ArrayRef>> {
180        PARENT_RULES.evaluate(array, parent, child_idx)
181    }
182
183    fn execute_parent(
184        array: &Self::Array,
185        parent: &ArrayRef,
186        child_idx: usize,
187        ctx: &mut ExecutionCtx,
188    ) -> VortexResult<Option<ArrayRef>> {
189        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
190    }
191}
vortex_array/arrays/listview/vtable/mod.rs

vortex_array/arrays/listview/vtable/
mod.rs
