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

vortex_array/arrays/listview/vtable/
mod.rs
