vortex_array/arrays/listview/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::hash::Hash;
5
6use vortex_error::VortexExpect;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_error::vortex_ensure;
10use vortex_error::vortex_panic;
11use vortex_session::VortexSession;
12
13use crate::ArrayRef;
14use crate::DeserializeMetadata;
15use crate::ExecutionCtx;
16use crate::Precision;
17use crate::ProstMetadata;
18use crate::SerializeMetadata;
19use crate::arrays::ListViewArray;
20use crate::arrays::listview::compute::rules::PARENT_RULES;
21use crate::buffer::BufferHandle;
22use crate::dtype::DType;
23use crate::dtype::Nullability;
24use crate::dtype::PType;
25use crate::hash::ArrayEq;
26use crate::hash::ArrayHash;
27use crate::serde::ArrayChildren;
28use crate::stats::StatsSetRef;
29use crate::validity::Validity;
30use crate::vtable;
31use crate::vtable::ArrayId;
32use crate::vtable::VTable;
33use crate::vtable::ValidityVTableFromValidityHelper;
34use crate::vtable::validity_nchildren;
35use crate::vtable::validity_to_child;
36mod operations;
37mod validity;
38vtable!(ListView);
39
40#[derive(Debug)]
41pub struct ListViewVTable;
42
43impl ListViewVTable {
44    pub const ID: ArrayId = ArrayId::new_ref("vortex.listview");
45}
46
47#[derive(Clone, prost::Message)]
48pub struct ListViewMetadata {
49    #[prost(uint64, tag = "1")]
50    elements_len: u64,
51    #[prost(enumeration = "PType", tag = "2")]
52    offset_ptype: i32,
53    #[prost(enumeration = "PType", tag = "3")]
54    size_ptype: i32,
55}
56
57impl VTable for ListViewVTable {
58    type Array = ListViewArray;
59
60    type Metadata = ProstMetadata<ListViewMetadata>;
61    type OperationsVTable = Self;
62    type ValidityVTable = ValidityVTableFromValidityHelper;
63    fn id(_array: &Self::Array) -> ArrayId {
64        Self::ID
65    }
66
67    fn len(array: &ListViewArray) -> usize {
68        debug_assert_eq!(array.offsets().len(), array.sizes().len());
69        array.offsets().len()
70    }
71
72    fn dtype(array: &ListViewArray) -> &DType {
73        &array.dtype
74    }
75
76    fn stats(array: &ListViewArray) -> StatsSetRef<'_> {
77        array.stats_set.to_ref(array.as_ref())
78    }
79
80    fn array_hash<H: std::hash::Hasher>(
81        array: &ListViewArray,
82        state: &mut H,
83        precision: Precision,
84    ) {
85        array.dtype.hash(state);
86        array.elements().array_hash(state, precision);
87        array.offsets().array_hash(state, precision);
88        array.sizes().array_hash(state, precision);
89        array.validity.array_hash(state, precision);
90    }
91
92    fn array_eq(array: &ListViewArray, other: &ListViewArray, precision: Precision) -> bool {
93        array.dtype == other.dtype
94            && array.elements().array_eq(other.elements(), precision)
95            && array.offsets().array_eq(other.offsets(), precision)
96            && array.sizes().array_eq(other.sizes(), precision)
97            && array.validity.array_eq(&other.validity, precision)
98    }
99
100    fn nbuffers(_array: &ListViewArray) -> usize {
101        0
102    }
103
104    fn buffer(_array: &ListViewArray, idx: usize) -> BufferHandle {
105        vortex_panic!("ListViewArray buffer index {idx} out of bounds")
106    }
107
108    fn buffer_name(_array: &ListViewArray, idx: usize) -> Option<String> {
109        vortex_panic!("ListViewArray buffer_name index {idx} out of bounds")
110    }
111
112    fn nchildren(array: &ListViewArray) -> usize {
113        3 + validity_nchildren(&array.validity)
114    }
115
116    fn child(array: &ListViewArray, idx: usize) -> ArrayRef {
117        match idx {
118            0 => array.elements().clone(),
119            1 => array.offsets().clone(),
120            2 => array.sizes().clone(),
121            3 => validity_to_child(&array.validity, array.len())
122                .vortex_expect("ListViewArray validity child out of bounds"),
123            _ => vortex_panic!("ListViewArray child index {idx} out of bounds"),
124        }
125    }
126
127    fn child_name(_array: &ListViewArray, idx: usize) -> String {
128        match idx {
129            0 => "elements".to_string(),
130            1 => "offsets".to_string(),
131            2 => "sizes".to_string(),
132            3 => "validity".to_string(),
133            _ => vortex_panic!("ListViewArray child_name index {idx} out of bounds"),
134        }
135    }
136
137    fn metadata(array: &ListViewArray) -> VortexResult<Self::Metadata> {
138        Ok(ProstMetadata(ListViewMetadata {
139            elements_len: array.elements().len() as u64,
140            offset_ptype: PType::try_from(array.offsets().dtype())? as i32,
141            size_ptype: PType::try_from(array.sizes().dtype())? as i32,
142        }))
143    }
144
145    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
146        Ok(Some(metadata.serialize()))
147    }
148
149    fn deserialize(
150        bytes: &[u8],
151        _dtype: &DType,
152        _len: usize,
153        _buffers: &[BufferHandle],
154        _session: &VortexSession,
155    ) -> VortexResult<Self::Metadata> {
156        let metadata = <Self::Metadata as DeserializeMetadata>::deserialize(bytes)?;
157        Ok(ProstMetadata(metadata))
158    }
159
160    fn build(
161        dtype: &DType,
162        len: usize,
163        metadata: &Self::Metadata,
164        buffers: &[BufferHandle],
165        children: &dyn ArrayChildren,
166    ) -> VortexResult<ListViewArray> {
167        vortex_ensure!(
168            buffers.is_empty(),
169            "`ListViewArray::build` expects no buffers"
170        );
171
172        let DType::List(element_dtype, _) = dtype else {
173            vortex_bail!("Expected List dtype, got {:?}", dtype);
174        };
175
176        let validity = if children.len() == 3 {
177            Validity::from(dtype.nullability())
178        } else if children.len() == 4 {
179            let validity = children.get(3, &Validity::DTYPE, len)?;
180            Validity::Array(validity)
181        } else {
182            vortex_bail!(
183                "`ListViewArray::build` expects 3 or 4 children, got {}",
184                children.len()
185            );
186        };
187
188        // Get elements with the correct length from metadata.
189        let elements = children.get(
190            0,
191            element_dtype.as_ref(),
192            usize::try_from(metadata.0.elements_len)?,
193        )?;
194
195        // Get offsets with proper type from metadata.
196        let offsets = children.get(
197            1,
198            &DType::Primitive(metadata.0.offset_ptype(), Nullability::NonNullable),
199            len,
200        )?;
201
202        // Get sizes with proper type from metadata.
203        let sizes = children.get(
204            2,
205            &DType::Primitive(metadata.0.size_ptype(), Nullability::NonNullable),
206            len,
207        )?;
208
209        ListViewArray::try_new(elements, offsets, sizes, validity)
210    }
211
212    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
213        vortex_ensure!(
214            children.len() == 3 || children.len() == 4,
215            "ListViewArray expects 3 or 4 children, got {}",
216            children.len()
217        );
218
219        let mut iter = children.into_iter();
220        let elements = iter
221            .next()
222            .vortex_expect("children length already validated");
223        let offsets = iter
224            .next()
225            .vortex_expect("children length already validated");
226        let sizes = iter
227            .next()
228            .vortex_expect("children length already validated");
229        let validity = if let Some(validity_array) = iter.next() {
230            Validity::Array(validity_array)
231        } else {
232            Validity::from(array.dtype.nullability())
233        };
234
235        let new_array = ListViewArray::try_new(elements, offsets, sizes, validity)?;
236        *array = new_array;
237        Ok(())
238    }
239
240    fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
241        Ok(array.to_array())
242    }
243
244    fn reduce_parent(
245        array: &Self::Array,
246        parent: &ArrayRef,
247        child_idx: usize,
248    ) -> VortexResult<Option<ArrayRef>> {
249        PARENT_RULES.evaluate(array, parent, child_idx)
250    }
251}
vortex_array/arrays/listview/vtable/mod.rs

vortex_array/arrays/listview/vtable/
mod.rs
