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

vortex_layout/
flatbuffers.rs
