vortex_layout/layouts/flat/
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::env;
8use std::sync::Arc;
9use std::sync::LazyLock;
10
11use vortex_array::DeserializeMetadata;
12use vortex_array::ProstMetadata;
13use vortex_array::dtype::DType;
14use vortex_buffer::ByteBuffer;
15use vortex_error::VortexResult;
16use vortex_error::vortex_bail;
17use vortex_error::vortex_panic;
18use vortex_session::VortexSession;
19use vortex_session::registry::CachedId;
20use vortex_session::registry::ReadContext;
21
22use crate::LayoutBuildContext;
23use crate::LayoutChildType;
24use crate::LayoutEncodingRef;
25use crate::LayoutId;
26use crate::LayoutReaderRef;
27use crate::LayoutRef;
28use crate::VTable;
29use crate::children::LayoutChildren;
30use crate::layouts::flat::reader::FlatReader;
31use crate::segments::SegmentId;
32use crate::segments::SegmentSource;
33use crate::vtable;
34
35/// Check if inline array node is enabled.
36pub(super) fn flat_layout_inline_array_node() -> bool {
37    static FLAT_LAYOUT_INLINE_ARRAY_NODE: LazyLock<bool> =
38        LazyLock::new(|| env::var("FLAT_LAYOUT_INLINE_ARRAY_NODE").is_ok_and(|v| v == "1"));
39    *FLAT_LAYOUT_INLINE_ARRAY_NODE
40}
41
42vtable!(Flat);
43
44impl VTable for Flat {
45    type Layout = FlatLayout;
46    type Encoding = FlatLayoutEncoding;
47    type Metadata = ProstMetadata<FlatLayoutMetadata>;
48
49    fn id(_encoding: &Self::Encoding) -> LayoutId {
50        static ID: CachedId = CachedId::new("vortex.flat");
51        *ID
52    }
53
54    fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef {
55        LayoutEncodingRef::new_ref(FlatLayoutEncoding.as_ref())
56    }
57
58    fn row_count(layout: &Self::Layout) -> u64 {
59        layout.row_count
60    }
61
62    fn dtype(layout: &Self::Layout) -> &DType {
63        &layout.dtype
64    }
65
66    fn metadata(layout: &Self::Layout) -> Self::Metadata {
67        ProstMetadata(FlatLayoutMetadata {
68            array_encoding_tree: layout.array_tree.as_ref().map(|bytes| bytes.to_vec()),
69        })
70    }
71
72    fn segment_ids(layout: &Self::Layout) -> Vec<SegmentId> {
73        vec![layout.segment_id]
74    }
75
76    fn nchildren(_layout: &Self::Layout) -> usize {
77        0
78    }
79
80    fn child(_layout: &Self::Layout, _idx: usize) -> VortexResult<LayoutRef> {
81        vortex_bail!("Flat layout has no children");
82    }
83
84    fn child_type(_layout: &Self::Layout, _idx: usize) -> LayoutChildType {
85        vortex_panic!("Flat layout has no children");
86    }
87
88    fn new_reader(
89        layout: &Self::Layout,
90        name: Arc<str>,
91        segment_source: Arc<dyn SegmentSource>,
92        session: &VortexSession,
93        _ctx: &crate::LayoutReaderContext,
94    ) -> VortexResult<LayoutReaderRef> {
95        Ok(Arc::new(FlatReader::new(
96            layout.clone(),
97            name,
98            segment_source,
99            session.clone(),
100        )))
101    }
102
103    fn build(
104        _encoding: &Self::Encoding,
105        dtype: &DType,
106        row_count: u64,
107        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
108        segment_ids: Vec<SegmentId>,
109        _children: &dyn LayoutChildren,
110        build_ctx: &LayoutBuildContext<'_>,
111    ) -> VortexResult<Self::Layout> {
112        if segment_ids.len() != 1 {
113            vortex_bail!("Flat layout must have exactly one segment ID");
114        }
115        Ok(FlatLayout::new_with_metadata(
116            row_count,
117            dtype.clone(),
118            segment_ids[0],
119            build_ctx.array_read_ctx.clone(),
120            metadata
121                .array_encoding_tree
122                .as_ref()
123                .map(|v| ByteBuffer::from(v.clone())),
124        ))
125    }
126
127    fn with_children(_layout: &mut Self::Layout, children: Vec<LayoutRef>) -> VortexResult<()> {
128        if !children.is_empty() {
129            vortex_bail!("Flat layout has no children, got {}", children.len());
130        }
131        Ok(())
132    }
133}
134
135#[derive(Debug)]
136pub struct FlatLayoutEncoding;
137
138/// The terminal node of a layout tree. Stores a single chunk of array data as one serialized
139/// segment on disk.
140#[derive(Clone, Debug)]
141pub struct FlatLayout {
142    row_count: u64,
143    dtype: DType,
144    segment_id: SegmentId,
145    ctx: ReadContext,
146    array_tree: Option<ByteBuffer>,
147}
148
149impl FlatLayout {
150    pub fn new(row_count: u64, dtype: DType, segment_id: SegmentId, ctx: ReadContext) -> Self {
151        Self {
152            row_count,
153            dtype,
154            segment_id,
155            ctx,
156            array_tree: None,
157        }
158    }
159
160    pub fn new_with_metadata(
161        row_count: u64,
162        dtype: DType,
163        segment_id: SegmentId,
164        ctx: ReadContext,
165        metadata: Option<ByteBuffer>,
166    ) -> Self {
167        Self {
168            row_count,
169            dtype,
170            segment_id,
171            ctx,
172            array_tree: metadata,
173        }
174    }
175
176    #[inline]
177    pub fn segment_id(&self) -> SegmentId {
178        self.segment_id
179    }
180
181    #[inline]
182    pub fn array_ctx(&self) -> &ReadContext {
183        &self.ctx
184    }
185
186    #[inline]
187    pub fn array_tree(&self) -> Option<&ByteBuffer> {
188        self.array_tree.as_ref()
189    }
190}
191
192#[derive(prost::Message)]
193pub struct FlatLayoutMetadata {
194    // We can optionally store the array encoding tree here to avoid needing to fetch the segment
195    // to plan array deserialization.
196    // This will be a `ArrayNode`.
197    #[prost(optional, bytes, tag = "1")]
198    pub array_encoding_tree: Option<Vec<u8>>,
199}
vortex_layout/layouts/flat/mod.rs

vortex_layout/layouts/flat/
mod.rs
