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    /// Flat readers only ever register the end of the requested range, so flat layouts are
69    /// indivisible and split collection can skip materializing flat children.
70    fn is_indivisible(&self) -> bool {
71        true
72    }
73
74    fn metadata(layout: &Layout<Self>) -> Self::Metadata {
75        ProstMetadata(FlatLayoutMetadata {
76            array_encoding_tree: layout.array_tree.as_ref().map(|bytes| bytes.to_vec()),
77        })
78    }
79
80    fn deserialize(
81        &self,
82        args: &LayoutDeserializeArgs<'_>,
83        metadata: &FlatLayoutMetadata,
84    ) -> VortexResult<Self::LayoutData> {
85        if args.segment_ids.len() != 1 {
86            vortex_bail!("Flat layout must have exactly one segment ID");
87        }
88        if args.children.nchildren() != 0 {
89            vortex_bail!("Flat layout must not have children");
90        }
91        Ok(FlatData {
92            segment_id: args.segment_ids[0],
93            ctx: args.array_read_ctx.clone(),
94            array_tree: metadata
95                .array_encoding_tree
96                .as_ref()
97                .map(|bytes| ByteBuffer::from(bytes.clone())),
98        })
99    }
100
101    fn child_dtype(_layout: &Layout<Self>, idx: usize) -> VortexResult<DType> {
102        vortex_bail!("Flat layout has no child {idx}")
103    }
104
105    fn child_type(_layout: &Layout<Self>, idx: usize) -> LayoutChildType {
106        vortex_panic!("Flat layout has no child {idx}")
107    }
108
109    fn new_reader(
110        layout: &Layout<Self>,
111        name: Arc<str>,
112        segment_source: Arc<dyn SegmentSource>,
113        session: &VortexSession,
114        _ctx: &LayoutReaderContext,
115    ) -> VortexResult<LayoutReaderRef> {
116        Ok(Arc::new(FlatReader::new(
117            layout.clone(),
118            name,
119            segment_source,
120            session.clone(),
121        )))
122    }
123}
124
125impl Layout<Flat> {
126    /// Construct a flat layout without an inline array encoding tree.
127    pub fn new(row_count: u64, dtype: DType, segment_id: SegmentId, ctx: ReadContext) -> Self {
128        Self::new_with_metadata(row_count, dtype, segment_id, ctx, None)
129    }
130
131    /// Construct a flat layout with optional inline array metadata.
132    pub fn new_with_metadata(
133        row_count: u64,
134        dtype: DType,
135        segment_id: SegmentId,
136        ctx: ReadContext,
137        array_tree: Option<ByteBuffer>,
138    ) -> Self {
139        LayoutParts::new(
140            Flat,
141            dtype,
142            row_count,
143            vec![segment_id],
144            OwnedLayoutChildren::layout_children(Vec::new()),
145            FlatData {
146                segment_id,
147                ctx,
148                array_tree,
149            },
150        )
151        .into_typed()
152    }
153
154    /// Returns the serialized array segment ID.
155    pub fn segment_id(&self) -> SegmentId {
156        self.segment_id
157    }
158
159    /// Returns the array read context.
160    pub fn array_ctx(&self) -> &ReadContext {
161        &self.ctx
162    }
163
164    /// Returns the optional inline array encoding tree.
165    pub fn array_tree(&self) -> Option<&ByteBuffer> {
166        self.array_tree.as_ref()
167    }
168}
169
170#[derive(prost::Message)]
171pub struct FlatLayoutMetadata {
172    #[prost(optional, bytes, tag = "1")]
173    pub array_encoding_tree: Option<Vec<u8>>,
174}
vortex_layout/layouts/flat/mod.rs

vortex_layout/layouts/flat/
mod.rs
