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::VortexResult;
14use vortex_error::vortex_err;
15use vortex_flatbuffers::FlatBuffer;
16use vortex_flatbuffers::FlatBufferRoot;
17use vortex_flatbuffers::WriteFlatBuffer;
18use vortex_flatbuffers::layout;
19
20use crate::Layout;
21use crate::LayoutContext;
22use crate::LayoutRef;
23use crate::children::ViewedLayoutChildren;
24use crate::segments::SegmentId;
25use crate::session::LayoutRegistry;
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    ctx: &ArrayContext,
50    layouts: &LayoutRegistry,
51) -> VortexResult<LayoutRef> {
52    let fb_layout = root_with_opts::<layout::Layout>(&LAYOUT_VERIFIER, &flatbuffer)?;
53    let encoding_id = layout_ctx
54        .resolve(fb_layout.encoding())
55        .ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?;
56    let encoding = layouts
57        .find(&encoding_id)
58        .ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?;
59
60    // SAFETY: we validate the flatbuffer above in the `root` call, and extract a loc.
61    let viewed_children = unsafe {
62        ViewedLayoutChildren::new_unchecked(
63            flatbuffer.clone(),
64            fb_layout._tab.loc(),
65            ctx.clone(),
66            layout_ctx.clone(),
67            layouts.clone(),
68        )
69    };
70
71    let layout = encoding.build(
72        dtype,
73        fb_layout.row_count(),
74        fb_layout
75            .metadata()
76            .map(|m| m.bytes())
77            .unwrap_or_else(|| &[]),
78        fb_layout
79            .segments()
80            .unwrap_or_default()
81            .iter()
82            .map(SegmentId::from)
83            .collect(),
84        &viewed_children,
85        ctx,
86    )?;
87
88    Ok(layout)
89}
90
91impl dyn Layout + '_ {
92    /// Serialize the layout into a [`FlatBufferBuilder`].
93    pub fn flatbuffer_writer<'a>(
94        &'a self,
95        ctx: &'a LayoutContext,
96    ) -> impl WriteFlatBuffer<Target<'a> = layout::Layout<'a>> + FlatBufferRoot + 'a {
97        LayoutFlatBufferWriter { layout: self, ctx }
98    }
99}
100
101/// An adapter struct for writing a layout to a FlatBuffer.
102struct LayoutFlatBufferWriter<'a> {
103    layout: &'a dyn Layout,
104    ctx: &'a LayoutContext,
105}
106
107impl FlatBufferRoot for LayoutFlatBufferWriter<'_> {}
108
109impl WriteFlatBuffer for LayoutFlatBufferWriter<'_> {
110    type Target<'fb> = layout::Layout<'fb>;
111
112    fn write_flatbuffer<'fb>(
113        &self,
114        fbb: &mut FlatBufferBuilder<'fb>,
115    ) -> VortexResult<WIPOffset<Self::Target<'fb>>> {
116        // First we recurse into the children and write them out
117        let child_layouts = self.layout.children()?;
118        let children = child_layouts
119            .iter()
120            .map(|layout| {
121                LayoutFlatBufferWriter {
122                    layout: layout.as_ref(),
123                    ctx: self.ctx,
124                }
125                .write_flatbuffer(fbb)
126            })
127            .collect::<VortexResult<Vec<_>>>()?;
128        let children = (!children.is_empty()).then(|| fbb.create_vector(&children));
129
130        // Next we write out the metadata if it's non-empty.
131        let metadata = self.layout.metadata();
132        let metadata = (!metadata.is_empty()).then(|| fbb.create_vector(&metadata));
133
134        let segments = self
135            .layout
136            .segment_ids()
137            .into_iter()
138            .map(|s| *s)
139            .collect::<Vec<_>>();
140        let segments = (!segments.is_empty()).then(|| fbb.create_vector(&segments));
141
142        // Dictionary-encode the layout ID
143        let encoding = self.ctx.intern(&self.layout.encoding_id()).ok_or_else(|| {
144            vortex_err!(
145                "Failed to intern layout encoding ID: {}",
146                self.layout.encoding_id()
147            )
148        })?;
149
150        Ok(layout::Layout::create(
151            fbb,
152            &layout::LayoutArgs {
153                encoding,
154                row_count: self.layout.row_count(),
155                metadata,
156                children,
157                segments,
158            },
159        ))
160    }
161}
vortex_layout/flatbuffers.rs

vortex_layout/
flatbuffers.rs
