vortex_ipc/messages/
encoder.rs

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

vortex_ipc/messages/
encoder.rs
