vortex_ipc/messages/
encoder.rs

1use bytes::{Bytes, BytesMut};
2use flatbuffers::FlatBufferBuilder;
3use vortex_array::serde::SerializeOptions;
4use vortex_array::{Array, ArrayContext};
5use vortex_buffer::ByteBuffer;
6use vortex_dtype::DType;
7use vortex_error::VortexExpect;
8use vortex_flatbuffers::{FlatBuffer, WriteFlatBufferExt, message as fb};
9
10/// An IPC message ready to be passed to the encoder.
11pub enum EncoderMessage<'a> {
12    Array(&'a dyn Array),
13    Buffer(&'a ByteBuffer),
14    DType(&'a DType),
15}
16
17pub struct MessageEncoder {
18    /// A reusable buffer of zeros used for padding.
19    zeros: Bytes,
20}
21
22impl Default for MessageEncoder {
23    fn default() -> Self {
24        Self {
25            zeros: BytesMut::zeroed(u16::MAX as usize).freeze(),
26        }
27    }
28}
29
30impl MessageEncoder {
31    /// Encode an IPC message for writing to a byte stream.
32    ///
33    /// The returned buffers should be written contiguously to the stream.
34    pub fn encode(&mut self, message: EncoderMessage) -> Vec<Bytes> {
35        let mut buffers = vec![];
36
37        // We'll push one buffer as a placeholder for the flatbuffer message length, and one
38        // for the flatbuffer itself.
39        buffers.push(self.zeros.clone());
40        buffers.push(self.zeros.clone());
41
42        // We initialize the flatbuffer builder with a 4-byte vector that we will use to store
43        // the flatbuffer length into. By passing this vector into the FlatBufferBuilder, the
44        // flatbuffers internal alignment mechanisms will handle everything else for us.
45        // TODO(ngates): again, this a ton of padding...
46        let mut fbb = FlatBufferBuilder::from_vec(vec![0u8; 4]);
47
48        let (header, body_len) = match message {
49            EncoderMessage::Array(array) => {
50                // Currently we include a Context in every message. We could convert this to
51                // sending deltas later.
52                let ctx = ArrayContext::empty();
53                let array_buffers = array.serialize(&ctx, &SerializeOptions::default());
54                let body_len = array_buffers.iter().map(|b| b.len() as u64).sum::<u64>();
55
56                let array_encodings = ctx
57                    .encodings()
58                    .iter()
59                    .map(|e| fbb.create_string(e.id().as_ref()))
60                    .collect::<Vec<_>>();
61                let array_encodings = fbb.create_vector(array_encodings.as_slice());
62
63                let header = fb::ArrayMessage::create(
64                    &mut fbb,
65                    &fb::ArrayMessageArgs {
66                        row_count: u32::try_from(array.len())
67                            .vortex_expect("Array length must fit into u32"),
68                        encodings: Some(array_encodings),
69                    },
70                )
71                .as_union_value();
72
73                buffers.extend(array_buffers.into_iter().map(|b| b.into_inner()));
74
75                (header, body_len)
76            }
77            EncoderMessage::Buffer(buffer) => {
78                let header = fb::BufferMessage::create(
79                    &mut fbb,
80                    &fb::BufferMessageArgs {
81                        alignment_exponent: buffer.alignment().exponent(),
82                    },
83                )
84                .as_union_value();
85                let body_len = buffer.len() as u64;
86                buffers.push(buffer.clone().into_inner());
87
88                (header, body_len)
89            }
90            EncoderMessage::DType(dtype) => {
91                let header =
92                    fb::DTypeMessage::create(&mut fbb, &fb::DTypeMessageArgs {}).as_union_value();
93
94                let buffer = dtype.write_flatbuffer_bytes().into_inner().into_inner();
95                let body_len = buffer.len() as u64;
96                buffers.push(buffer);
97
98                (header, body_len)
99            }
100        };
101
102        let mut msg = fb::MessageBuilder::new(&mut fbb);
103        msg.add_version(Default::default());
104        msg.add_header_type(match message {
105            EncoderMessage::Array(_) => fb::MessageHeader::ArrayMessage,
106            EncoderMessage::Buffer(_) => fb::MessageHeader::BufferMessage,
107            EncoderMessage::DType(_) => fb::MessageHeader::DTypeMessage,
108        });
109        msg.add_header(header);
110        msg.add_body_size(body_len);
111        let msg = msg.finish();
112
113        // Finish the flatbuffer and swap it out for the placeholder buffer.
114        fbb.finish_minimal(msg);
115        let (fbv, pos) = fbb.collapse();
116        let fb_buffer = FlatBuffer::copy_from(&fbv[pos..]);
117        let fb_buffer_len = u32::try_from(fb_buffer.len())
118            .vortex_expect("IPC flatbuffer headers must fit into u32 bytes");
119
120        buffers[0] = Bytes::from(fb_buffer_len.to_le_bytes().to_vec());
121        buffers[1] = fb_buffer.into_inner().into_inner();
122
123        buffers
124    }
125}
vortex_ipc/messages/encoder.rs

vortex_ipc/messages/
encoder.rs
