vortex_layout/layouts/struct_/
mod.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
mod eval_expr;
mod eval_stats;
mod reader;
pub mod writer;

use std::collections::BTreeSet;
use std::sync::Arc;

use reader::StructReader;
use vortex_array::ContextRef;
use vortex_dtype::{DType, FieldMask};
use vortex_error::{vortex_bail, vortex_err, VortexResult};

use crate::data::Layout;
use crate::reader::{LayoutReader, LayoutReaderExt};
use crate::segments::AsyncSegmentReader;
use crate::vtable::LayoutVTable;
use crate::{LayoutId, COLUMNAR_LAYOUT_ID};

#[derive(Debug)]
pub struct StructLayout;

impl LayoutVTable for StructLayout {
    fn id(&self) -> LayoutId {
        COLUMNAR_LAYOUT_ID
    }

    fn reader(
        &self,
        layout: Layout,
        ctx: ContextRef,
        segments: Arc<dyn AsyncSegmentReader>,
    ) -> VortexResult<Arc<dyn LayoutReader>> {
        Ok(StructReader::try_new(layout, segments, ctx)?.into_arc())
    }

    fn register_splits(
        &self,
        layout: &Layout,
        field_mask: &[FieldMask],
        row_offset: u64,
        splits: &mut BTreeSet<u64>,
    ) -> VortexResult<()> {
        let DType::Struct(dtype, _) = layout.dtype() else {
            vortex_bail!("Mismatched dtype {} for struct layout", layout.dtype());
        };

        // If the field mask contains an `All` fields, then register splits for all fields.
        if field_mask.iter().any(|mask| mask.matches_all()) {
            for (idx, field_dtype) in dtype.dtypes().enumerate() {
                let child = layout.child(idx, field_dtype)?;
                child.register_splits(&[FieldMask::All], row_offset, splits)?;
            }
            return Ok(());
        }

        // Register the splits for each field in the mask
        for path in field_mask {
            let Some(field) = path.starting_field()? else {
                // skip fields not in mask
                continue;
            };

            let idx = dtype
                .find(field)
                .ok_or_else(|| vortex_err!("Field not found: {:?}", path))?;

            let child = layout.child(idx, dtype.field_dtype(idx)?)?;
            child.register_splits(&[path.clone().step_into()?], row_offset, splits)?;
        }

        Ok(())
    }
}