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;
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!(
71                child_row_count == args.row_count,
72                "Struct child {idx} row count does not match parent"
73            );
74        }
75        Ok(())
76    }
77
78    fn nslots(layout: &Layout<Self>) -> usize {
79        // Slot 0 is always reserved for validity (present only when nullable); the remaining
80        // slots are the struct fields, so field `i` always occupies slot `i + 1`.
81        layout.struct_fields().nfields() + 1
82    }
83
84    fn slot_to_child(layout: &Layout<Self>, slot: usize) -> Option<usize> {
85        let nullable = layout.dtype().is_nullable();
86        match slot {
87            0 => nullable.then_some(0),
88            _ => Some(slot - 1 + usize::from(nullable)),
89        }
90    }
91
92    fn child_dtype(layout: &Layout<Self>, slot: usize) -> VortexResult<DType> {
93        StructLayout::slot_dtype(layout.dtype(), slot)
94    }
95
96    fn child_type(layout: &Layout<Self>, slot: usize) -> LayoutChildType {
97        if slot == 0 {
98            LayoutChildType::Auxiliary("validity".into())
99        } else {
100            LayoutChildType::Field(
101                layout
102                    .struct_fields()
103                    .field_name(slot - 1)
104                    .vortex_expect("Field index out of bounds")
105                    .clone(),
106            )
107        }
108    }
109
110    fn new_reader(
111        layout: &Layout<Self>,
112        name: Arc<str>,
113        segment_source: Arc<dyn SegmentSource>,
114        session: &VortexSession,
115        ctx: &LayoutReaderContext,
116    ) -> VortexResult<LayoutReaderRef> {
117        Ok(Arc::new(StructReader::try_new(
118            layout.clone(),
119            name,
120            segment_source,
121            session.session(),
122            ctx.clone(),
123        )?))
124    }
125}
126
127impl Layout<Struct> {
128    /// Construct a struct layout from owned children.
129    pub fn new(row_count: u64, dtype: DType, children: Vec<LayoutRef>) -> Self {
130        Self::validate_children(&dtype, children.len()).vortex_expect("invalid struct children");
131        LayoutParts::new(
132            Struct,
133            dtype,
134            row_count,
135            Vec::new(),
136            OwnedLayoutChildren::layout_children(children),
137            (),
138        )
139        .into_typed()
140    }
141
142    /// Returns the struct fields.
143    pub fn struct_fields(&self) -> &StructFields {
144        self.dtype()
145            .as_struct_fields_opt()
146            .vortex_expect("Struct layout dtype must be a struct")
147    }
148
149    /// Invokes `per_child` for fields selected by `field_mask`.
150    pub fn matching_fields<F>(&self, field_mask: &[FieldMask], mut per_child: F) -> VortexResult<()>
151    where
152        F: FnMut(FieldMask, usize) -> VortexResult<()>,
153    {
154        if field_mask.iter().any(|mask| mask.matches_all()) {
155            for idx in 0..self.struct_fields().nfields() {
156                per_child(FieldMask::All, idx)?;
157            }
158            return Ok(());
159        }
160
161        for path in field_mask {
162            let Some(field) = path.starting_field()? else {
163                continue;
164            };
165            let Field::Name(field_name) = field else {
166                vortex_bail!("Expected field name, got {field:?}");
167            };
168            let idx = self
169                .struct_fields()
170                .find(field_name)
171                .ok_or_else(|| vortex_err!("Field not found: {field_name}"))?;
172            per_child(path.clone().step_into()?, idx)?;
173        }
174        Ok(())
175    }
176
177    fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> {
178        let fields = dtype
179            .as_struct_fields_opt()
180            .ok_or_else(|| vortex_err!("Expected struct dtype"))?;
181        let expected = fields.nfields() + usize::from(dtype.is_nullable());
182        vortex_ensure!(
183            nchildren == expected,
184            "Struct layout has {nchildren} children, expected {expected}"
185        );
186        Ok(())
187    }
188
189    /// Returns the dtype of the child in logical `slot`: slot 0 is the non-nullable validity
190    /// bitmap, and slot `i + 1` is struct field `i`.
191    fn slot_dtype(dtype: &DType, slot: usize) -> VortexResult<DType> {
192        if slot == 0 {
193            Ok(DType::Bool(Nullability::NonNullable))
194        } else {
195            dtype
196                .as_struct_fields_opt()
197                .and_then(|fields| fields.field_by_index(slot - 1))
198                .ok_or_else(|| vortex_err!("Missing field {}", slot - 1))
199        }
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use vortex_array::dtype::FieldName;
206    use vortex_array::dtype::PType;
207    use vortex_session::registry::ReadContext;
208
209    use super::*;
210    use crate::layouts::flat::FlatLayout;
211    use crate::segments::SegmentId;
212
213    fn flat_child(dtype: DType, segment: u32) -> LayoutRef {
214        FlatLayout::new(3, dtype, SegmentId::from(segment), ReadContext::new([])).into_layout()
215    }
216
217    fn two_field_struct(nullability: Nullability) -> DType {
218        let i32 = DType::Primitive(PType::I32, Nullability::NonNullable);
219        DType::Struct(
220            StructFields::from_iter([("a", i32.clone()), ("b", i32)]),
221            nullability,
222        )
223    }
224
225    /// Field `i` must occupy logical slot `i + 1` regardless of whether the struct is nullable,
226    /// with slot 0 reserved for validity (absent when non-nullable).
227    #[test]
228    fn field_slots_are_stable_across_nullability() -> VortexResult<()> {
229        let i32 = DType::Primitive(PType::I32, Nullability::NonNullable);
230        let bool_ = DType::Bool(Nullability::NonNullable);
231
232        let non_null = StructLayout::new(
233            3,
234            two_field_struct(Nullability::NonNullable),
235            vec![flat_child(i32.clone(), 0), flat_child(i32.clone(), 1)],
236        );
237        // Validity + two fields, but validity is absent so only two serialized children.
238        assert_eq!(non_null.nslots(), 3);
239        assert_eq!(non_null.nchildren(), 2);
240        assert_eq!(non_null.slot_to_child(0), None);
241        assert_eq!(non_null.slot_to_child(1), Some(0));
242        assert_eq!(non_null.slot_to_child(2), Some(1));
243        // The validity slot is absent, so materializing it and its type both yield `None`.
244        assert!(non_null.slot(0)?.is_none());
245        assert_eq!(non_null.slot_type(0), None);
246
247        let nullable = StructLayout::new(
248            3,
249            two_field_struct(Nullability::Nullable),
250            vec![
251                flat_child(bool_, 0),
252                flat_child(i32.clone(), 1),
253                flat_child(i32, 2),
254            ],
255        );
256        assert_eq!(nullable.nslots(), 3);
257        assert_eq!(nullable.nchildren(), 3);
258        assert_eq!(nullable.slot_to_child(0), Some(0));
259        assert_eq!(nullable.slot_to_child(1), Some(1));
260        assert_eq!(nullable.slot_to_child(2), Some(2));
261        // The validity slot is present only for the nullable struct.
262        assert!(nullable.slot(0)?.is_some());
263        assert_eq!(
264            nullable.slot_type(0),
265            Some(LayoutChildType::Auxiliary("validity".into()))
266        );
267
268        // Field "a" is slot 1 in both layouts, even though its serialized index differs.
269        for layout in [&non_null, &nullable] {
270            assert_eq!(
271                layout.slot_type(1),
272                Some(LayoutChildType::Field(FieldName::from("a")))
273            );
274        }
275
276        Ok(())
277    }
278}
vortex_layout/layouts/struct_/mod.rs

vortex_layout/layouts/struct_/
mod.rs
