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::{ArrayContext, DeserializeMetadata, EmptyMetadata};
11use vortex_dtype::{DType, Field, FieldMask, StructFields};
12use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err, vortex_panic};
13
14use crate::children::{LayoutChildren, OwnedLayoutChildren};
15use crate::segments::{SegmentId, SegmentSource};
16use crate::{
17    LayoutChildType, LayoutEncodingRef, LayoutId, LayoutReaderRef, LayoutRef, VTable, vtable,
18};
19
20vtable!(Struct);
21
22impl VTable for StructVTable {
23    type Layout = StructLayout;
24    type Encoding = StructLayoutEncoding;
25    type Metadata = EmptyMetadata;
26
27    fn id(_encoding: &Self::Encoding) -> LayoutId {
28        LayoutId::new_ref("vortex.struct")
29    }
30
31    fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef {
32        LayoutEncodingRef::new_ref(StructLayoutEncoding.as_ref())
33    }
34
35    fn row_count(layout: &Self::Layout) -> u64 {
36        layout.row_count
37    }
38
39    fn dtype(layout: &Self::Layout) -> &DType {
40        &layout.dtype
41    }
42
43    fn metadata(_layout: &Self::Layout) -> Self::Metadata {
44        EmptyMetadata
45    }
46
47    fn segment_ids(_layout: &Self::Layout) -> Vec<SegmentId> {
48        vec![]
49    }
50
51    fn nchildren(layout: &Self::Layout) -> usize {
52        layout.struct_fields().nfields()
53    }
54
55    fn child(layout: &Self::Layout, idx: usize) -> VortexResult<LayoutRef> {
56        layout.children.child(
57            idx,
58            &layout
59                .struct_fields()
60                .field_by_index(idx)
61                .ok_or_else(|| vortex_err!("Missing field {idx}"))?,
62        )
63    }
64
65    fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType {
66        LayoutChildType::Field(
67            layout
68                .struct_fields()
69                .field_name(idx)
70                .vortex_expect("Field index out of bounds")
71                .clone(),
72        )
73    }
74
75    fn new_reader(
76        layout: &Self::Layout,
77        name: Arc<str>,
78        segment_source: Arc<dyn SegmentSource>,
79    ) -> VortexResult<LayoutReaderRef> {
80        Ok(Arc::new(StructReader::try_new(
81            layout.clone(),
82            name,
83            segment_source,
84        )?))
85    }
86
87    fn build(
88        _encoding: &Self::Encoding,
89        dtype: &DType,
90        row_count: u64,
91        _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
92        _segment_ids: Vec<SegmentId>,
93        children: &dyn LayoutChildren,
94        _ctx: ArrayContext,
95    ) -> VortexResult<Self::Layout> {
96        let struct_dt = dtype
97            .as_struct_opt()
98            .ok_or_else(|| vortex_err!("Expected struct dtype"))?;
99        if children.nchildren() != struct_dt.nfields() {
100            vortex_bail!(
101                "Struct layout has {} children, but dtype has {} fields",
102                children.nchildren(),
103                struct_dt.nfields()
104            );
105        }
106        Ok(StructLayout {
107            row_count,
108            dtype: dtype.clone(),
109            children: children.to_arc(),
110        })
111    }
112}
113
114#[derive(Debug)]
115pub struct StructLayoutEncoding;
116
117#[derive(Clone, Debug)]
118pub struct StructLayout {
119    row_count: u64,
120    dtype: DType,
121    children: Arc<dyn LayoutChildren>,
122}
123
124impl StructLayout {
125    pub fn new(row_count: u64, dtype: DType, children: Vec<LayoutRef>) -> Self {
126        Self {
127            row_count,
128            dtype,
129            children: OwnedLayoutChildren::layout_children(children),
130        }
131    }
132
133    pub fn struct_fields(&self) -> &StructFields {
134        let DType::Struct(dtype, _) = self.dtype() else {
135            vortex_panic!("Mismatched dtype {} for struct layout", self.dtype());
136        };
137        dtype
138    }
139
140    pub fn matching_fields<F>(&self, field_mask: &[FieldMask], mut per_child: F) -> VortexResult<()>
141    where
142        F: FnMut(FieldMask, usize) -> VortexResult<()>,
143    {
144        // If the field mask contains an `All` fields, then enumerate all fields.
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        // Enumerate each field in the mask
153        for path in field_mask {
154            let Some(field) = path.starting_field()? else {
155                // skip fields not in mask
156                continue;
157            };
158            let Field::Name(field_name) = field else {
159                vortex_bail!("Expected field name, got {field:?}");
160            };
161            let idx = self
162                .struct_fields()
163                .find(field_name)
164                .ok_or_else(|| vortex_err!("Field not found: {field_name}"))?;
165
166            per_child(path.clone().step_into()?, idx)?;
167        }
168
169        Ok(())
170    }
171}
vortex_layout/layouts/struct_/mod.rs

vortex_layout/layouts/struct_/
mod.rs
