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

vortex_array/arrays/list/vtable/
mod.rs
