vortex_layout/
flatbuffers.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::env;
5use std::sync::LazyLock;
6
7use flatbuffers::FlatBufferBuilder;
8use flatbuffers::VerifierOptions;
9use flatbuffers::WIPOffset;
10use flatbuffers::root_with_opts;
11use vortex_array::ArrayContext;
12use vortex_dtype::DType;
13use vortex_error::VortexExpect;
14use vortex_error::VortexResult;
15use vortex_error::vortex_err;
16use vortex_flatbuffers::FlatBuffer;
17use vortex_flatbuffers::FlatBufferRoot;
18use vortex_flatbuffers::WriteFlatBuffer;
19use vortex_flatbuffers::layout;
20
21use crate::Layout;
22use crate::LayoutContext;
23use crate::LayoutRef;
24use crate::children::ViewedLayoutChildren;
25use crate::segments::SegmentId;
26
27static LAYOUT_VERIFIER: LazyLock<VerifierOptions> = LazyLock::new(|| {
28    VerifierOptions {
29        // Overridden
30        max_tables: env::var("VORTEX_MAX_LAYOUT_TABLES")
31            .ok()
32            .and_then(|lmt| lmt.parse::<usize>().ok())
33            .unwrap_or(1000000),
34        max_depth: env::var("VORTEX_MAX_LAYOUT_DEPTH")
35            .ok()
36            .and_then(|lmt| lmt.parse::<usize>().ok())
37            .unwrap_or(64),
38        // Defaults from flatbuffers
39        max_apparent_size: 1 << 31,
40        ignore_missing_null_terminator: false,
41    }
42});
43
44/// Parse a [`LayoutRef`] from a layout flatbuffer.
45pub fn layout_from_flatbuffer(
46    flatbuffer: FlatBuffer,
47    dtype: &DType,
48    layout_ctx: &LayoutContext,
49    array_ctx: &ArrayContext,
50) -> VortexResult<LayoutRef> {
51    let fb_layout = root_with_opts::<layout::Layout>(&LAYOUT_VERIFIER, &flatbuffer)?;
52    let encoding = layout_ctx
53        .lookup_encoding(fb_layout.encoding())
54        .ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?;
55
56    // SAFETY: we validate the flatbuffer above in the `root` call, and extract a loc.
57    let viewed_children = unsafe {
58        ViewedLayoutChildren::new_unchecked(
59            flatbuffer.clone(),
60            fb_layout._tab.loc(),
61            array_ctx.clone(),
62            layout_ctx.clone(),
63        )
64    };
65
66    let layout = encoding.build(
67        dtype,
68        fb_layout.row_count(),
69        fb_layout
70            .metadata()
71            .map(|m| m.bytes())
72            .unwrap_or_else(|| &[]),
73        fb_layout
74            .segments()
75            .unwrap_or_default()
76            .iter()
77            .map(SegmentId::from)
78            .collect(),
79        &viewed_children,
80        array_ctx.clone(),
81    )?;
82
83    Ok(layout)
84}
85
86impl dyn Layout + '_ {
87    /// Serialize the layout into a [`FlatBufferBuilder`].
88    pub fn flatbuffer_writer<'a>(
89        &'a self,
90        ctx: &'a LayoutContext,
91    ) -> impl WriteFlatBuffer<Target<'a> = layout::Layout<'a>> + FlatBufferRoot + 'a {
92        LayoutFlatBufferWriter { layout: self, ctx }
93    }
94}
95
96/// An adapter struct for writing a layout to a FlatBuffer.
97struct LayoutFlatBufferWriter<'a> {
98    layout: &'a dyn Layout,
99    ctx: &'a LayoutContext,
100}
101
102impl FlatBufferRoot for LayoutFlatBufferWriter<'_> {}
103
104impl WriteFlatBuffer for LayoutFlatBufferWriter<'_> {
105    type Target<'t> = layout::Layout<'t>;
106
107    fn write_flatbuffer<'fb>(
108        &self,
109        fbb: &mut FlatBufferBuilder<'fb>,
110    ) -> WIPOffset<Self::Target<'fb>> {
111        // First we recurse into the children and write them out
112        let child_layouts = self
113            .layout
114            .children()
115            .vortex_expect("Failed to load layout children");
116        let children = child_layouts
117            .iter()
118            .map(|layout| {
119                LayoutFlatBufferWriter {
120                    layout: layout.as_ref(),
121                    ctx: self.ctx,
122                }
123                .write_flatbuffer(fbb)
124            })
125            .collect::<Vec<_>>();
126        let children = (!children.is_empty()).then(|| fbb.create_vector(&children));
127
128        // Next we write out the metadata if it's non-empty.
129        let metadata = self.layout.metadata();
130        let metadata = (!metadata.is_empty()).then(|| fbb.create_vector(&metadata));
131
132        let segments = self
133            .layout
134            .segment_ids()
135            .into_iter()
136            .map(|s| *s)
137            .collect::<Vec<_>>();
138        let segments = (!segments.is_empty()).then(|| fbb.create_vector(&segments));
139
140        // Dictionary-encode the layout ID
141        let encoding = self.ctx.encoding_idx(&self.layout.encoding());
142
143        layout::Layout::create(
144            fbb,
145            &layout::LayoutArgs {
146                encoding,
147                row_count: self.layout.row_count(),
148                metadata,
149                children,
150                segments,
151            },
152        )
153    }
154}
vortex_layout/flatbuffers.rs

vortex_layout/
flatbuffers.rs
