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

vortex_array/arrays/list/vtable/
mod.rs
