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

vortex_array/arrays/struct_/vtable/
mod.rs
