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

vortex_layout/layouts/flat/
mod.rs
