vortex_array/flatbuffers/
traits.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Traits for reading and writing Vortex types as flatbuffers.
5
6use flatbuffers::FlatBufferBuilder;
7use flatbuffers::Follow;
8use flatbuffers::InvalidFlatbuffer;
9use flatbuffers::Verifiable;
10use flatbuffers::WIPOffset;
11use flatbuffers::root;
12use vortex_buffer::ByteBuffer;
13use vortex_buffer::ConstByteBuffer;
14use vortex_error::VortexResult;
15
16/// We define a const-aligned byte buffer for flatbuffers with 8-byte alignment.
17///
18/// This is based on the assumption that the maximum primitive type is 8 bytes.
19/// See: <https://groups.google.com/g/flatbuffers/c/PSgQeWeTx_g>
20pub type FlatBuffer = ConstByteBuffer<8>;
21
22/// Marker trait for types that can be the root of a FlatBuffer.
23pub trait FlatBufferRoot {}
24
25/// Trait for reading a type from a FlatBuffer.
26pub trait ReadFlatBuffer: Sized {
27    /// The FlatBuffer type that this type can be read from.
28    type Source<'a>: Verifiable + Follow<'a>;
29    /// The error type returned when reading fails.
30    type Error: From<InvalidFlatbuffer>;
31
32    /// Reads this type from a FlatBuffer source.
33    fn read_flatbuffer<'buf>(
34        fb: &<Self::Source<'buf> as Follow<'buf>>::Inner,
35    ) -> Result<Self, Self::Error>;
36
37    /// Reads this type from bytes representing a FlatBuffer source.
38    fn read_flatbuffer_bytes<'buf>(bytes: &'buf [u8]) -> Result<Self, Self::Error>
39    where
40        <Self as ReadFlatBuffer>::Source<'buf>: 'buf,
41    {
42        let fb = root::<Self::Source<'buf>>(bytes)?;
43        Self::read_flatbuffer(&fb)
44    }
45}
46
47/// Trait for writing a type to a FlatBuffer.
48pub trait WriteFlatBuffer {
49    /// The FlatBuffer type that this type can be written to.
50    type Target<'a>;
51
52    /// Writes this type to a FlatBuffer builder.
53    fn write_flatbuffer<'fb>(
54        &self,
55        fbb: &mut FlatBufferBuilder<'fb>,
56    ) -> VortexResult<WIPOffset<Self::Target<'fb>>>;
57}
58
59/// Extension trait for types that can be written as FlatBuffer root objects.
60pub trait WriteFlatBufferExt: WriteFlatBuffer + FlatBufferRoot {
61    /// Writes self as a FlatBuffer root object into a [`FlatBuffer`] byte buffer.
62    fn write_flatbuffer_bytes(&self) -> VortexResult<FlatBuffer>;
63}
64
65impl<F: WriteFlatBuffer + FlatBufferRoot> WriteFlatBufferExt for F {
66    fn write_flatbuffer_bytes(&self) -> VortexResult<FlatBuffer> {
67        let mut fbb = FlatBufferBuilder::new();
68        let root_offset = self.write_flatbuffer(&mut fbb)?;
69        fbb.finish_minimal(root_offset);
70        let (vec, start) = fbb.collapse();
71        let end = vec.len();
72        Ok(FlatBuffer::align_from(
73            ByteBuffer::from(vec).slice(start..end),
74        ))
75    }
76}
vortex_array/flatbuffers/traits.rs

vortex_array/flatbuffers/
traits.rs
