laminar_core/serialization/
mod.rs

1//! Shared serialization helpers.
2//!
3//! - Arrow IPC: `RecordBatch` ↔ bytes conversion using the Arrow IPC stream format.
4//! - `jsonb_tags`: Canonical JSONB binary format type tag constants.
5
6/// Canonical JSONB binary format type tag constants.
7pub mod jsonb_tags;
8
9use arrow::buffer::Buffer;
10use arrow_array::RecordBatch;
11use arrow_ipc::reader::{StreamDecoder, StreamReader};
12use arrow_ipc::writer::{IpcWriteOptions, StreamWriter};
13use arrow_schema::{ArrowError, Schema};
14
15/// A growable byte writer that rejects payloads and retained capacities above a fixed limit.
16/// This is shared by Arrow IPC and archive encoders at checkpoint/shuffle boundaries.
17pub struct BoundedBytesWriter {
18    bytes: Vec<u8>,
19    limit: usize,
20}
21
22impl BoundedBytesWriter {
23    /// Create an empty bounded writer.
24    #[must_use]
25    pub fn new(limit: usize) -> Self {
26        Self::with_capacity(limit, 0)
27    }
28
29    fn with_capacity(limit: usize, initial_capacity: usize) -> Self {
30        Self {
31            bytes: Vec::with_capacity(initial_capacity.min(limit)),
32            limit,
33        }
34    }
35
36    /// Consume the writer and return its retained bytes without copying.
37    #[must_use]
38    pub fn into_vec(self) -> Vec<u8> {
39        self.bytes
40    }
41}
42
43impl std::io::Write for BoundedBytesWriter {
44    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
45        let next_len = self
46            .bytes
47            .len()
48            .checked_add(buf.len())
49            .filter(|next_len| *next_len <= self.limit)
50            .ok_or_else(|| {
51                std::io::Error::new(
52                    std::io::ErrorKind::OutOfMemory,
53                    "serialized payload exceeds its configured bound",
54                )
55            })?;
56        if next_len > self.bytes.capacity() {
57            let target_capacity = self
58                .bytes
59                .capacity()
60                .saturating_mul(2)
61                .max(next_len)
62                .min(self.limit);
63            self.bytes
64                .try_reserve_exact(target_capacity - self.bytes.len())
65                .map_err(std::io::Error::other)?;
66            if self.bytes.capacity() < next_len || self.bytes.capacity() > self.limit {
67                return Err(std::io::Error::new(
68                    std::io::ErrorKind::OutOfMemory,
69                    "serialized allocation exceeds its configured bound",
70                ));
71            }
72        }
73        self.bytes.extend_from_slice(buf);
74        debug_assert!(self.bytes.capacity() <= self.limit);
75        Ok(buf.len())
76    }
77
78    fn flush(&mut self) -> std::io::Result<()> {
79        Ok(())
80    }
81}
82
83/// Serializes a single [`RecordBatch`] to Arrow IPC stream bytes.
84///
85/// # Errors
86///
87/// Returns [`arrow_schema::ArrowError`] if IPC encoding fails.
88pub fn serialize_batch_stream(batch: &RecordBatch) -> Result<Vec<u8>, arrow_schema::ArrowError> {
89    let mut buf = Vec::new();
90    {
91        let mut writer = StreamWriter::try_new(&mut buf, &batch.schema())?;
92        writer.write(batch)?;
93        writer.finish()?;
94    }
95    Ok(buf)
96}
97
98/// Deserializes a single [`RecordBatch`] from Arrow IPC stream bytes.
99///
100/// # Errors
101///
102/// Returns [`arrow_schema::ArrowError`] if the bytes are invalid or contain no batches.
103pub fn deserialize_batch_stream(bytes: &[u8]) -> Result<RecordBatch, arrow_schema::ArrowError> {
104    let cursor = std::io::Cursor::new(bytes);
105    let mut reader = StreamReader::try_new(cursor, None)?;
106    reader.next().ok_or_else(|| {
107        arrow_schema::ArrowError::IpcError("no record batch in IPC stream".to_string())
108    })?
109}
110
111/// Test encoder for exercising incrementally chunked Arrow IPC streams.
112#[cfg(test)]
113pub(crate) struct BatchStreamEncoder {
114    writer: StreamWriter<Vec<u8>>,
115}
116
117#[cfg(test)]
118impl BatchStreamEncoder {
119    /// Encoder for `schema`; the schema message flushes out with the first batch.
120    ///
121    /// # Errors
122    /// [`ArrowError`] if the schema header can't be IPC-encoded.
123    pub(crate) fn new(schema: &Schema) -> Result<Self, ArrowError> {
124        Ok(Self {
125            writer: StreamWriter::try_new(Vec::new(), schema)?,
126        })
127    }
128
129    /// Encode one batch, returning the bytes written since the last call (the
130    /// first call also carries the schema).
131    ///
132    /// # Errors
133    /// [`ArrowError`] if IPC encoding fails.
134    pub(crate) fn encode(&mut self, batch: &RecordBatch) -> Result<Vec<u8>, ArrowError> {
135        self.writer.write(batch)?;
136        Ok(std::mem::take(self.writer.get_mut()))
137    }
138
139    /// Finish the stream, returning the end-of-stream marker to append after the
140    /// last [`encode`](Self::encode). Also lets the decoder flush a trailing
141    /// zero-row batch; no batches may be encoded after this.
142    ///
143    /// # Errors
144    /// [`ArrowError`] if writing the marker fails.
145    pub(crate) fn finish(&mut self) -> Result<Vec<u8>, ArrowError> {
146        self.writer.finish()?;
147        Ok(std::mem::take(self.writer.get_mut()))
148    }
149}
150
151/// Decoder for an incrementally chunked Arrow IPC stream. The first chunk's
152/// schema decodes all later schema-less chunks.
153#[derive(Debug, Default)]
154pub struct BatchStreamDecoder {
155    decoder: StreamDecoder,
156}
157
158impl BatchStreamDecoder {
159    /// Creates an empty decoder that has not yet seen a schema.
160    #[must_use]
161    pub fn new() -> Self {
162        Self::default()
163    }
164
165    /// Decode every complete batch in one chunk; a batch straddling a chunk
166    /// boundary is buffered until the rest arrives, preserving order.
167    ///
168    /// # Errors
169    /// [`ArrowError`] if the bytes aren't a valid continuation (e.g. a batch
170    /// before any schema).
171    pub fn decode_chunk(&mut self, bytes: Vec<u8>) -> Result<Vec<RecordBatch>, ArrowError> {
172        let mut buffer = Buffer::from_vec(bytes);
173        let mut batches = Vec::new();
174        // Drain the chunk: `decode` yields a batch each time one completes.
175        while !buffer.is_empty() {
176            if let Some(batch) = self.decoder.decode(&mut buffer)? {
177                batches.push(batch);
178            }
179        }
180        Ok(batches)
181    }
182
183    /// Verify that the last chunk ended between IPC messages rather than in a
184    /// buffered header, metadata block, or body. This does not finish or reset
185    /// the decoder.
186    #[cfg(feature = "cluster")]
187    pub(crate) fn ensure_message_boundary(&mut self) -> Result<(), ArrowError> {
188        self.decoder.finish()
189    }
190}
191
192/// Serializes batches as one Arrow IPC stream without allowing the writer to allocate beyond
193/// `max_bytes`.
194///
195/// # Errors
196///
197/// Returns [`arrow_schema::ArrowError`] if encoding fails or crosses the byte bound.
198pub fn serialize_batches_stream_bounded<'a, I>(
199    schema: &Schema,
200    batches: I,
201    max_bytes: usize,
202) -> Result<Vec<u8>, arrow_schema::ArrowError>
203where
204    I: IntoIterator<Item = &'a RecordBatch>,
205{
206    serialize_batches_stream_bounded_with_options(
207        schema,
208        batches,
209        max_bytes,
210        0,
211        IpcWriteOptions::default(),
212    )
213}
214
215/// Serializes batches as one LZ4-compressed Arrow IPC stream without allowing the output writer
216/// to allocate beyond `max_bytes`.
217///
218/// # Errors
219///
220/// Returns [`arrow_schema::ArrowError`] if compression or encoding fails, or if the encoded stream
221/// crosses the byte bound.
222pub fn serialize_batches_stream_lz4_bounded<'a, I>(
223    schema: &Schema,
224    batches: I,
225    max_bytes: usize,
226) -> Result<Vec<u8>, arrow_schema::ArrowError>
227where
228    I: IntoIterator<Item = &'a RecordBatch>,
229{
230    let options = IpcWriteOptions::default()
231        .try_with_compression(Some(arrow_ipc::CompressionType::LZ4_FRAME))?;
232    serialize_batches_stream_bounded_with_options(schema, batches, max_bytes, 0, options)
233}
234
235fn serialize_batches_stream_bounded_with_options<'a, I>(
236    schema: &Schema,
237    batches: I,
238    max_bytes: usize,
239    initial_capacity: usize,
240    options: IpcWriteOptions,
241) -> Result<Vec<u8>, arrow_schema::ArrowError>
242where
243    I: IntoIterator<Item = &'a RecordBatch>,
244{
245    let mut bounded = BoundedBytesWriter::with_capacity(max_bytes, initial_capacity);
246    {
247        let mut writer = StreamWriter::try_new_with_options(&mut bounded, schema, options)?;
248        for batch in batches {
249            writer.write(batch)?;
250        }
251        writer.finish()?;
252    }
253    Ok(bounded.bytes)
254}
255
256#[cfg(feature = "cluster")]
257/// Single-batch adapter that preserves the shuffle path's measured initial-capacity hint.
258pub(crate) fn serialize_batch_stream_bounded(
259    batch: &RecordBatch,
260    max_bytes: usize,
261    initial_capacity: usize,
262) -> Result<Vec<u8>, arrow_schema::ArrowError> {
263    serialize_batches_stream_bounded_with_options(
264        batch.schema().as_ref(),
265        std::iter::once(batch),
266        max_bytes,
267        initial_capacity,
268        IpcWriteOptions::default(),
269    )
270}
271
272#[cfg(test)]
273mod tests;
laminar_core/serialization/mod.rs

laminar_core/serialization/
mod.rs
