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::ProstMetadata;
12use vortex_array::dtype::DType;
13use vortex_buffer::ByteBuffer;
14use vortex_error::VortexResult;
15use vortex_error::vortex_bail;
16use vortex_error::vortex_panic;
17use vortex_session::VortexSession;
18use vortex_session::registry::CachedId;
19use vortex_session::registry::ReadContext;
20
21use crate::Layout;
22use crate::LayoutChildType;
23use crate::LayoutDeserializeArgs;
24use crate::LayoutId;
25use crate::LayoutParts;
26use crate::LayoutReaderContext;
27use crate::LayoutReaderRef;
28use crate::VTable;
29use crate::children::OwnedLayoutChildren;
30use crate::layouts::flat::reader::FlatReader;
31use crate::segments::SegmentId;
32use crate::segments::SegmentSource;
33
34/// Check if inline array node is enabled.
35pub(super) fn flat_layout_inline_array_node() -> bool {
36    static FLAT_LAYOUT_INLINE_ARRAY_NODE: LazyLock<bool> =
37        LazyLock::new(|| env::var("FLAT_LAYOUT_INLINE_ARRAY_NODE").is_ok_and(|v| v == "1"));
38    *FLAT_LAYOUT_INLINE_ARRAY_NODE
39}
40
41/// Flat layout vtable.
42#[derive(Clone, Debug)]
43pub struct Flat;
44
45/// Backwards-compatible name for the flat layout plugin.
46pub use Flat as FlatLayoutEncoding;
47
48/// Flat-layout-specific data.
49#[derive(Clone, Debug)]
50pub struct FlatData {
51    segment_id: SegmentId,
52    ctx: ReadContext,
53    array_tree: Option<ByteBuffer>,
54}
55
56/// A terminal layout storing one serialized array segment.
57pub type FlatLayout = Layout<Flat>;
58
59impl VTable for Flat {
60    type LayoutData = FlatData;
61    type Metadata = ProstMetadata<FlatLayoutMetadata>;
62
63    fn id(&self) -> LayoutId {
64        static ID: CachedId = CachedId::new("vortex.flat");
65        *ID
66    }
67
68    fn metadata(layout: &Layout<Self>) -> Self::Metadata {
69        ProstMetadata(FlatLayoutMetadata {
70            array_encoding_tree: layout.array_tree.as_ref().map(|bytes| bytes.to_vec()),
71        })
72    }
73
74    fn deserialize(
75        &self,
76        args: &LayoutDeserializeArgs<'_>,
77        metadata: &FlatLayoutMetadata,
78    ) -> VortexResult<Self::LayoutData> {
79        if args.segment_ids.len() != 1 {
80            vortex_bail!("Flat layout must have exactly one segment ID");
81        }
82        if args.children.nchildren() != 0 {
83            vortex_bail!("Flat layout must not have children");
84        }
85        Ok(FlatData {
86            segment_id: args.segment_ids[0],
87            ctx: args.array_read_ctx.clone(),
88            array_tree: metadata
89                .array_encoding_tree
90                .as_ref()
91                .map(|bytes| ByteBuffer::from(bytes.clone())),
92        })
93    }
94
95    fn child_dtype(_layout: &Layout<Self>, idx: usize) -> VortexResult<DType> {
96        vortex_bail!("Flat layout has no child {idx}")
97    }
98
99    fn child_type(_layout: &Layout<Self>, idx: usize) -> LayoutChildType {
100        vortex_panic!("Flat layout has no child {idx}")
101    }
102
103    fn new_reader(
104        layout: &Layout<Self>,
105        name: Arc<str>,
106        segment_source: Arc<dyn SegmentSource>,
107        session: &VortexSession,
108        _ctx: &LayoutReaderContext,
109    ) -> VortexResult<LayoutReaderRef> {
110        Ok(Arc::new(FlatReader::new(
111            layout.clone(),
112            name,
113            segment_source,
114            session.clone(),
115        )))
116    }
117}
118
119impl Layout<Flat> {
120    /// Construct a flat layout without an inline array encoding tree.
121    pub fn new(row_count: u64, dtype: DType, segment_id: SegmentId, ctx: ReadContext) -> Self {
122        Self::new_with_metadata(row_count, dtype, segment_id, ctx, None)
123    }
124
125    /// Construct a flat layout with optional inline array metadata.
126    pub fn new_with_metadata(
127        row_count: u64,
128        dtype: DType,
129        segment_id: SegmentId,
130        ctx: ReadContext,
131        array_tree: Option<ByteBuffer>,
132    ) -> Self {
133        LayoutParts::new(
134            Flat,
135            dtype,
136            row_count,
137            vec![segment_id],
138            OwnedLayoutChildren::layout_children(Vec::new()),
139            FlatData {
140                segment_id,
141                ctx,
142                array_tree,
143            },
144        )
145        .into_typed()
146    }
147
148    /// Returns the serialized array segment ID.
149    pub fn segment_id(&self) -> SegmentId {
150        self.segment_id
151    }
152
153    /// Returns the array read context.
154    pub fn array_ctx(&self) -> &ReadContext {
155        &self.ctx
156    }
157
158    /// Returns the optional inline array encoding tree.
159    pub fn array_tree(&self) -> Option<&ByteBuffer> {
160        self.array_tree.as_ref()
161    }
162}
163
164#[derive(prost::Message)]
165pub struct FlatLayoutMetadata {
166    #[prost(optional, bytes, tag = "1")]
167    pub array_encoding_tree: Option<Vec<u8>>,
168}
vortex_layout/layouts/flat/mod.rs

vortex_layout/layouts/flat/
mod.rs
