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 child_dtype(layout: &Layout<Self>, index: usize) -> VortexResult<DType> {
79        StructLayout::child_dtype(layout.dtype(), index)
80    }
81
82    fn child_type(layout: &Layout<Self>, idx: usize) -> LayoutChildType {
83        let schema_index = if layout.dtype().is_nullable() {
84            idx.saturating_sub(1)
85        } else {
86            idx
87        };
88        if idx == 0 && layout.dtype().is_nullable() {
89            LayoutChildType::Auxiliary("validity".into())
90        } else {
91            LayoutChildType::Field(
92                layout
93                    .struct_fields()
94                    .field_name(schema_index)
95                    .vortex_expect("Field index out of bounds")
96                    .clone(),
97            )
98        }
99    }
100
101    fn new_reader(
102        layout: &Layout<Self>,
103        name: Arc<str>,
104        segment_source: Arc<dyn SegmentSource>,
105        session: &VortexSession,
106        ctx: &LayoutReaderContext,
107    ) -> VortexResult<LayoutReaderRef> {
108        Ok(Arc::new(StructReader::try_new(
109            layout.clone(),
110            name,
111            segment_source,
112            session.session(),
113            ctx.clone(),
114        )?))
115    }
116}
117
118impl Layout<Struct> {
119    /// Construct a struct layout from owned children.
120    pub fn new(row_count: u64, dtype: DType, children: Vec<LayoutRef>) -> Self {
121        Self::validate_children(&dtype, children.len()).vortex_expect("invalid struct children");
122        LayoutParts::new(
123            Struct,
124            dtype,
125            row_count,
126            Vec::new(),
127            OwnedLayoutChildren::layout_children(children),
128            (),
129        )
130        .into_typed()
131    }
132
133    /// Returns the struct fields.
134    pub fn struct_fields(&self) -> &StructFields {
135        self.dtype()
136            .as_struct_fields_opt()
137            .vortex_expect("Struct layout dtype must be a struct")
138    }
139
140    /// Invokes `per_child` for fields selected by `field_mask`.
141    pub fn matching_fields<F>(&self, field_mask: &[FieldMask], mut per_child: F) -> VortexResult<()>
142    where
143        F: FnMut(FieldMask, usize) -> VortexResult<()>,
144    {
145        if field_mask.iter().any(|mask| mask.matches_all()) {
146            for idx in 0..self.struct_fields().nfields() {
147                per_child(FieldMask::All, idx)?;
148            }
149            return Ok(());
150        }
151
152        for path in field_mask {
153            let Some(field) = path.starting_field()? else {
154                continue;
155            };
156            let Field::Name(field_name) = field else {
157                vortex_bail!("Expected field name, got {field:?}");
158            };
159            let idx = self
160                .struct_fields()
161                .find(field_name)
162                .ok_or_else(|| vortex_err!("Field not found: {field_name}"))?;
163            per_child(path.clone().step_into()?, idx)?;
164        }
165        Ok(())
166    }
167
168    fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> {
169        let fields = dtype
170            .as_struct_fields_opt()
171            .ok_or_else(|| vortex_err!("Expected struct dtype"))?;
172        let expected = fields.nfields() + usize::from(dtype.is_nullable());
173        vortex_ensure!(
174            nchildren == expected,
175            "Struct layout has {nchildren} children, expected {expected}"
176        );
177        Ok(())
178    }
179
180    fn child_dtype(dtype: &DType, index: usize) -> VortexResult<DType> {
181        let schema_index = if dtype.is_nullable() {
182            index.saturating_sub(1)
183        } else {
184            index
185        };
186        if index == 0 && dtype.is_nullable() {
187            Ok(DType::Bool(Nullability::NonNullable))
188        } else {
189            dtype
190                .as_struct_fields_opt()
191                .and_then(|fields| fields.field_by_index(schema_index))
192                .ok_or_else(|| vortex_err!("Missing field {schema_index}"))
193        }
194    }
195}
vortex_layout/layouts/struct_/mod.rs

vortex_layout/layouts/struct_/
mod.rs
