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

vortex_array/arrays/list/vtable/
mod.rs
