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_vector::Vector;
12
13use crate::ArrayRef;
14use crate::ExecutionCtx;
15use crate::ProstMetadata;
16use crate::arrays::ListArray;
17use crate::arrays::list::vtable::kernel::PARENT_KERNELS;
18use crate::buffer::BufferHandle;
19use crate::metadata::DeserializeMetadata;
20use crate::metadata::SerializeMetadata;
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 kernel;
34mod operations;
35mod validity;
36mod visitor;
37
38vtable!(List);
39
40#[derive(Clone, prost::Message)]
41pub struct ListMetadata {
42    #[prost(uint64, tag = "1")]
43    elements_len: u64,
44    #[prost(enumeration = "PType", tag = "2")]
45    offset_ptype: i32,
46}
47
48impl VTable for ListVTable {
49    type Array = ListArray;
50
51    type Metadata = ProstMetadata<ListMetadata>;
52
53    type ArrayVTable = Self;
54    type CanonicalVTable = Self;
55    type OperationsVTable = Self;
56    type ValidityVTable = ValidityVTableFromValidityHelper;
57    type VisitorVTable = Self;
58    type ComputeVTable = NotSupported;
59    type EncodeVTable = NotSupported;
60
61    fn id(&self) -> ArrayId {
62        ArrayId::new_ref("vortex.list")
63    }
64
65    fn encoding(_array: &Self::Array) -> ArrayVTable {
66        ListVTable.as_vtable()
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(bytes: &[u8]) -> VortexResult<Self::Metadata> {
81        Ok(ProstMetadata(
82            <ProstMetadata<ListMetadata> as DeserializeMetadata>::deserialize(bytes)?,
83        ))
84    }
85
86    fn build(
87        &self,
88        dtype: &DType,
89        len: usize,
90        metadata: &Self::Metadata,
91        _buffers: &[BufferHandle],
92        children: &dyn ArrayChildren,
93    ) -> VortexResult<ListArray> {
94        let validity = if children.len() == 2 {
95            Validity::from(dtype.nullability())
96        } else if children.len() == 3 {
97            let validity = children.get(2, &Validity::DTYPE, len)?;
98            Validity::Array(validity)
99        } else {
100            vortex_bail!("Expected 2 or 3 children, got {}", children.len());
101        };
102
103        let DType::List(element_dtype, _) = &dtype else {
104            vortex_bail!("Expected List dtype, got {:?}", dtype);
105        };
106        let elements = children.get(
107            0,
108            element_dtype.as_ref(),
109            usize::try_from(metadata.0.elements_len)?,
110        )?;
111
112        let offsets = children.get(
113            1,
114            &DType::Primitive(metadata.0.offset_ptype(), Nullability::NonNullable),
115            len + 1,
116        )?;
117
118        ListArray::try_new(elements, offsets, validity)
119    }
120
121    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
122        vortex_ensure!(
123            children.len() == 2 || children.len() == 3,
124            "ListArray expects 2 or 3 children, got {}",
125            children.len()
126        );
127
128        let mut iter = children.into_iter();
129        let elements = iter
130            .next()
131            .vortex_expect("children length already validated");
132        let offsets = iter
133            .next()
134            .vortex_expect("children length already validated");
135        let validity = if let Some(validity_array) = iter.next() {
136            Validity::Array(validity_array)
137        } else {
138            Validity::from(array.dtype.nullability())
139        };
140
141        let new_array = ListArray::try_new(elements, offsets, validity)?;
142        *array = new_array;
143        Ok(())
144    }
145
146    fn execute_parent(
147        array: &Self::Array,
148        parent: &ArrayRef,
149        child_idx: usize,
150        ctx: &mut ExecutionCtx,
151    ) -> VortexResult<Option<Vector>> {
152        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
153    }
154}
155
156#[derive(Debug)]
157pub struct ListVTable;
vortex_array/arrays/list/vtable/mod.rs

vortex_array/arrays/list/vtable/
mod.rs
