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_buffer::BufferHandle;
8use vortex_dtype::DType;
9use vortex_error::VortexExpect;
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_vector::Vector;
13use vortex_vector::struct_::StructVector;
14
15use crate::EmptyMetadata;
16use crate::arrays::struct_::StructArray;
17use crate::execution::ExecutionCtx;
18use crate::serde::ArrayChildren;
19use crate::validity::Validity;
20use crate::vtable;
21use crate::vtable::ArrayVTableExt;
22use crate::vtable::NotSupported;
23use crate::vtable::VTable;
24use crate::vtable::ValidityVTableFromValidityHelper;
25
26mod array;
27mod canonical;
28mod operations;
29mod validity;
30mod visitor;
31
32use crate::vtable::ArrayId;
33use crate::vtable::ArrayVTable;
34
35vtable!(Struct);
36
37impl VTable for StructVTable {
38    type Array = StructArray;
39
40    type Metadata = EmptyMetadata;
41
42    type ArrayVTable = Self;
43    type CanonicalVTable = Self;
44    type OperationsVTable = Self;
45    type ValidityVTable = ValidityVTableFromValidityHelper;
46    type VisitorVTable = Self;
47    type ComputeVTable = NotSupported;
48    type EncodeVTable = NotSupported;
49
50    fn id(&self) -> ArrayId {
51        ArrayId::new_ref("vortex.struct")
52    }
53
54    fn encoding(_array: &Self::Array) -> ArrayVTable {
55        StructVTable.as_vtable()
56    }
57
58    fn metadata(_array: &StructArray) -> VortexResult<Self::Metadata> {
59        Ok(EmptyMetadata)
60    }
61
62    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
63        Ok(Some(vec![]))
64    }
65
66    fn deserialize(_buffer: &[u8]) -> VortexResult<Self::Metadata> {
67        Ok(EmptyMetadata)
68    }
69
70    fn build(
71        &self,
72        dtype: &DType,
73        len: usize,
74        _metadata: &Self::Metadata,
75        _buffers: &[BufferHandle],
76        children: &dyn ArrayChildren,
77    ) -> VortexResult<StructArray> {
78        let DType::Struct(struct_dtype, nullability) = dtype else {
79            vortex_bail!("Expected struct dtype, found {:?}", dtype)
80        };
81
82        let (validity, non_data_children) = if children.len() == struct_dtype.nfields() {
83            (Validity::from(*nullability), 0_usize)
84        } else if children.len() == struct_dtype.nfields() + 1 {
85            // Validity is the first child if it exists.
86            let validity = children.get(0, &Validity::DTYPE, len)?;
87            (Validity::Array(validity), 1_usize)
88        } else {
89            vortex_bail!(
90                "Expected {} or {} children, found {}",
91                struct_dtype.nfields(),
92                struct_dtype.nfields() + 1,
93                children.len()
94            );
95        };
96
97        let children: Vec<_> = (0..struct_dtype.nfields())
98            .map(|i| {
99                let child_dtype = struct_dtype
100                    .field_by_index(i)
101                    .vortex_expect("no out of bounds");
102                children.get(non_data_children + i, &child_dtype, len)
103            })
104            .try_collect()?;
105
106        StructArray::try_new_with_dtype(children, struct_dtype.clone(), len, validity)
107    }
108
109    fn batch_execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
110        let fields: Box<[_]> = array
111            .fields()
112            .iter()
113            .map(|field| field.batch_execute(ctx))
114            .try_collect()?;
115        // SAFETY: we know that all field lengths match the struct array length, and the validity
116        Ok(unsafe { StructVector::new_unchecked(Arc::new(fields), array.validity_mask()) }.into())
117    }
118}
119
120#[derive(Debug)]
121pub struct StructVTable;