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::dtype::DType;
12use vortex_error::VortexResult;
13use vortex_error::vortex_err;
14use vortex_flatbuffers::FlatBuffer;
15use vortex_flatbuffers::FlatBufferRoot;
16use vortex_flatbuffers::WriteFlatBuffer;
17use vortex_flatbuffers::layout;
18use vortex_session::VortexSession;
19use vortex_session::registry::ReadContext;
20
21use crate::Layout;
22use crate::LayoutBuildContext;
23use crate::LayoutContext;
24use crate::LayoutRef;
25use crate::children::ViewedLayoutChildren;
26use crate::layouts::foreign::new_foreign_layout;
27use crate::segments::SegmentId;
28use crate::session::LayoutSessionExt;
29
30static LAYOUT_VERIFIER: LazyLock<VerifierOptions> = LazyLock::new(|| {
31    VerifierOptions {
32        // Overridden
33        max_tables: env::var("VORTEX_MAX_LAYOUT_TABLES")
34            .ok()
35            .and_then(|lmt| lmt.parse::<usize>().ok())
36            .unwrap_or(1000000),
37        max_depth: env::var("VORTEX_MAX_LAYOUT_DEPTH")
38            .ok()
39            .and_then(|lmt| lmt.parse::<usize>().ok())
40            .unwrap_or(64),
41        // Defaults from flatbuffers
42        max_apparent_size: 1 << 31,
43        ignore_missing_null_terminator: false,
44    }
45});
46
47/// Parse a [`LayoutRef`] from a layout flatbuffer.
48pub fn layout_from_flatbuffer(
49    flatbuffer: FlatBuffer,
50    dtype: &DType,
51    layout_ctx: &ReadContext,
52    ctx: &ReadContext,
53    session: &VortexSession,
54) -> VortexResult<LayoutRef> {
55    layout_from_flatbuffer_with_options(flatbuffer, dtype, layout_ctx, ctx, session, false)
56}
57
58/// Parse a [`LayoutRef`] from a layout flatbuffer with unknown-encoding behavior control.
59pub fn layout_from_flatbuffer_with_options(
60    flatbuffer: FlatBuffer,
61    dtype: &DType,
62    layout_ctx: &ReadContext,
63    ctx: &ReadContext,
64    session: &VortexSession,
65    allow_unknown: bool,
66) -> VortexResult<LayoutRef> {
67    let layout_session = session.layouts();
68    let layouts = layout_session.registry();
69    let fb_layout = root_with_opts::<layout::Layout>(&LAYOUT_VERIFIER, &flatbuffer)?;
70    let encoding_id = layout_ctx
71        .resolve(fb_layout.encoding())
72        .ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?;
73    let encoding = layouts.find(&encoding_id);
74
75    if encoding.is_none() && allow_unknown {
76        return foreign_layout_from_fb(fb_layout, dtype, layout_ctx);
77    }
78    let encoding =
79        encoding.ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?;
80
81    // SAFETY: we validate the flatbuffer above in the `root` call, and extract a loc.
82    let viewed_children = unsafe {
83        ViewedLayoutChildren::new_unchecked(
84            flatbuffer.clone(),
85            fb_layout._tab.loc(),
86            ctx.clone(),
87            layout_ctx.clone(),
88            layouts.clone(),
89            allow_unknown,
90            session.clone(),
91        )
92    };
93
94    let build_ctx = LayoutBuildContext {
95        session,
96        array_read_ctx: ctx,
97    };
98    let layout = encoding.build(
99        dtype,
100        fb_layout.row_count(),
101        fb_layout
102            .metadata()
103            .map(|m| m.bytes())
104            .unwrap_or_else(|| &[]),
105        fb_layout
106            .segments()
107            .unwrap_or_default()
108            .iter()
109            .map(SegmentId::from)
110            .collect(),
111        &viewed_children,
112        &build_ctx,
113    )?;
114
115    Ok(layout)
116}
117
118fn foreign_layout_from_fb(
119    fb_layout: layout::Layout<'_>,
120    dtype: &DType,
121    layout_ctx: &ReadContext,
122) -> VortexResult<LayoutRef> {
123    let encoding_id = layout_ctx
124        .resolve(fb_layout.encoding())
125        .ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?;
126
127    let children = fb_layout
128        .children()
129        .unwrap_or_default()
130        .iter()
131        .map(|child| foreign_layout_from_fb(child, dtype, layout_ctx))
132        .collect::<VortexResult<Vec<_>>>()?;
133
134    Ok(new_foreign_layout(
135        encoding_id,
136        dtype.clone(),
137        fb_layout.row_count(),
138        fb_layout
139            .metadata()
140            .map(|m| m.bytes().to_vec())
141            .unwrap_or_default(),
142        fb_layout
143            .segments()
144            .unwrap_or_default()
145            .iter()
146            .map(SegmentId::from)
147            .collect(),
148        children,
149    ))
150}
151
152impl dyn Layout + '_ {
153    /// Serialize the layout into a [`FlatBufferBuilder`].
154    pub fn flatbuffer_writer<'a>(
155        &'a self,
156        ctx: &'a LayoutContext,
157    ) -> impl WriteFlatBuffer<Target<'a> = layout::Layout<'a>> + FlatBufferRoot + 'a {
158        LayoutFlatBufferWriter { layout: self, ctx }
159    }
160}
161
162/// An adapter struct for writing a layout to a FlatBuffer.
163struct LayoutFlatBufferWriter<'a> {
164    layout: &'a dyn Layout,
165    ctx: &'a LayoutContext,
166}
167
168impl FlatBufferRoot for LayoutFlatBufferWriter<'_> {}
169
170impl WriteFlatBuffer for LayoutFlatBufferWriter<'_> {
171    type Target<'fb> = layout::Layout<'fb>;
172
173    fn write_flatbuffer<'fb>(
174        &self,
175        fbb: &mut FlatBufferBuilder<'fb>,
176    ) -> VortexResult<WIPOffset<Self::Target<'fb>>> {
177        // First we recurse into the children and write them out
178        let child_layouts = self.layout.children()?;
179        let children = child_layouts
180            .iter()
181            .map(|layout| {
182                LayoutFlatBufferWriter {
183                    layout: layout.as_ref(),
184                    ctx: self.ctx,
185                }
186                .write_flatbuffer(fbb)
187            })
188            .collect::<VortexResult<Vec<_>>>()?;
189        let children = (!children.is_empty()).then(|| fbb.create_vector(&children));
190
191        // Next we write out the metadata if it's non-empty.
192        let metadata = self.layout.metadata();
193        let metadata = (!metadata.is_empty()).then(|| fbb.create_vector(&metadata));
194
195        let segments = self
196            .layout
197            .segment_ids()
198            .into_iter()
199            .map(|s| *s)
200            .collect::<Vec<_>>();
201        let segments = (!segments.is_empty()).then(|| fbb.create_vector(&segments));
202
203        // Dictionary-encode the layout ID
204        let encoding = self.ctx.intern(&self.layout.encoding_id()).ok_or_else(|| {
205            vortex_err!(
206                "Failed to intern layout encoding ID: {}",
207                self.layout.encoding_id()
208            )
209        })?;
210
211        Ok(layout::Layout::create(
212            fbb,
213            &layout::LayoutArgs {
214                encoding,
215                row_count: self.layout.row_count(),
216                metadata,
217                children,
218                segments,
219            },
220        ))
221    }
222}
223
224#[cfg(test)]
225mod tests {
226    use flatbuffers::FlatBufferBuilder;
227    use vortex_array::dtype::DType;
228    use vortex_array::dtype::Nullability;
229    use vortex_flatbuffers::layout as fbl;
230    use vortex_session::registry::ReadContext;
231
232    use super::layout_from_flatbuffer_with_options;
233    use crate::LayoutEncodingId;
234    use crate::session::LayoutSession;
235
236    #[expect(clippy::disallowed_methods, reason = "test-only id")]
237    #[test]
238    fn unknown_layout_encoding_allow_unknown() {
239        let mut fbb = FlatBufferBuilder::new();
240
241        let child_metadata = fbb.create_vector(&[9u8]);
242        let child = fbl::Layout::create(
243            &mut fbb,
244            &fbl::LayoutArgs {
245                encoding: 1,
246                row_count: 3,
247                metadata: Some(child_metadata),
248                children: None,
249                segments: None,
250            },
251        );
252
253        let children = fbb.create_vector(&[child]);
254        let metadata = fbb.create_vector(&[1u8, 2, 3]);
255        let segments = fbb.create_vector(&[7u32]);
256        let root = fbl::Layout::create(
257            &mut fbb,
258            &fbl::LayoutArgs {
259                encoding: 0,
260                row_count: 10,
261                metadata: Some(metadata),
262                children: Some(children),
263                segments: Some(segments),
264            },
265        );
266        fbb.finish_minimal(root);
267        let (buf, start) = fbb.collapse();
268        let layout_buffer = vortex_flatbuffers::FlatBuffer::align_from(
269            vortex_buffer::ByteBuffer::from(buf).slice(start..),
270        );
271
272        let layout_ctx = ReadContext::new([
273            LayoutEncodingId::new("vortex.test.foreign_layout"),
274            LayoutEncodingId::new("vortex.test.foreign_child_layout"),
275        ]);
276        let array_ctx = ReadContext::new([]);
277        let session = vortex_array::array_session().with::<LayoutSession>();
278
279        let layout = layout_from_flatbuffer_with_options(
280            layout_buffer,
281            &DType::Variant(Nullability::Nullable),
282            &layout_ctx,
283            &array_ctx,
284            &session,
285            true,
286        )
287        .unwrap();
288
289        assert_eq!(layout.encoding_id().as_ref(), "vortex.test.foreign_layout");
290        assert_eq!(layout.row_count(), 10);
291        assert_eq!(layout.metadata(), vec![1, 2, 3]);
292        assert_eq!(layout.segment_ids().len(), 1);
293        assert_eq!(*layout.segment_ids()[0], 7);
294        assert_eq!(layout.nchildren(), 1);
295
296        let child = layout.child(0).unwrap();
297        assert_eq!(
298            child.encoding_id().as_ref(),
299            "vortex.test.foreign_child_layout"
300        );
301        assert_eq!(child.metadata(), vec![9]);
302    }
303}
vortex_layout/flatbuffers.rs

vortex_layout/
flatbuffers.rs
