vortex_layout/layouts/flat/
mod.rs

1mod reader;
2pub mod writer;
3
4use std::collections::BTreeSet;
5use std::sync::Arc;
6
7use vortex_array::{ArrayContext, DeserializeMetadata, EmptyMetadata};
8use vortex_dtype::{DType, FieldMask};
9use vortex_error::{VortexResult, vortex_bail, vortex_panic};
10
11use crate::children::LayoutChildren;
12use crate::layouts::flat::reader::FlatReader;
13use crate::segments::{SegmentId, SegmentSource};
14use crate::{
15    LayoutChildType, LayoutEncodingRef, LayoutId, LayoutReaderRef, LayoutRef, VTable, vtable,
16};
17
18vtable!(Flat);
19
20impl VTable for FlatVTable {
21    type Layout = FlatLayout;
22    type Encoding = FlatLayoutEncoding;
23    type Metadata = EmptyMetadata;
24
25    fn id(_encoding: &Self::Encoding) -> LayoutId {
26        LayoutId::new_ref("vortex.flat")
27    }
28
29    fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef {
30        LayoutEncodingRef::new_ref(FlatLayoutEncoding.as_ref())
31    }
32
33    fn row_count(layout: &Self::Layout) -> u64 {
34        layout.row_count
35    }
36
37    fn dtype(layout: &Self::Layout) -> &DType {
38        &layout.dtype
39    }
40
41    fn metadata(_layout: &Self::Layout) -> Self::Metadata {
42        EmptyMetadata
43    }
44
45    fn segment_ids(layout: &Self::Layout) -> Vec<SegmentId> {
46        vec![layout.segment_id]
47    }
48
49    fn nchildren(_layout: &Self::Layout) -> usize {
50        0
51    }
52
53    fn child(_layout: &Self::Layout, _idx: usize) -> VortexResult<LayoutRef> {
54        vortex_bail!("Flat layout has no children");
55    }
56
57    fn child_type(_layout: &Self::Layout, _idx: usize) -> LayoutChildType {
58        vortex_panic!("Flat layout has no children");
59    }
60
61    fn register_splits(
62        layout: &Self::Layout,
63        field_mask: &[FieldMask],
64        row_offset: u64,
65        splits: &mut BTreeSet<u64>,
66    ) -> VortexResult<()> {
67        for path in field_mask {
68            if path.matches_root() {
69                splits.insert(row_offset + layout.row_count());
70                break;
71            }
72        }
73        Ok(())
74    }
75
76    fn new_reader(
77        layout: &Self::Layout,
78        name: &Arc<str>,
79        segment_source: &Arc<dyn SegmentSource>,
80        ctx: &ArrayContext,
81    ) -> VortexResult<LayoutReaderRef> {
82        Ok(Arc::new(FlatReader::new(
83            layout.clone(),
84            name.clone(),
85            segment_source.clone(),
86            ctx.clone(),
87        )))
88    }
89
90    fn build(
91        _encoding: &Self::Encoding,
92        dtype: &DType,
93        row_count: u64,
94        _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
95        segment_ids: Vec<SegmentId>,
96        _children: &dyn LayoutChildren,
97    ) -> VortexResult<Self::Layout> {
98        if segment_ids.len() != 1 {
99            vortex_bail!("Flat layout must have exactly one segment ID");
100        }
101        Ok(FlatLayout {
102            row_count,
103            dtype: dtype.clone(),
104            segment_id: segment_ids[0],
105        })
106    }
107}
108
109#[derive(Debug)]
110pub struct FlatLayoutEncoding;
111
112#[derive(Clone, Debug)]
113pub struct FlatLayout {
114    row_count: u64,
115    dtype: DType,
116    segment_id: SegmentId,
117}
118
119impl FlatLayout {
120    pub fn new(row_count: u64, dtype: DType, segment_id: SegmentId) -> Self {
121        Self {
122            row_count,
123            dtype,
124            segment_id,
125        }
126    }
127
128    pub fn segment_id(&self) -> SegmentId {
129        self.segment_id
130    }
131}