vortex_array/arrays/list/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::DType;
5use vortex_dtype::Nullability;
6use vortex_dtype::PType;
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9use vortex_error::vortex_bail;
10use vortex_error::vortex_ensure;
11use vortex_session::VortexSession;
12
13use crate::Array;
14use crate::ArrayRef;
15use crate::ExecutionCtx;
16use crate::IntoArray;
17use crate::ProstMetadata;
18use crate::arrays::ListArray;
19use crate::arrays::list::compute::PARENT_KERNELS;
20use crate::arrays::list::compute::rules::PARENT_RULES;
21use crate::arrays::list_view_from_list;
22use crate::buffer::BufferHandle;
23use crate::metadata::DeserializeMetadata;
24use crate::metadata::SerializeMetadata;
25use crate::serde::ArrayChildren;
26use crate::validity::Validity;
27use crate::vtable;
28use crate::vtable::ArrayId;
29use crate::vtable::VTable;
30use crate::vtable::ValidityVTableFromValidityHelper;
31
32mod array;
33mod operations;
34mod validity;
35mod visitor;
36
37vtable!(List);
38
39#[derive(Clone, prost::Message)]
40pub struct ListMetadata {
41    #[prost(uint64, tag = "1")]
42    elements_len: u64,
43    #[prost(enumeration = "PType", tag = "2")]
44    offset_ptype: i32,
45}
46
47impl VTable for ListVTable {
48    type Array = ListArray;
49
50    type Metadata = ProstMetadata<ListMetadata>;
51
52    type ArrayVTable = Self;
53    type OperationsVTable = Self;
54    type ValidityVTable = ValidityVTableFromValidityHelper;
55    type VisitorVTable = Self;
56
57    fn id(_array: &Self::Array) -> ArrayId {
58        Self::ID
59    }
60
61    fn reduce_parent(
62        array: &Self::Array,
63        parent: &ArrayRef,
64        child_idx: usize,
65    ) -> VortexResult<Option<ArrayRef>> {
66        PARENT_RULES.evaluate(array, parent, child_idx)
67    }
68
69    fn metadata(array: &ListArray) -> VortexResult<Self::Metadata> {
70        Ok(ProstMetadata(ListMetadata {
71            elements_len: array.elements().len() as u64,
72            offset_ptype: PType::try_from(array.offsets().dtype())? as i32,
73        }))
74    }
75
76    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
77        Ok(Some(SerializeMetadata::serialize(metadata)))
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        Ok(ProstMetadata(
88            <ProstMetadata<ListMetadata> as DeserializeMetadata>::deserialize(bytes)?,
89        ))
90    }
91
92    fn build(
93        dtype: &DType,
94        len: usize,
95        metadata: &Self::Metadata,
96        _buffers: &[BufferHandle],
97        children: &dyn ArrayChildren,
98    ) -> VortexResult<ListArray> {
99        let validity = if children.len() == 2 {
100            Validity::from(dtype.nullability())
101        } else if children.len() == 3 {
102            let validity = children.get(2, &Validity::DTYPE, len)?;
103            Validity::Array(validity)
104        } else {
105            vortex_bail!("Expected 2 or 3 children, got {}", children.len());
106        };
107
108        let DType::List(element_dtype, _) = &dtype else {
109            vortex_bail!("Expected List dtype, got {:?}", dtype);
110        };
111        let elements = children.get(
112            0,
113            element_dtype.as_ref(),
114            usize::try_from(metadata.0.elements_len)?,
115        )?;
116
117        let offsets = children.get(
118            1,
119            &DType::Primitive(metadata.0.offset_ptype(), Nullability::NonNullable),
120            len + 1,
121        )?;
122
123        ListArray::try_new(elements, offsets, validity)
124    }
125
126    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
127        vortex_ensure!(
128            children.len() == 2 || children.len() == 3,
129            "ListArray expects 2 or 3 children, got {}",
130            children.len()
131        );
132
133        let mut iter = children.into_iter();
134        let elements = iter
135            .next()
136            .vortex_expect("children length already validated");
137        let offsets = iter
138            .next()
139            .vortex_expect("children length already validated");
140        let validity = if let Some(validity_array) = iter.next() {
141            Validity::Array(validity_array)
142        } else {
143            Validity::from(array.dtype.nullability())
144        };
145
146        let new_array = ListArray::try_new(elements, offsets, validity)?;
147        *array = new_array;
148        Ok(())
149    }
150
151    fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
152        Ok(list_view_from_list(array.clone(), ctx)?.into_array())
153    }
154
155    fn execute_parent(
156        array: &Self::Array,
157        parent: &ArrayRef,
158        child_idx: usize,
159        ctx: &mut ExecutionCtx,
160    ) -> VortexResult<Option<ArrayRef>> {
161        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
162    }
163}
164
165#[derive(Debug)]
166pub struct ListVTable;
167
168impl ListVTable {
169    pub const ID: ArrayId = ArrayId::new_ref("vortex.list");
170}
vortex_array/arrays/list/vtable/mod.rs

vortex_array/arrays/list/vtable/
mod.rs
