vortex_layout/
flatbuffers.rs

1use flatbuffers::{FlatBufferBuilder, WIPOffset, root};
2use vortex_dtype::DType;
3use vortex_error::{VortexExpect, VortexResult, vortex_err};
4use vortex_flatbuffers::{FlatBuffer, FlatBufferRoot, WriteFlatBuffer, layout};
5
6use crate::children::ViewedLayoutChildren;
7use crate::segments::SegmentId;
8use crate::{Layout, LayoutContext, LayoutRef};
9
10/// Parse a [`LayoutRef`] from a layout flatbuffer.
11pub fn layout_from_flatbuffer(
12    flatbuffer: FlatBuffer,
13    dtype: &DType,
14    ctx: &LayoutContext,
15) -> VortexResult<LayoutRef> {
16    let fb_layout = root::<layout::Layout>(&flatbuffer)?;
17    let encoding = ctx
18        .lookup_encoding(fb_layout.encoding())
19        .ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?;
20
21    // SAFETY: we validate the flatbuffer above in the `root` call, and extract a loc.
22    let viewed_children = unsafe {
23        ViewedLayoutChildren::new_unchecked(flatbuffer.clone(), fb_layout._tab.loc(), ctx.clone())
24    };
25
26    let layout = encoding.build(
27        dtype,
28        fb_layout.row_count(),
29        fb_layout
30            .metadata()
31            .map(|m| m.bytes())
32            .unwrap_or_else(|| &[]),
33        fb_layout
34            .segments()
35            .unwrap_or_default()
36            .iter()
37            .map(SegmentId::from)
38            .collect(),
39        &viewed_children,
40    )?;
41
42    Ok(layout)
43}
44
45impl dyn Layout + '_ {
46    /// Serialize the layout into a [`FlatBufferBuilder`].
47    pub fn flatbuffer_writer<'a>(
48        &'a self,
49        ctx: &'a LayoutContext,
50    ) -> impl WriteFlatBuffer<Target<'a> = layout::Layout<'a>> + FlatBufferRoot + 'a {
51        LayoutFlatBufferWriter { layout: self, ctx }
52    }
53}
54
55/// An adapter struct for writing a layout to a FlatBuffer.
56struct LayoutFlatBufferWriter<'a> {
57    layout: &'a dyn Layout,
58    ctx: &'a LayoutContext,
59}
60
61impl FlatBufferRoot for LayoutFlatBufferWriter<'_> {}
62
63impl WriteFlatBuffer for LayoutFlatBufferWriter<'_> {
64    type Target<'t> = layout::Layout<'t>;
65
66    fn write_flatbuffer<'fb>(
67        &self,
68        fbb: &mut FlatBufferBuilder<'fb>,
69    ) -> WIPOffset<Self::Target<'fb>> {
70        // First we recurse into the children and write them out
71        let child_layouts = self
72            .layout
73            .children()
74            .vortex_expect("Failed to load layout children");
75        let children = child_layouts
76            .iter()
77            .map(|layout| {
78                LayoutFlatBufferWriter {
79                    layout: layout.as_ref(),
80                    ctx: self.ctx,
81                }
82                .write_flatbuffer(fbb)
83            })
84            .collect::<Vec<_>>();
85        let children = (!children.is_empty()).then(|| fbb.create_vector(&children));
86
87        // Next we write out the metadata if it's non-empty.
88        let metadata = self.layout.metadata();
89        let metadata = (!metadata.is_empty()).then(|| fbb.create_vector(&metadata));
90
91        let segments = self
92            .layout
93            .segment_ids()
94            .into_iter()
95            .map(|s| *s)
96            .collect::<Vec<_>>();
97        let segments = (!segments.is_empty()).then(|| fbb.create_vector(&segments));
98
99        // Dictionary-encode the layout ID
100        let encoding = self.ctx.encoding_idx(&self.layout.encoding());
101
102        layout::Layout::create(
103            fbb,
104            &layout::LayoutArgs {
105                encoding,
106                row_count: self.layout.row_count(),
107                metadata,
108                children,
109                segments,
110            },
111        )
112    }
113}
vortex_layout/flatbuffers.rs

vortex_layout/
flatbuffers.rs
