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

vortex_layout/layouts/flat/
mod.rs
