vortex_layout/layouts/struct_/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod reader;
5pub mod writer;
6
7use std::sync::Arc;
8
9use reader::StructReader;
10use vortex_array::EmptyMetadata;
11use vortex_array::dtype::DType;
12use vortex_array::dtype::Field;
13use vortex_array::dtype::FieldMask;
14use vortex_array::dtype::Nullability;
15use vortex_array::dtype::StructFields;
16use vortex_error::VortexExpect;
17use vortex_error::VortexResult;
18use vortex_error::vortex_bail;
19use vortex_error::vortex_ensure_eq;
20use vortex_error::vortex_err;
21use vortex_session::SessionExt;
22use vortex_session::VortexSession;
23use vortex_session::registry::CachedId;
24pub use writer::StructStrategy;
25
26use crate::Layout;
27use crate::LayoutChildType;
28use crate::LayoutDeserializeArgs;
29use crate::LayoutId;
30use crate::LayoutParts;
31use crate::LayoutReaderContext;
32use crate::LayoutReaderRef;
33use crate::LayoutRef;
34use crate::VTable;
35use crate::children::OwnedLayoutChildren;
36use crate::segments::SegmentSource;
37
38/// Struct layout vtable.
39#[derive(Clone, Debug)]
40pub struct Struct;
41
42/// Backwards-compatible name for the struct layout plugin.
43pub use Struct as StructLayoutEncoding;
44
45/// A layout decomposing a struct into one child per field and optional validity.
46pub type StructLayout = Layout<Struct>;
47
48impl VTable for Struct {
49    type LayoutData = ();
50    type Metadata = EmptyMetadata;
51
52    fn id(&self) -> LayoutId {
53        static ID: CachedId = CachedId::new("vortex.struct");
54        *ID
55    }
56
57    fn metadata(_layout: &Layout<Self>) -> Self::Metadata {
58        EmptyMetadata
59    }
60
61    fn deserialize(
62        &self,
63        args: &LayoutDeserializeArgs<'_>,
64        _metadata: &EmptyMetadata,
65    ) -> VortexResult<Self::LayoutData> {
66        Layout::<Struct>::validate_children(args.dtype, args.children.nchildren())?;
67
68        for idx in 0..args.children.nchildren() {
69            let child_row_count = args.children.child_row_count(idx);
70            vortex_ensure_eq!(
71                child_row_count,
72                args.row_count,
73                "Struct child {idx} row count does not match parent"
74            );
75        }
76        Ok(())
77    }
78
79    fn nslots(layout: &Layout<Self>) -> usize {
80        // Slot 0 is always reserved for validity (present only when nullable); the remaining
81        // slots are the struct fields, so field `i` always occupies slot `i + 1`.
82        layout.struct_fields().nfields() + 1
83    }
84
85    fn slot_to_child(layout: &Layout<Self>, slot: usize) -> Option<usize> {
86        let nullable = layout.dtype().is_nullable();
87        match slot {
88            0 => nullable.then_some(0),
89            _ => Some(slot - 1 + usize::from(nullable)),
90        }
91    }
92
93    fn child_dtype(layout: &Layout<Self>, slot: usize) -> VortexResult<DType> {
94        StructLayout::slot_dtype(layout.dtype(), slot)
95    }
96
97    fn child_type(layout: &Layout<Self>, slot: usize) -> LayoutChildType {
98        if slot == 0 {
99            LayoutChildType::Auxiliary("validity".into())
100        } else {
101            LayoutChildType::Field(
102                layout
103                    .struct_fields()
104                    .field_name(slot - 1)
105                    .vortex_expect("Field index out of bounds")
106                    .clone(),
107            )
108        }
109    }
110
111    fn new_reader(
112        layout: &Layout<Self>,
113        name: Arc<str>,
114        segment_source: Arc<dyn SegmentSource>,
115        session: &VortexSession,
116        ctx: &LayoutReaderContext,
117    ) -> VortexResult<LayoutReaderRef> {
118        Ok(Arc::new(StructReader::try_new(
119            layout.clone(),
120            name,
121            segment_source,
122            session.session(),
123            ctx.clone(),
124        )?))
125    }
126}
127
128impl Layout<Struct> {
129    /// Construct a struct layout from owned children.
130    pub fn new(row_count: u64, dtype: DType, children: Vec<LayoutRef>) -> Self {
131        Self::validate_children(&dtype, children.len()).vortex_expect("invalid struct children");
132        LayoutParts::new(
133            Struct,
134            dtype,
135            row_count,
136            Vec::new(),
137            OwnedLayoutChildren::layout_children(children),
138            (),
139        )
140        .into_typed()
141    }
142
143    /// Returns the struct fields.
144    pub fn struct_fields(&self) -> &StructFields {
145        self.dtype()
146            .as_struct_fields_opt()
147            .vortex_expect("Struct layout dtype must be a struct")
148    }
149
150    /// Invokes `per_child` for fields selected by `field_mask`.
151    pub fn matching_fields<F>(&self, field_mask: &[FieldMask], mut per_child: F) -> VortexResult<()>
152    where
153        F: FnMut(FieldMask, usize) -> VortexResult<()>,
154    {
155        if field_mask.iter().any(|mask| mask.matches_all()) {
156            for idx in 0..self.struct_fields().nfields() {
157                per_child(FieldMask::All, idx)?;
158            }
159            return Ok(());
160        }
161
162        for path in field_mask {
163            let Some(field) = path.starting_field()? else {
164                continue;
165            };
166            let Field::Name(field_name) = field else {
167                vortex_bail!("Expected field name, got {field:?}");
168            };
169            let idx = self
170                .struct_fields()
171                .find(field_name)
172                .ok_or_else(|| vortex_err!("Field not found: {field_name}"))?;
173            per_child(path.clone().step_into()?, idx)?;
174        }
175        Ok(())
176    }
177
178    fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> {
179        let fields = dtype
180            .as_struct_fields_opt()
181            .ok_or_else(|| vortex_err!("Expected struct dtype"))?;
182        let expected = fields.nfields() + usize::from(dtype.is_nullable());
183        vortex_ensure_eq!(nchildren, expected);
184        Ok(())
185    }
186
187    /// Returns the dtype of the child in logical `slot`: slot 0 is the non-nullable validity
188    /// bitmap, and slot `i + 1` is struct field `i`.
189    fn slot_dtype(dtype: &DType, slot: usize) -> VortexResult<DType> {
190        if slot == 0 {
191            Ok(DType::Bool(Nullability::NonNullable))
192        } else {
193            dtype
194                .as_struct_fields_opt()
195                .and_then(|fields| fields.field_by_index(slot - 1))
196                .ok_or_else(|| vortex_err!("Missing field {}", slot - 1))
197        }
198    }
199}
200
201#[cfg(test)]
202mod tests {
203    use vortex_array::dtype::FieldName;
204    use vortex_array::dtype::PType;
205    use vortex_session::registry::ReadContext;
206
207    use super::*;
208    use crate::layouts::flat::FlatLayout;
209    use crate::segments::SegmentId;
210
211    fn flat_child(dtype: DType, segment: u32) -> LayoutRef {
212        FlatLayout::new(3, dtype, SegmentId::from(segment), ReadContext::new([])).into_layout()
213    }
214
215    fn two_field_struct(nullability: Nullability) -> DType {
216        let i32 = DType::Primitive(PType::I32, Nullability::NonNullable);
217        DType::Struct(
218            StructFields::from_iter([("a", i32.clone()), ("b", i32)]),
219            nullability,
220        )
221    }
222
223    /// Field `i` must occupy logical slot `i + 1` regardless of whether the struct is nullable,
224    /// with slot 0 reserved for validity (absent when non-nullable).
225    #[test]
226    fn field_slots_are_stable_across_nullability() -> VortexResult<()> {
227        let i32 = DType::Primitive(PType::I32, Nullability::NonNullable);
228        let bool_ = DType::Bool(Nullability::NonNullable);
229
230        let non_null = StructLayout::new(
231            3,
232            two_field_struct(Nullability::NonNullable),
233            vec![flat_child(i32.clone(), 0), flat_child(i32.clone(), 1)],
234        );
235        // Validity + two fields, but validity is absent so only two serialized children.
236        assert_eq!(non_null.nslots(), 3);
237        assert_eq!(non_null.nchildren(), 2);
238        assert_eq!(non_null.slot_to_child(0), None);
239        assert_eq!(non_null.slot_to_child(1), Some(0));
240        assert_eq!(non_null.slot_to_child(2), Some(1));
241        // The validity slot is absent, so materializing it and its type both yield `None`.
242        assert!(non_null.slot(0)?.is_none());
243        assert_eq!(non_null.slot_type(0), None);
244
245        let nullable = StructLayout::new(
246            3,
247            two_field_struct(Nullability::Nullable),
248            vec![
249                flat_child(bool_, 0),
250                flat_child(i32.clone(), 1),
251                flat_child(i32, 2),
252            ],
253        );
254        assert_eq!(nullable.nslots(), 3);
255        assert_eq!(nullable.nchildren(), 3);
256        assert_eq!(nullable.slot_to_child(0), Some(0));
257        assert_eq!(nullable.slot_to_child(1), Some(1));
258        assert_eq!(nullable.slot_to_child(2), Some(2));
259        // The validity slot is present only for the nullable struct.
260        assert!(nullable.slot(0)?.is_some());
261        assert_eq!(
262            nullable.slot_type(0),
263            Some(LayoutChildType::Auxiliary("validity".into()))
264        );
265
266        // Field "a" is slot 1 in both layouts, even though its serialized index differs.
267        for layout in [&non_null, &nullable] {
268            assert_eq!(
269                layout.slot_type(1),
270                Some(LayoutChildType::Field(FieldName::from("a")))
271            );
272        }
273
274        Ok(())
275    }
276}
vortex_layout/layouts/struct_/mod.rs

vortex_layout/layouts/struct_/
mod.rs
