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 vortex_dtype::DType;
8use vortex_error::VortexExpect;
9use vortex_error::VortexResult;
10use vortex_error::vortex_bail;
11use vortex_error::vortex_ensure;
12use vortex_vector::Vector;
13use vortex_vector::struct_::StructVector;
14
15use crate::ArrayRef;
16use crate::EmptyMetadata;
17use crate::VectorExecutor;
18use crate::arrays::struct_::StructArray;
19use crate::arrays::struct_::vtable::rules::PARENT_RULES;
20use crate::buffer::BufferHandle;
21use crate::executor::ExecutionCtx;
22use crate::serde::ArrayChildren;
23use crate::validity::Validity;
24use crate::vtable;
25use crate::vtable::ArrayVTableExt;
26use crate::vtable::NotSupported;
27use crate::vtable::VTable;
28use crate::vtable::ValidityVTableFromValidityHelper;
29
30mod array;
31mod canonical;
32mod operations;
33mod rules;
34mod validity;
35mod visitor;
36
37use crate::vtable::ArrayId;
38use crate::vtable::ArrayVTable;
39
40vtable!(Struct);
41
42impl VTable for StructVTable {
43    type Array = StructArray;
44
45    type Metadata = EmptyMetadata;
46
47    type ArrayVTable = Self;
48    type CanonicalVTable = Self;
49    type OperationsVTable = Self;
50    type ValidityVTable = ValidityVTableFromValidityHelper;
51    type VisitorVTable = Self;
52    type ComputeVTable = NotSupported;
53    type EncodeVTable = NotSupported;
54
55    fn id(&self) -> ArrayId {
56        ArrayId::new_ref("vortex.struct")
57    }
58
59    fn encoding(_array: &Self::Array) -> ArrayVTable {
60        StructVTable.as_vtable()
61    }
62
63    fn metadata(_array: &StructArray) -> VortexResult<Self::Metadata> {
64        Ok(EmptyMetadata)
65    }
66
67    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
68        Ok(Some(vec![]))
69    }
70
71    fn deserialize(_buffer: &[u8]) -> VortexResult<Self::Metadata> {
72        Ok(EmptyMetadata)
73    }
74
75    fn build(
76        &self,
77        dtype: &DType,
78        len: usize,
79        _metadata: &Self::Metadata,
80        _buffers: &[BufferHandle],
81        children: &dyn ArrayChildren,
82    ) -> VortexResult<StructArray> {
83        let DType::Struct(struct_dtype, nullability) = dtype else {
84            vortex_bail!("Expected struct dtype, found {:?}", dtype)
85        };
86
87        let (validity, non_data_children) = if children.len() == struct_dtype.nfields() {
88            (Validity::from(*nullability), 0_usize)
89        } else if children.len() == struct_dtype.nfields() + 1 {
90            // Validity is the first child if it exists.
91            let validity = children.get(0, &Validity::DTYPE, len)?;
92            (Validity::Array(validity), 1_usize)
93        } else {
94            vortex_bail!(
95                "Expected {} or {} children, found {}",
96                struct_dtype.nfields(),
97                struct_dtype.nfields() + 1,
98                children.len()
99            );
100        };
101
102        let children: Vec<_> = (0..struct_dtype.nfields())
103            .map(|i| {
104                let child_dtype = struct_dtype
105                    .field_by_index(i)
106                    .vortex_expect("no out of bounds");
107                children.get(non_data_children + i, &child_dtype, len)
108            })
109            .try_collect()?;
110
111        StructArray::try_new_with_dtype(children, struct_dtype.clone(), len, validity)
112    }
113
114    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
115        let DType::Struct(struct_dtype, _nullability) = &array.dtype else {
116            vortex_bail!("Expected struct dtype, found {:?}", array.dtype)
117        };
118
119        // First child is validity (if present), followed by fields
120        let (validity, non_data_children) = if children.len() == struct_dtype.nfields() {
121            (array.validity.clone(), 0_usize)
122        } else if children.len() == struct_dtype.nfields() + 1 {
123            (Validity::Array(children[0].clone()), 1_usize)
124        } else {
125            vortex_bail!(
126                "Expected {} or {} children, found {}",
127                struct_dtype.nfields(),
128                struct_dtype.nfields() + 1,
129                children.len()
130            );
131        };
132
133        let fields: Arc<[ArrayRef]> = children.into_iter().skip(non_data_children).collect();
134        vortex_ensure!(
135            fields.len() == struct_dtype.nfields(),
136            "Expected {} field children, found {}",
137            struct_dtype.nfields(),
138            fields.len()
139        );
140
141        array.fields = fields;
142        array.validity = validity;
143        Ok(())
144    }
145
146    fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
147        let fields: Box<[_]> = array
148            .fields()
149            .iter()
150            .map(|field| field.execute(ctx))
151            .try_collect()?;
152        let validity_mask = array.validity_mask();
153
154        // SAFETY: we know that all field lengths match the struct array length, and the validity
155        Ok(unsafe { StructVector::new_unchecked(Arc::new(fields), validity_mask) }.into())
156    }
157
158    fn reduce_parent(
159        array: &Self::Array,
160        parent: &ArrayRef,
161        child_idx: usize,
162    ) -> VortexResult<Option<ArrayRef>> {
163        PARENT_RULES.evaluate(array, parent, child_idx)
164    }
165}
166
167#[derive(Debug)]
168pub struct StructVTable;
vortex_array/arrays/struct_/vtable/mod.rs

vortex_array/arrays/struct_/vtable/
mod.rs
