vortex_layout/layouts/struct_/
mod.rs

1mod eval_expr;
2mod reader;
3pub mod writer;
4
5use std::collections::BTreeSet;
6use std::sync::Arc;
7
8use reader::StructReader;
9use vortex_array::ArrayContext;
10use vortex_dtype::{DType, Field, FieldMask};
11use vortex_error::{VortexResult, vortex_bail};
12
13use crate::data::Layout;
14use crate::reader::{LayoutReader, LayoutReaderExt};
15use crate::segments::{AsyncSegmentReader, RequiredSegmentKind, SegmentCollector};
16use crate::vtable::LayoutVTable;
17use crate::{LayoutId, STRUCT_LAYOUT_ID};
18
19#[derive(Debug)]
20pub struct StructLayout;
21
22impl LayoutVTable for StructLayout {
23    fn id(&self) -> LayoutId {
24        STRUCT_LAYOUT_ID
25    }
26
27    fn reader(
28        &self,
29        layout: Layout,
30        ctx: ArrayContext,
31        segment_reader: Arc<dyn AsyncSegmentReader>,
32    ) -> VortexResult<Arc<dyn LayoutReader>> {
33        Ok(StructReader::try_new(layout, ctx, segment_reader)?.into_arc())
34    }
35
36    fn register_splits(
37        &self,
38        layout: &Layout,
39        field_mask: &[FieldMask],
40        row_offset: u64,
41        splits: &mut BTreeSet<u64>,
42    ) -> VortexResult<()> {
43        for_all_matching_children(layout, field_mask, |mask, child| {
44            child.register_splits(&[mask], row_offset, splits)
45        })?;
46        Ok(())
47    }
48
49    fn required_segments(
50        &self,
51        layout: &Layout,
52        row_offset: u64,
53        filter_field_mask: &[FieldMask],
54        projection_field_mask: &[FieldMask],
55        segments: &mut SegmentCollector,
56    ) -> VortexResult<()> {
57        for_all_matching_children(layout, filter_field_mask, |field_mask, child| {
58            child.required_segments(
59                row_offset,
60                &[field_mask],
61                &[],
62                &mut segments.with_priority_hint(RequiredSegmentKind::Filter),
63            )
64        })?;
65        for_all_matching_children(layout, projection_field_mask, |field_mask, child| {
66            child.required_segments(
67                row_offset,
68                &[],
69                &[field_mask],
70                &mut segments.with_priority_hint(RequiredSegmentKind::Projection),
71            )
72        })?;
73        Ok(())
74    }
75}
76
77fn for_all_matching_children<F>(
78    layout: &Layout,
79    field_mask: &[FieldMask],
80    mut per_child: F,
81) -> VortexResult<()>
82where
83    F: FnMut(FieldMask, Layout) -> VortexResult<()>,
84{
85    let DType::Struct(dtype, _) = layout.dtype() else {
86        vortex_bail!("Mismatched dtype {} for struct layout", layout.dtype());
87    };
88
89    // If the field mask contains an `All` fields, then enumerate all fields.
90    if field_mask.iter().any(|mask| mask.matches_all()) {
91        for (idx, field_dtype) in dtype.fields().enumerate() {
92            let child = layout.child(idx, field_dtype, dtype.field_name(idx)?)?;
93            per_child(FieldMask::All, child)?;
94        }
95        return Ok(());
96    }
97
98    // Enumerate each field in the mask
99    for path in field_mask {
100        let Some(field) = path.starting_field()? else {
101            // skip fields not in mask
102            continue;
103        };
104        let Field::Name(field_name) = field else {
105            vortex_bail!("Expected field name, got {:?}", field);
106        };
107
108        let idx = dtype.find(field_name)?;
109        let child = layout.child(idx, dtype.field_by_index(idx)?, field_name)?;
110        per_child(path.clone().step_into()?, child)?;
111    }
112
113    Ok(())
114}
vortex_layout/layouts/struct_/mod.rs

vortex_layout/layouts/struct_/
mod.rs
