vortex_array/arrays/struct_/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use itertools::Itertools;
7use kernel::PARENT_KERNELS;
8use vortex_error::VortexExpect;
9use vortex_error::VortexResult;
10use vortex_error::vortex_bail;
11use vortex_error::vortex_ensure;
12use vortex_error::vortex_panic;
13use vortex_session::VortexSession;
14
15use crate::ArrayRef;
16use crate::EmptyMetadata;
17use crate::ExecutionCtx;
18use crate::ExecutionStep;
19use crate::IntoArray;
20use crate::arrays::StructArray;
21use crate::arrays::struct_::compute::rules::PARENT_RULES;
22use crate::buffer::BufferHandle;
23use crate::dtype::DType;
24use crate::serde::ArrayChildren;
25use crate::validity::Validity;
26use crate::vtable;
27use crate::vtable::VTable;
28use crate::vtable::ValidityVTableFromValidityHelper;
29use crate::vtable::validity_nchildren;
30use crate::vtable::validity_to_child;
31mod kernel;
32mod operations;
33mod validity;
34use std::hash::Hash;
35
36use crate::Precision;
37use crate::hash::ArrayEq;
38use crate::hash::ArrayHash;
39use crate::stats::StatsSetRef;
40use crate::vtable::ArrayId;
41
42vtable!(Struct);
43
44impl VTable for StructVTable {
45    type Array = StructArray;
46
47    type Metadata = EmptyMetadata;
48    type OperationsVTable = Self;
49    type ValidityVTable = ValidityVTableFromValidityHelper;
50    fn id(_array: &Self::Array) -> ArrayId {
51        Self::ID
52    }
53
54    fn len(array: &StructArray) -> usize {
55        array.len
56    }
57
58    fn dtype(array: &StructArray) -> &DType {
59        &array.dtype
60    }
61
62    fn stats(array: &StructArray) -> StatsSetRef<'_> {
63        array.stats_set.to_ref(array.as_ref())
64    }
65
66    fn array_hash<H: std::hash::Hasher>(array: &StructArray, state: &mut H, precision: Precision) {
67        array.len.hash(state);
68        array.dtype.hash(state);
69        for field in array.fields.iter() {
70            field.array_hash(state, precision);
71        }
72        array.validity.array_hash(state, precision);
73    }
74
75    fn array_eq(array: &StructArray, other: &StructArray, precision: Precision) -> bool {
76        array.len == other.len
77            && array.dtype == other.dtype
78            && array.fields.len() == other.fields.len()
79            && array
80                .fields
81                .iter()
82                .zip(other.fields.iter())
83                .all(|(a, b)| a.array_eq(b, precision))
84            && array.validity.array_eq(&other.validity, precision)
85    }
86
87    fn nbuffers(_array: &StructArray) -> usize {
88        0
89    }
90
91    fn buffer(_array: &StructArray, idx: usize) -> BufferHandle {
92        vortex_panic!("StructArray buffer index {idx} out of bounds")
93    }
94
95    fn buffer_name(_array: &StructArray, idx: usize) -> Option<String> {
96        vortex_panic!("StructArray buffer_name index {idx} out of bounds")
97    }
98
99    fn nchildren(array: &StructArray) -> usize {
100        validity_nchildren(&array.validity) + array.unmasked_fields().len()
101    }
102
103    fn child(array: &StructArray, idx: usize) -> ArrayRef {
104        let vc = validity_nchildren(&array.validity);
105        if idx < vc {
106            validity_to_child(&array.validity, array.len())
107                .vortex_expect("StructArray validity child out of bounds")
108        } else {
109            array.unmasked_fields()[idx - vc].clone()
110        }
111    }
112
113    fn child_name(array: &StructArray, idx: usize) -> String {
114        let vc = validity_nchildren(&array.validity);
115        if idx < vc {
116            "validity".to_string()
117        } else {
118            array.names()[idx - vc].as_ref().to_string()
119        }
120    }
121
122    fn metadata(_array: &StructArray) -> VortexResult<Self::Metadata> {
123        Ok(EmptyMetadata)
124    }
125
126    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
127        Ok(Some(vec![]))
128    }
129
130    fn deserialize(
131        _bytes: &[u8],
132        _dtype: &DType,
133        _len: usize,
134        _buffers: &[BufferHandle],
135        _session: &VortexSession,
136    ) -> VortexResult<Self::Metadata> {
137        Ok(EmptyMetadata)
138    }
139
140    fn build(
141        dtype: &DType,
142        len: usize,
143        _metadata: &Self::Metadata,
144        _buffers: &[BufferHandle],
145        children: &dyn ArrayChildren,
146    ) -> VortexResult<StructArray> {
147        let DType::Struct(struct_dtype, nullability) = dtype else {
148            vortex_bail!("Expected struct dtype, found {:?}", dtype)
149        };
150
151        let (validity, non_data_children) = if children.len() == struct_dtype.nfields() {
152            (Validity::from(*nullability), 0_usize)
153        } else if children.len() == struct_dtype.nfields() + 1 {
154            // Validity is the first child if it exists.
155            let validity = children.get(0, &Validity::DTYPE, len)?;
156            (Validity::Array(validity), 1_usize)
157        } else {
158            vortex_bail!(
159                "Expected {} or {} children, found {}",
160                struct_dtype.nfields(),
161                struct_dtype.nfields() + 1,
162                children.len()
163            );
164        };
165
166        let children: Vec<_> = (0..struct_dtype.nfields())
167            .map(|i| {
168                let child_dtype = struct_dtype
169                    .field_by_index(i)
170                    .vortex_expect("no out of bounds");
171                children.get(non_data_children + i, &child_dtype, len)
172            })
173            .try_collect()?;
174
175        StructArray::try_new_with_dtype(children, struct_dtype.clone(), len, validity)
176    }
177
178    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
179        let DType::Struct(struct_dtype, _nullability) = &array.dtype else {
180            vortex_bail!("Expected struct dtype, found {:?}", array.dtype)
181        };
182
183        // First child is validity (if present), followed by fields
184        let (validity, non_data_children) = if children.len() == struct_dtype.nfields() {
185            (array.validity.clone(), 0_usize)
186        } else if children.len() == struct_dtype.nfields() + 1 {
187            (Validity::Array(children[0].clone()), 1_usize)
188        } else {
189            vortex_bail!(
190                "Expected {} or {} children, found {}",
191                struct_dtype.nfields(),
192                struct_dtype.nfields() + 1,
193                children.len()
194            );
195        };
196
197        let fields: Arc<[ArrayRef]> = children.into_iter().skip(non_data_children).collect();
198        vortex_ensure!(
199            fields.len() == struct_dtype.nfields(),
200            "Expected {} field children, found {}",
201            struct_dtype.nfields(),
202            fields.len()
203        );
204
205        array.fields = fields;
206        array.validity = validity;
207        Ok(())
208    }
209
210    fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionStep> {
211        Ok(ExecutionStep::Done(array.clone().into_array()))
212    }
213
214    fn reduce_parent(
215        array: &Self::Array,
216        parent: &ArrayRef,
217        child_idx: usize,
218    ) -> VortexResult<Option<ArrayRef>> {
219        PARENT_RULES.evaluate(array, parent, child_idx)
220    }
221
222    fn execute_parent(
223        array: &Self::Array,
224        parent: &ArrayRef,
225        child_idx: usize,
226        ctx: &mut ExecutionCtx,
227    ) -> VortexResult<Option<ArrayRef>> {
228        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
229    }
230}
231
232#[derive(Debug)]
233pub struct StructVTable;
234
235impl StructVTable {
236    pub const ID: ArrayId = ArrayId::new_ref("vortex.struct");
237}
vortex_array/arrays/struct_/vtable/mod.rs

vortex_array/arrays/struct_/vtable/
mod.rs
