Skip to main content

dora_node_api/node/arrow_utils/
ipc_encode.rs

1//! Hand-rolled, 1-copy Arrow IPC stream encoder.
2//!
3//! Arrow's official [`StreamWriter`](arrow::ipc::writer::StreamWriter) always
4//! stages the record-batch body in an internal `Vec` before writing it out, so
5//! encoding a message through it copies the payload at least twice. This module
6//! provides a *fast path* that writes the IPC flatbuffer headers and copies each
7//! array buffer **directly into a caller-provided `&mut [u8]`** — exactly one
8//! copy of the payload, straight into the (shared-memory) sample. For types the
9//! fast path does not handle it falls back to the official writer.
10//!
11//! The fast-path output is a normal Arrow IPC stream and decodes through the
12//! official [`StreamDecoder`](arrow::ipc::reader::StreamDecoder) — including the
13//! zero-copy [`decode_arrow_ipc_zero_copy`](super::decode_arrow_ipc_zero_copy)
14//! receive path — because every body buffer is placed at a 64-byte-aligned
15//! offset, matching what the official writer produces with its default
16//! `alignment = 64`.
17//!
18//! ## How the fast path stays correct without slice-truncation logic
19//!
20//! Arrow's writer contains a lot of per-type code to *truncate* buffers for
21//! sliced arrays. We sidestep all of it with two rules:
22//!  * **Require `offset() == 0` on every node** — so logical element `i` lives at
23//!    physical position `i`. Any array (or child) with a non-zero offset routes
24//!    to the fallback.
25//!  * **Copy each data buffer in full.** Arrow tolerates buffers that are larger
26//!    than strictly required for `len` elements, so copying the whole buffer (a
27//!    freshly built array's buffers are exactly sized anyway) always decodes to
28//!    a logically-equal array. The only generated buffer is the all-ones
29//!    validity bitmap for a node with no nulls, exactly as arrow emits.
30//!
31//! Two types need their children sliced before recursion because the child's
32//! IPC length is the parent's rather than the child's own: `Struct` (each field
33//! to the struct's `len`) and `FixedSizeList` (its child to `len * value_size`).
34//! `List`/`LargeList` children are bounded by an offsets buffer and recursed at
35//! full length.
36
37use arrow::array::ArrayData;
38use arrow::buffer::Buffer as ArrowBuffer;
39use arrow::ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteOptions};
40use arrow::ipc::{Buffer as IpcBuffer, FieldNode, MessageHeader, MetadataVersion};
41use arrow_schema::{DataType, Field, Schema};
42use eyre::{Context, bail, eyre};
43
44use super::ARROW_BUFFER_ALIGNMENT as ALIGN;
45
46/// IPC stream continuation marker (precedes every message length prefix).
47const CONTINUATION_MARKER: [u8; 4] = [0xff, 0xff, 0xff, 0xff];
48/// Bytes of the message length/continuation prefix (continuation + i32 length).
49const PREFIX_LEN: usize = 8;
50
51#[inline]
52fn round_up(n: usize, align: usize) -> usize {
53    debug_assert!(align.is_power_of_two());
54    (n + align - 1) & !(align - 1)
55}
56
57/// Whether the fast path handles this `DataType` directly (per node — children
58/// are validated by recursion). Everything else uses the official-writer
59/// fallback: dictionary, any `*View`, `Union`, `Map`, run-end-encoded, etc.
60fn is_fast_path_type(data_type: &DataType) -> bool {
61    use DataType::*;
62    matches!(
63        data_type,
64        Null | Boolean
65            | Int8
66            | Int16
67            | Int32
68            | Int64
69            | UInt8
70            | UInt16
71            | UInt32
72            | UInt64
73            | Float16
74            | Float32
75            | Float64
76            | Timestamp(_, _)
77            | Date32
78            | Date64
79            | Time32(_)
80            | Time64(_)
81            | Duration(_)
82            | Interval(_)
83            | Decimal128(_, _)
84            | Decimal256(_, _)
85            | FixedSizeBinary(_)
86            | Binary
87            | LargeBinary
88            | Utf8
89            | LargeUtf8
90            | List(_)
91            | LargeList(_)
92            | FixedSizeList(_, _)
93            | Struct(_)
94    )
95}
96
97/// One body buffer to emit: a 64-aligned body offset, its byte length, and where
98/// the bytes come from.
99struct Desc {
100    offset: usize,
101    len: usize,
102    src: BufferSrc,
103}
104
105enum BufferSrc {
106    /// Copy `len` bytes from this buffer starting at the given byte offset.
107    Bytes(ArrowBuffer, usize),
108    /// Fill `len` bytes with `0xff` (an all-valid validity bitmap).
109    AllOnes,
110}
111
112#[derive(Default)]
113struct Layout {
114    nodes: Vec<FieldNode>,
115    ipc_buffers: Vec<IpcBuffer>,
116    descs: Vec<Desc>,
117    body_len: usize,
118}
119
120impl Layout {
121    /// Append a buffer descriptor at the current (64-aligned) offset and advance
122    /// past it (padded to 64). Returns `None` if the running body size would
123    /// overflow `usize` or exceed [`super::MAX_IPC_BYTES`] — such arrays route to
124    /// the fallback rather than wrapping to a too-small offset.
125    fn push_buffer(&mut self, off: &mut usize, len: usize, src: BufferSrc) -> Option<()> {
126        // `*off` is always 64-aligned here (the invariant `push_buffer` keeps).
127        let padded = len.checked_add(ALIGN - 1).map(|n| n & !(ALIGN - 1))?;
128        let next = off.checked_add(padded)?;
129        if next > super::MAX_IPC_BYTES {
130            return None;
131        }
132        self.ipc_buffers
133            .push(IpcBuffer::new(*off as i64, len as i64));
134        self.descs.push(Desc {
135            offset: *off,
136            len,
137            src,
138        });
139        *off = next;
140        Some(())
141    }
142}
143
144/// Walk `array` building the IPC field-node list, buffer descriptors, and body
145/// length. Returns `None` if any node is not fast-path eligible (unsupported
146/// type or non-zero offset), so the caller falls back to the official writer.
147fn build_layout(array: &ArrayData) -> Option<Layout> {
148    let mut layout = Layout::default();
149    let mut off = 0usize;
150    build_layout_rec(array, &mut layout, &mut off)?;
151    layout.body_len = off;
152    Some(layout)
153}
154
155fn build_layout_rec(array: &ArrayData, layout: &mut Layout, off: &mut usize) -> Option<()> {
156    let data_type = array.data_type();
157    if !is_fast_path_type(data_type) || array.offset() != 0 {
158        return None;
159    }
160
161    let len = array.len();
162    // NullArray reports `null_count == 0` on its `ArrayData`, but the IPC field
163    // node records every element as null (matching arrow's writer).
164    let null_count = if matches!(data_type, DataType::Null) {
165        len
166    } else {
167        array.null_count()
168    };
169    layout
170        .nodes
171        .push(FieldNode::new(len as i64, null_count as i64));
172
173    // Validity bitmap (every type except `Null` carries one in IPC V5; the
174    // fast-path set excludes the other no-validity types — Union/RunEndEncoded).
175    if !matches!(data_type, DataType::Null) {
176        match array.nulls() {
177            Some(nulls) => {
178                if nulls.inner().offset() != 0 {
179                    return None;
180                }
181                let sliced = nulls.inner().sliced();
182                let bytes = sliced.len();
183                layout.push_buffer(off, bytes, BufferSrc::Bytes(sliced, 0))?;
184            }
185            None => {
186                let bytes = len.div_ceil(8);
187                layout.push_buffer(off, bytes, BufferSrc::AllOnes)?;
188            }
189        }
190    }
191
192    // Data buffers, copied in full (empty for Struct/FixedSizeList).
193    for buffer in array.buffers() {
194        layout.push_buffer(off, buffer.len(), BufferSrc::Bytes(buffer.clone(), 0))?;
195    }
196
197    // Children.
198    match data_type {
199        DataType::FixedSizeList(_, value_size) => {
200            // The child length is implied (`len * value_size`), not carried by
201            // an offsets buffer, so slice it to exactly that before recursing.
202            let n = len.checked_mul(*value_size as usize)?;
203            let child = array.child_data().first()?;
204            if child.len() < n {
205                return None;
206            }
207            build_layout_rec(&child.slice(0, n), layout, off)?;
208        }
209        DataType::Struct(_) => {
210            // Every struct field's IPC field node must report the struct's row
211            // count. A child `ArrayData` may legally be *longer* than the struct
212            // (e.g. built via the low-level builder), so slice each child to the
213            // struct's `len` before recursing — otherwise the child node would
214            // declare a different length and the stream would be un-decodable.
215            for child in array.child_data() {
216                if child.len() < len {
217                    return None;
218                }
219                build_layout_rec(&child.slice(0, len), layout, off)?;
220            }
221        }
222        _ => {
223            // List/LargeList: the single child is the values array, bounded by
224            // the offsets buffer (its length is independent of the parent's), so
225            // it is recursed at full length.
226            for child in array.child_data() {
227                build_layout_rec(child, layout, off)?;
228            }
229        }
230    }
231
232    Some(())
233}
234
235/// Everything needed to write the stream, plus the exact total length.
236struct Prepared {
237    layout: Layout,
238    schema_message: Vec<u8>,
239    record_batch_message: Vec<u8>,
240    schema_block: usize,
241    record_batch_block: usize,
242    total: usize,
243}
244
245fn ipc_write_options() -> eyre::Result<IpcWriteOptions> {
246    IpcWriteOptions::try_new(ALIGN, false, MetadataVersion::V5)
247        .map_err(|e| eyre!("failed to build Arrow IPC write options: {e}"))
248}
249
250/// Build the schema IPC message flatbuffer (one nullable field named `data`),
251/// matching what `encode_arrow_ipc` / the official writer emit.
252fn build_schema_message(data_type: &DataType) -> eyre::Result<Vec<u8>> {
253    let schema = Schema::new(vec![Field::new("data", data_type.clone(), true)]);
254    let options = ipc_write_options()?;
255    let mut tracker = DictionaryTracker::new(false);
256    let encoded = IpcDataGenerator {}.schema_to_bytes_with_dictionary_tracker(
257        &schema,
258        &mut tracker,
259        &options,
260    );
261    Ok(encoded.ipc_message)
262}
263
264/// Hand-build the RecordBatch IPC message flatbuffer (header only — no body),
265/// mirroring arrow's `record_batch_to_bytes`.
266fn build_record_batch_message(
267    num_rows: usize,
268    nodes: &[FieldNode],
269    buffers: &[IpcBuffer],
270    body_len: usize,
271) -> Vec<u8> {
272    use flatbuffers::FlatBufferBuilder;
273
274    let mut fbb = FlatBufferBuilder::new();
275    let buffers_fb = fbb.create_vector(buffers);
276    let nodes_fb = fbb.create_vector(nodes);
277
278    let record_batch = {
279        let mut builder = arrow::ipc::RecordBatchBuilder::new(&mut fbb);
280        builder.add_length(num_rows as i64);
281        builder.add_nodes(nodes_fb);
282        builder.add_buffers(buffers_fb);
283        builder.finish()
284    };
285
286    let message = {
287        let mut builder = arrow::ipc::MessageBuilder::new(&mut fbb);
288        builder.add_version(MetadataVersion::V5);
289        builder.add_header_type(MessageHeader::RecordBatch);
290        builder.add_bodyLength(body_len as i64);
291        builder.add_header(record_batch.as_union_value());
292        builder.finish()
293    };
294
295    fbb.finish(message, None);
296    fbb.finished_data().to_vec()
297}
298
299fn prepare(array: &ArrayData) -> Option<Prepared> {
300    let layout = build_layout(array)?;
301    let schema_message = build_schema_message(array.data_type()).ok()?;
302    let record_batch_message = build_record_batch_message(
303        array.len(),
304        &layout.nodes,
305        &layout.ipc_buffers,
306        layout.body_len,
307    );
308    let schema_block = round_up(PREFIX_LEN + schema_message.len(), ALIGN);
309    let record_batch_block = round_up(PREFIX_LEN + record_batch_message.len(), ALIGN);
310    // schema block + record-batch header block + body + end-of-stream (8 bytes).
311    let total = schema_block + record_batch_block + layout.body_len + PREFIX_LEN;
312    Some(Prepared {
313        layout,
314        schema_message,
315        record_batch_message,
316        schema_block,
317        record_batch_block,
318        total,
319    })
320}
321
322/// Exact byte length of the IPC stream [`encode_ipc_into`] would write, or
323/// `None` if `array` is not fast-path eligible (use [`encode_ipc_to_vec`]).
324///
325/// Lets a caller size the (shared-memory) sample before encoding into it.
326pub fn ipc_fast_path_len(array: &ArrayData) -> Option<usize> {
327    prepare(array).map(|p| p.total)
328}
329
330/// Frame one IPC message into `dst[at..]`: continuation marker, i32-LE metadata
331/// length, the flatbuffer, then zero padding so the block is 64-aligned (which
332/// makes the following body — or next message — start 64-aligned). Returns the
333/// block size written.
334fn write_framed_message(dst: &mut [u8], at: usize, flatbuffer: &[u8]) -> usize {
335    let block = round_up(PREFIX_LEN + flatbuffer.len(), ALIGN);
336    let metadata_len = (block - PREFIX_LEN) as i32;
337    dst[at..at + 4].copy_from_slice(&CONTINUATION_MARKER);
338    dst[at + 4..at + 8].copy_from_slice(&metadata_len.to_le_bytes());
339    dst[at + 8..at + 8 + flatbuffer.len()].copy_from_slice(flatbuffer);
340    // Zero the trailing padding (the sample buffer may be uninitialized SHM).
341    dst[at + 8 + flatbuffer.len()..at + block].fill(0);
342    block
343}
344
345/// Encode `array` as a complete Arrow IPC stream directly into `dst`, copying
346/// each array buffer exactly once.
347///
348/// `dst.len()` must equal [`ipc_fast_path_len(array)`](ipc_fast_path_len);
349/// `array` must be fast-path eligible (it is when `ipc_fast_path_len` returned
350/// `Some`).
351pub fn encode_ipc_into(array: &ArrayData, dst: &mut [u8]) -> eyre::Result<()> {
352    let prepared =
353        prepare(array).ok_or_else(|| eyre!("array is not Arrow IPC fast-path eligible"))?;
354    if dst.len() != prepared.total {
355        bail!(
356            "destination size {} does not match required IPC length {}",
357            dst.len(),
358            prepared.total
359        );
360    }
361
362    let mut at = 0;
363    at += write_framed_message(dst, at, &prepared.schema_message);
364    debug_assert_eq!(at, prepared.schema_block);
365    at += write_framed_message(dst, at, &prepared.record_batch_message);
366    debug_assert_eq!(at, prepared.schema_block + prepared.record_batch_block);
367
368    let body_start = at;
369    write_body(dst, body_start, &prepared.layout);
370    at = body_start + prepared.layout.body_len;
371
372    // End-of-stream: continuation marker + zero length.
373    dst[at..at + 4].copy_from_slice(&CONTINUATION_MARKER);
374    dst[at + 4..at + 8].copy_from_slice(&0i32.to_le_bytes());
375    debug_assert_eq!(at + PREFIX_LEN, prepared.total);
376
377    Ok(())
378}
379
380/// Copy each body buffer to its 64-aligned offset within `dst[body_start..]`,
381/// zeroing per-buffer alignment padding (the destination may be uninitialized
382/// SHM). Shared by the full-stream and schema-less-batch encoders.
383fn write_body(dst: &mut [u8], body_start: usize, layout: &Layout) {
384    for desc in &layout.descs {
385        let start = body_start + desc.offset;
386        let end = start + desc.len;
387        match &desc.src {
388            BufferSrc::Bytes(buffer, src_off) => {
389                dst[start..end].copy_from_slice(&buffer.as_slice()[*src_off..*src_off + desc.len]);
390            }
391            BufferSrc::AllOnes => dst[start..end].fill(0xff),
392        }
393        let padded_end = body_start + desc.offset + round_up(desc.len, ALIGN);
394        dst[end..padded_end].fill(0);
395    }
396}
397
398/// Encode the IPC **schema message** for `data_type` as a framed,
399/// 64-byte-aligned block (one nullable field named `data`).
400///
401/// This is the schema prefix of a stream, sent once. A receiver feeds it to a
402/// persistent [`StreamDecoder`](arrow::ipc::reader::StreamDecoder) to prime it,
403/// then decodes a sequence of schema-less batch messages from
404/// [`encode_batch_into`] against the same decoder (the W3 schema-once path).
405pub fn encode_schema_message(data_type: &DataType) -> eyre::Result<Vec<u8>> {
406    let schema_message = build_schema_message(data_type)?;
407    let block = round_up(PREFIX_LEN + schema_message.len(), ALIGN);
408    let mut dst = vec![0u8; block];
409    write_framed_message(&mut dst, 0, &schema_message);
410    Ok(dst)
411}
412
413/// Exact byte length of the schema-less batch message [`encode_batch_into`]
414/// would write (record-batch header block + body + 8-byte end-of-stream marker,
415/// **no** schema prefix), or `None` if `array` is not fast-path eligible.
416pub fn batch_fast_path_len(array: &ArrayData) -> Option<usize> {
417    let prepared = prepare(array)?;
418    Some(prepared.record_batch_block + prepared.layout.body_len + PREFIX_LEN)
419}
420
421/// Encode `array` as a schema-less Arrow IPC **batch message** into `dst`:
422/// the record-batch header block, the body, and a trailing 8-byte end-of-stream
423/// marker (no schema prefix). Decoded by a [`StreamDecoder`] already primed with
424/// the matching [`encode_schema_message`]. The trailing marker lets the decoder
425/// flush a 0-row batch (empty body); a non-empty batch never reads it. `dst.len()`
426/// must equal [`batch_fast_path_len(array)`](batch_fast_path_len).
427pub fn encode_batch_into(array: &ArrayData, dst: &mut [u8]) -> eyre::Result<()> {
428    let prepared =
429        prepare(array).ok_or_else(|| eyre!("array is not Arrow IPC fast-path eligible"))?;
430    let expected = prepared.record_batch_block + prepared.layout.body_len + PREFIX_LEN;
431    if dst.len() != expected {
432        bail!(
433            "destination size {} does not match required batch length {expected}",
434            dst.len(),
435        );
436    }
437    let at = write_framed_message(dst, 0, &prepared.record_batch_message);
438    debug_assert_eq!(at, prepared.record_batch_block);
439    write_body(dst, at, &prepared.layout);
440    // End-of-stream marker (continuation + zero length), matching the tail of a
441    // full stream so an empty batch flushes through the persistent decoder.
442    let body_end = at + prepared.layout.body_len;
443    dst[body_end..body_end + 4].copy_from_slice(&CONTINUATION_MARKER);
444    dst[body_end + 4..body_end + PREFIX_LEN].copy_from_slice(&0i32.to_le_bytes());
445    Ok(())
446}
447
448/// Fallback encoder for any array (including non-fast-path types): produce a
449/// full Arrow IPC stream `Vec` via the official writer. The caller copies this
450/// into the sample, so this path costs two payload copies.
451pub fn encode_ipc_to_vec(array: &ArrayData) -> eyre::Result<Vec<u8>> {
452    super::encode_arrow_ipc(array).context("Arrow IPC fallback encode")
453}
454
455/// IPC layout of a `UInt8` array of `data_len` elements (no nulls): the byte
456/// offsets and message blocks, computed directly without materializing the
457/// data buffer (so it is cheap for a large image/tensor).
458struct Uint8Layout {
459    schema_message: Vec<u8>,
460    record_batch_message: Vec<u8>,
461    validity_len: usize,
462    validity_padded: usize,
463    body_len: usize,
464    total: usize,
465    data_offset: usize,
466}
467
468fn uint8_layout(data_len: usize) -> eyre::Result<Uint8Layout> {
469    if data_len > super::MAX_IPC_BYTES {
470        bail!(
471            "UInt8 payload too large: {data_len} bytes (max {})",
472            super::MAX_IPC_BYTES
473        );
474    }
475    let validity_len = data_len.div_ceil(8);
476    let validity_padded = round_up(validity_len, ALIGN);
477    let body_len = validity_padded + round_up(data_len, ALIGN);
478    // Matches what `encode_ipc_into` emits for a no-null UInt8Array: one field
479    // node, then [validity all-ones | data].
480    let nodes = [FieldNode::new(data_len as i64, 0)];
481    let buffers = [
482        IpcBuffer::new(0, validity_len as i64),
483        IpcBuffer::new(validity_padded as i64, data_len as i64),
484    ];
485    let record_batch_message = build_record_batch_message(data_len, &nodes, &buffers, body_len);
486    let schema_message = build_schema_message(&DataType::UInt8)?;
487    let schema_block = round_up(PREFIX_LEN + schema_message.len(), ALIGN);
488    let record_batch_block = round_up(PREFIX_LEN + record_batch_message.len(), ALIGN);
489    let total = schema_block + record_batch_block + body_len + PREFIX_LEN;
490    let data_offset = schema_block + record_batch_block + validity_padded;
491    Ok(Uint8Layout {
492        schema_message,
493        record_batch_message,
494        validity_len,
495        validity_padded,
496        body_len,
497        total,
498        data_offset,
499    })
500}
501
502/// Total IPC stream length for a no-null `UInt8` array of `data_len` elements.
503/// Lets a caller size the sample before constructing the message in place via
504/// [`encode_uint8_ipc_header`].
505pub fn uint8_ipc_len(data_len: usize) -> eyre::Result<usize> {
506    Ok(uint8_layout(data_len)?.total)
507}
508
509/// Write a complete no-null `UInt8` IPC stream into `dst` **except the data
510/// region**, and return the byte offset at which the caller must write the
511/// `data_len` data bytes (the buffer-protocol "construct in place" path).
512///
513/// `dst.len()` must equal [`uint8_ipc_len(data_len)`](uint8_ipc_len). After the
514/// caller fills `dst[offset..offset + data_len]`, `dst` is a valid IPC stream
515/// that decodes to the user's bytes as a `UInt8Array` — with zero payload
516/// copies.
517pub fn encode_uint8_ipc_header(dst: &mut [u8], data_len: usize) -> eyre::Result<usize> {
518    let layout = uint8_layout(data_len)?;
519    if dst.len() != layout.total {
520        bail!(
521            "destination size {} does not match required UInt8 IPC length {}",
522            dst.len(),
523            layout.total
524        );
525    }
526    let mut at = 0;
527    at += write_framed_message(dst, at, &layout.schema_message);
528    at += write_framed_message(dst, at, &layout.record_batch_message);
529    let body_start = at;
530
531    // Validity: all-ones bitmap, then padding to the data buffer.
532    dst[body_start..body_start + layout.validity_len].fill(0xff);
533    dst[body_start + layout.validity_len..body_start + layout.validity_padded].fill(0);
534
535    // The data region [data_offset .. data_offset + data_len] is left for the
536    // caller. Zero only its trailing alignment padding.
537    let data_end = layout.data_offset + data_len;
538    let body_end = body_start + layout.body_len;
539    dst[data_end..body_end].fill(0);
540
541    // End-of-stream marker.
542    dst[body_end..body_end + 4].copy_from_slice(&CONTINUATION_MARKER);
543    dst[body_end + 4..body_end + 8].copy_from_slice(&0i32.to_le_bytes());
544    debug_assert_eq!(body_end + PREFIX_LEN, layout.total);
545
546    Ok(layout.data_offset)
547}
548
549/// Length of the leading schema-message block of a full IPC stream produced by
550/// [`encode_ipc_into`] (so a receiver can split it into the schema prefix and
551/// the record-batch+body), or `None` if `stream` is not a framed IPC message.
552pub fn schema_block_len(stream: &[u8]) -> Option<usize> {
553    if stream.len() < PREFIX_LEN || stream[0..4] != CONTINUATION_MARKER {
554        return None;
555    }
556    let metadata_len = i32::from_le_bytes(stream[4..8].try_into().ok()?);
557    let block = PREFIX_LEN.checked_add(usize::try_from(metadata_len).ok()?)?;
558    (block <= stream.len()).then_some(block)
559}
560
561/// Schema identity of a full IPC stream: the FNV-1a hash of its leading schema
562/// block, plus the block itself. This pairing IS the schema-once wire contract
563/// — the producer (`publish_schema_once`), the node receive path (in-band
564/// priming), and the daemon's `dora topic` rebuild all derive the hash from
565/// exactly these bytes via this function; independent re-derivations could
566/// drift and silently break batch↔schema matching.
567pub fn schema_block_and_hash(stream: &[u8]) -> Option<(u64, &[u8])> {
568    let block = schema_block_len(stream)?;
569    let schema = stream.get(..block)?;
570    Some((dora_message::metadata::fnv1a(schema), schema))
571}
572
573/// Given a full IPC stream, return the schema-less record-batch slice —
574/// everything after the schema block, **including** the trailing 8-byte
575/// end-of-stream marker. The marker is what lets the receiver's persistent
576/// decoder flush a 0-row batch (whose body is empty); a non-empty batch never
577/// reads it. `None` if `stream` is malformed. See [`decode_one_batch`].
578pub fn batch_slice(stream: &[u8]) -> Option<&[u8]> {
579    let block = schema_block_len(stream)?;
580    // The stream is [schema block][record-batch block][body][EOS(8)]; keep
581    // everything from the record-batch block onward.
582    (stream.len() >= block + PREFIX_LEN).then(|| &stream[block..])
583}
584
585/// Maximum number of distinct schemas an [`InputDecoder`] retains for local
586/// re-priming. Bounds the memory a misbehaving producer (rotating through many
587/// schemas on one output) can pin in a receiver; beyond it the oldest schema is
588/// evicted and its batches drop until it is re-installed.
589const MAX_RETAINED_SCHEMAS: usize = 8;
590
591/// Per-input receive state for the schema-once zenoh path: one persistent
592/// [`StreamDecoder`](arrow::ipc::reader::StreamDecoder) primed from the schema
593/// published on the output's `@schema` subtopic (or in-band, from the schema
594/// block of a full self-describing stream on the data topic), then reused to
595/// decode the schema-less batch messages that flow on the data topic.
596pub struct InputDecoder {
597    /// Live decoder, primed with the schema for [`schema_hash`](Self::schema_hash).
598    decoder: arrow::ipc::reader::StreamDecoder,
599    /// Hash of the schema the live `decoder` is currently primed with.
600    schema_hash: Option<u64>,
601    /// The framed schema messages installed via [`set_schema`](Self::set_schema),
602    /// keyed by hash (most-recent last, bounded by [`MAX_RETAINED_SCHEMAS`]).
603    /// Retained so the decoder can be re-primed locally (no network round-trip)
604    /// after a failed batch decode soft-resets the live decoder, or when batches
605    /// reference a schema seen earlier (e.g. after in-band priming from a
606    /// different schema's full stream re-primed the live decoder in between).
607    schemas: Vec<(u64, ArrowBuffer)>,
608}
609
610impl Default for InputDecoder {
611    fn default() -> Self {
612        Self::new()
613    }
614}
615
616impl InputDecoder {
617    /// Create an unprimed decoder. It decodes nothing until a schema is
618    /// installed via [`set_schema`](Self::set_schema) (delivered from the output's
619    /// `@schema` subtopic).
620    pub fn new() -> Self {
621        Self {
622            decoder: arrow::ipc::reader::StreamDecoder::new(),
623            schema_hash: None,
624            schemas: Vec::new(),
625        }
626    }
627
628    /// Forget all state — the primed decoder AND the retained schemas. The next
629    /// schema message must re-establish priming. Used on producer restart and
630    /// for lock-poison recovery.
631    pub fn reset(&mut self) {
632        self.decoder = arrow::ipc::reader::StreamDecoder::new();
633        self.schema_hash = None;
634        self.schemas.clear();
635    }
636
637    /// Whether a schema with this hash is already available — live or retained.
638    /// The in-band priming path skips the schema-block copy and the eager
639    /// re-prime for known schemas; [`decode_batch`](Self::decode_batch)
640    /// re-primes lazily from the retained set when a batch actually needs one,
641    /// so eagerly re-priming a retained schema would only churn the live
642    /// decoder (e.g. a large full-stream message of schema B clobbering the
643    /// live prime of schema A between A's schema-less batches).
644    pub fn knows_schema(&self, hash: u64) -> bool {
645        self.schema_hash == Some(hash) || self.schemas.iter().any(|(h, _)| *h == hash)
646    }
647
648    /// Install the schema for `hash` from a framed IPC **schema message** (the
649    /// payload published on the `@schema` subtopic, or the schema block of a
650    /// full stream received in-band): prime a fresh decoder with it and retain
651    /// the bytes for later local re-priming. A no-op when the live decoder is
652    /// already primed with `hash`.
653    pub fn set_schema(&mut self, hash: u64, schema: ArrowBuffer) -> eyre::Result<()> {
654        check_ipc_size(schema.len())?;
655        if self.schema_hash == Some(hash) {
656            return Ok(());
657        }
658        self.prime(hash, schema)
659    }
660
661    /// Prime a fresh `StreamDecoder` with `schema` and record it as the current
662    /// and retained schema for `hash`.
663    fn prime(&mut self, hash: u64, schema: ArrowBuffer) -> eyre::Result<()> {
664        // Create a fresh decoder before priming. In Arrow 59+, even consuming a
665        // schema message can leave the decoder close to terminal state when the
666        // EOS marker is encountered. A fresh decoder ensures we start clean.
667        let mut decoder = arrow::ipc::reader::StreamDecoder::new();
668        prime_with_schema(&mut decoder, schema.clone())?;
669        self.decoder = decoder;
670        self.schema_hash = Some(hash);
671        self.schemas.retain(|(h, _)| *h != hash);
672        self.schemas.push((hash, schema));
673        if self.schemas.len() > MAX_RETAINED_SCHEMAS {
674            self.schemas.remove(0);
675        }
676        Ok(())
677    }
678
679    /// Decode a schema-less batch message against the decoder primed for `hash`.
680    ///
681    /// If the live decoder is not primed for `hash` but a matching schema was
682    /// previously installed (e.g. after a soft reset), it re-primes from the
683    /// retained bytes first. Returns `Ok(None)` when no schema for `hash` is
684    /// known yet — the caller drops the message, which is fine on the lossy
685    /// `CongestionControl::Drop` data plane (the `@schema` history query, the
686    /// next schema publish, or the producer's periodic full-stream refresh —
687    /// which primes in-band — will prime it).
688    pub fn decode_batch(
689        &mut self,
690        buffer: ArrowBuffer,
691        hash: u64,
692    ) -> eyre::Result<Option<arrow::array::ArrayData>> {
693        check_ipc_size(buffer.len())?;
694        if self.schema_hash != Some(hash) {
695            match self.schemas.iter().find(|(h, _)| *h == hash) {
696                Some((_, schema)) => {
697                    let schema = schema.clone();
698                    self.prime(hash, schema)?;
699                }
700                // No schema for this hash known yet — drop (lossy plane).
701                None => return Ok(None),
702            }
703        }
704        match decode_one_batch(&mut self.decoder, buffer) {
705            Ok(array) => {
706                // Arrow 59+ reaches a terminal state after decoding each message's
707                // EOS marker, rejecting further input. Reset the decoder so the next
708                // batch (another independent IPC stream) can be decoded against the
709                // same schema. Clearing schema_hash forces the next batch to re-prime
710                // from the retained schemas (instant lookup, no network round-trip).
711                self.decoder = arrow::ipc::reader::StreamDecoder::new();
712                self.schema_hash = None;
713                Ok(Some(array))
714            }
715            Err(e) => {
716                // A failed batch decode (truncated/corrupt payload — zenoh
717                // tail-loss or a malicious peer) leaves the persistent decoder
718                // mid-message. Soft-reset the LIVE decoder so the next batch is
719                // not fed into that poisoned state (which would misinterpret it
720                // against this message's stale buffers and deliver corrupt data).
721                // The retained schema is kept, so the next batch re-primes
722                // locally — instant recovery, with no wait for a schema refresh.
723                self.decoder = arrow::ipc::reader::StreamDecoder::new();
724                self.schema_hash = None;
725                Err(e)
726            }
727        }
728    }
729}
730
731/// Reject an IPC payload larger than [`super::MAX_IPC_BYTES`] before decoding.
732/// Defense-in-depth against an oversized peer-controlled zenoh payload, mirroring
733/// the guard in [`decode_arrow_ipc_zero_copy`](super::decode_arrow_ipc_zero_copy)
734/// (the persistent-decoder paths receive the same untrusted bytes).
735fn check_ipc_size(len: usize) -> eyre::Result<()> {
736    if len > super::MAX_IPC_BYTES {
737        bail!(
738            "Arrow IPC payload too large: {len} bytes (max {})",
739            super::MAX_IPC_BYTES
740        );
741    }
742    Ok(())
743}
744
745/// Feed a schema message to `decoder`; it must yield no batch.
746fn prime_with_schema(
747    decoder: &mut arrow::ipc::reader::StreamDecoder,
748    mut buffer: ArrowBuffer,
749) -> eyre::Result<()> {
750    while !buffer.is_empty() {
751        let before = buffer.len();
752        if decoder
753            .decode(&mut buffer)
754            .map_err(|e| eyre!("failed to decode IPC schema message: {e}"))?
755            .is_some()
756        {
757            bail!("expected a schema message but got a record batch");
758        }
759        // Guard against a crafted/truncated payload that decodes to no batch
760        // without consuming bytes — otherwise this loop spins forever.
761        if buffer.len() == before {
762            bail!("IPC schema decoder made no progress on a partial/corrupt message");
763        }
764    }
765    Ok(())
766}
767
768/// Feed a record-batch message to `decoder` and return the single decoded array.
769///
770/// The schema-less batch is terminated by an 8-byte end-of-stream marker (see
771/// [`encode_batch_into`]/[`batch_slice`]). For a non-empty batch the body bytes
772/// flush it on their own; the marker matters for a **0-row** batch, whose
773/// zero-length body arrow's `StreamDecoder` only emits when polled with more
774/// input — without the trailing marker an empty array is silently dropped
775/// (PR #2366). The decoder yields the pending batch *before* consuming the
776/// marker's zero length, so it never reaches its terminal state and stays
777/// usable for the next batch.
778fn decode_one_batch(
779    decoder: &mut arrow::ipc::reader::StreamDecoder,
780    mut buffer: ArrowBuffer,
781) -> eyre::Result<arrow::array::ArrayData> {
782    while !buffer.is_empty() {
783        let before = buffer.len();
784        if let Some(batch) = decoder
785            .decode(&mut buffer)
786            .map_err(|e| eyre!("failed to decode IPC record batch: {e}"))?
787        {
788            if batch.num_columns() != 1 {
789                bail!(
790                    "expected 1 column in IPC record batch, got {}",
791                    batch.num_columns()
792                );
793            }
794            return Ok(batch.column(0).to_data());
795        }
796        // Guard against a crafted/truncated payload that decodes to no batch
797        // without consuming bytes — otherwise this loop spins forever.
798        if buffer.len() == before {
799            bail!("IPC batch decoder made no progress on a partial/corrupt message");
800        }
801    }
802    bail!("IPC batch message yielded no record batch")
803}
804
805#[cfg(test)]
806mod tests {
807    use super::*;
808    use crate::arrow_utils::decode_arrow_ipc_zero_copy;
809    use arrow::array::{
810        Array, ArrayRef, BooleanArray, FixedSizeBinaryArray, Float32Array, Int32Array,
811        LargeStringArray, ListArray, NullArray, StringArray, StructArray, UInt8Array, UInt64Array,
812    };
813    use arrow::buffer::Buffer;
814    use arrow::ipc::reader::{StreamDecoder, StreamReader};
815    use arrow_schema::{DataType, Field};
816    use std::io::Cursor;
817    use std::sync::Arc;
818
819    /// Encode via the fast path into a fresh `Vec` sized by `ipc_fast_path_len`.
820    fn fast_encode(array: &ArrayData) -> Vec<u8> {
821        let len = ipc_fast_path_len(array).expect("array should be fast-path eligible");
822        let mut buf = vec![0u8; len];
823        encode_ipc_into(array, &mut buf).expect("fast-path encode");
824        buf
825    }
826
827    /// Decode a complete IPC stream with the OFFICIAL arrow `StreamReader` (the
828    /// oracle for wire correctness — independent of our decoder).
829    fn read_official(bytes: &[u8]) -> ArrayData {
830        let mut reader = StreamReader::try_new(Cursor::new(bytes), None).expect("open IPC stream");
831        let batch = reader
832            .next()
833            .expect("one batch")
834            .expect("batch decodes via official reader");
835        assert_eq!(batch.num_columns(), 1);
836        batch.column(0).to_data()
837    }
838
839    /// Copy `bytes` into a 128-byte-aligned Arrow `Buffer`, mirroring how the
840    /// receive path backs SHM payloads — the precondition for zero-copy decode.
841    fn aligned_buffer(bytes: &[u8]) -> (Buffer, usize, usize) {
842        use aligned_vec::{AVec, ConstAlign};
843        use std::ptr::NonNull;
844        let mut aligned: AVec<u8, ConstAlign<128>> = AVec::__from_elem(128, 0, bytes.len());
845        aligned.copy_from_slice(bytes);
846        let base = aligned.as_ptr() as usize;
847        let len = aligned.len();
848        let ptr = NonNull::new(aligned.as_ptr() as *mut u8).unwrap();
849        // SAFETY: ptr/len describe `aligned`; the Arc keeps it alive.
850        let buffer =
851            unsafe { Buffer::from_custom_allocation(ptr, len, std::sync::Arc::new(aligned)) };
852        (buffer, base, len)
853    }
854
855    /// The core assertion: the fast-path stream (1) decodes via the official
856    /// reader to an equal array, and (2) `ipc_fast_path_len` matches the bytes
857    /// written and the stream the reader fully consumes.
858    fn assert_fast_roundtrip(array: &ArrayData) {
859        let encoded = fast_encode(array);
860        let decoded = read_official(&encoded);
861        assert_eq!(array, &decoded, "fast-path stream must decode to the input");
862        // Our own zero-copy decoder must agree too.
863        let (buffer, _, _) = aligned_buffer(&encoded);
864        let zc = decode_arrow_ipc_zero_copy(buffer).expect("zero-copy decode");
865        assert_eq!(array, &zc, "zero-copy decode must equal the input");
866    }
867
868    /// Encode `array` as a schema-less batch message (the fast path's
869    /// `batch_slice` equivalent), as shipped on the schema-once data plane.
870    fn batch_bytes(array: &ArrayData) -> Vec<u8> {
871        let len = batch_fast_path_len(array).unwrap();
872        let mut buf = vec![0u8; len];
873        encode_batch_into(array, &mut buf).unwrap();
874        buf
875    }
876
877    fn batch_buf(array: &ArrayData) -> Buffer {
878        Buffer::from_vec(batch_bytes(array))
879    }
880
881    /// The schema-once receive contract: an `InputDecoder` is primed by a schema
882    /// message (as delivered from the `@schema` subtopic), then decodes the
883    /// schema-less batches that follow, and drops a batch whose hash it isn't
884    /// primed for.
885    #[test]
886    fn input_decoder_schema_then_batches() {
887        let f32_schema = || Buffer::from_vec(encode_schema_message(&DataType::Float32).unwrap());
888
889        let mut dec = InputDecoder::new();
890
891        // A batch arriving before any schema is dropped (not primed).
892        let early = Float32Array::from(vec![9.0]).into_data();
893        assert!(dec.decode_batch(batch_buf(&early), 7).unwrap().is_none());
894
895        // Installing the schema primes the decoder; following batches decode.
896        dec.set_schema(7, f32_schema()).unwrap();
897        for vals in [vec![1.0f32, 2.0, 3.0], vec![4.0], vec![5.0, 6.0, 7.0]] {
898            let array = Float32Array::from(vals).into_data();
899            assert_eq!(
900                dec.decode_batch(batch_buf(&array), 7).unwrap().unwrap(),
901                array
902            );
903        }
904
905        // A batch tagged with a hash the decoder isn't primed for is dropped.
906        let other = Float32Array::from(vec![8.0]).into_data();
907        assert!(dec.decode_batch(batch_buf(&other), 99).unwrap().is_none());
908
909        // Installing a schema under the new hash primes it; its batches decode.
910        dec.set_schema(99, f32_schema()).unwrap();
911        let after = Float32Array::from(vec![10.0, 11.0]).into_data();
912        assert_eq!(
913            dec.decode_batch(batch_buf(&after), 99).unwrap().unwrap(),
914            after
915        );
916    }
917
918    /// Priming a new schema must not forget previously seen ones: batches for a
919    /// schema installed earlier still decode after the live decoder was re-primed
920    /// with a different schema in between (re-primed locally from the retained
921    /// set). Without this, in-band priming from a full stream (e.g. a large
922    /// message or a schema-change message on the same output) would clobber the
923    /// schema that later schema-less batches reference, silently dropping them.
924    #[test]
925    fn input_decoder_retains_multiple_schemas() {
926        let schema_msg = |dt: &DataType| Buffer::from_vec(encode_schema_message(dt).unwrap());
927
928        let mut dec = InputDecoder::new();
929        dec.set_schema(1, schema_msg(&DataType::Float32)).unwrap();
930        dec.set_schema(2, schema_msg(&DataType::Int32)).unwrap();
931
932        // Live decoder is primed for hash 2 …
933        let ints = Int32Array::from(vec![1, 2, 3]).into_data();
934        assert_eq!(
935            dec.decode_batch(batch_buf(&ints), 2).unwrap().unwrap(),
936            ints
937        );
938
939        // … but a batch for hash 1 must still decode (retained schema).
940        let floats = Float32Array::from(vec![4.0, 5.0]).into_data();
941        assert_eq!(
942            dec.decode_batch(batch_buf(&floats), 1).unwrap().unwrap(),
943            floats,
944            "a schema installed earlier must be retained across later primes"
945        );
946
947        // And switching back again also works.
948        let more_ints = Int32Array::from(vec![6]).into_data();
949        assert_eq!(
950            dec.decode_batch(batch_buf(&more_ints), 2).unwrap().unwrap(),
951            more_ints
952        );
953    }
954
955    /// The retained-schema set is bounded: schemas beyond the cap evict the
956    /// oldest, whose batches then drop until it is re-installed.
957    /// Test that schema-less batches can be decoded repeatedly after the
958    /// decoder is soft-reset between batches (Arrow 59+ terminal state fix).
959    /// This is the regression test for PR #2445/#2366 interaction with Arrow 59.
960    #[test]
961    fn input_decoder_handles_sequential_batches_arrow_59_terminal_state() {
962        let f32_schema = || Buffer::from_vec(encode_schema_message(&DataType::Float32).unwrap());
963
964        let mut dec = InputDecoder::new();
965        // Prime with a schema.
966        dec.set_schema(7, f32_schema()).unwrap();
967
968        // Decode 5 schema-less batches sequentially. Each one should decode
969        // correctly despite the decoder being reset between batches (soft-reset
970        // for Arrow 59 terminal state handling).
971        let batches = vec![
972            Float32Array::from(vec![1.0, 2.0]).into_data(),
973            Float32Array::from(vec![3.0]).into_data(),
974            Float32Array::from(vec![4.0, 5.0, 6.0]).into_data(),
975            Float32Array::from(vec![7.0, 8.0]).into_data(),
976            Float32Array::from(vec![9.0, 10.0, 11.0, 12.0]).into_data(),
977        ];
978
979        for (i, batch) in batches.iter().enumerate() {
980            let result = dec.decode_batch(batch_buf(batch), 7);
981            assert!(
982                result.is_ok(),
983                "batch {} decode failed: {:?}",
984                i,
985                result.err()
986            );
987            let decoded = result.unwrap().unwrap();
988            assert_eq!(
989                &decoded, batch,
990                "batch {} mismatch: expected {:?}, got {:?}",
991                i, batch, decoded
992            );
993        }
994    }
995
996    #[test]
997    fn input_decoder_evicts_oldest_schema_beyond_cap() {
998        let f32_schema = || Buffer::from_vec(encode_schema_message(&DataType::Float32).unwrap());
999
1000        let mut dec = InputDecoder::new();
1001        // Install cap + 1 distinct hashes (same schema bytes — only the hash
1002        // keys retention); hash 0 must be evicted, the rest retained.
1003        for hash in 0..=(MAX_RETAINED_SCHEMAS as u64) {
1004            dec.set_schema(hash, f32_schema()).unwrap();
1005        }
1006        let array = Float32Array::from(vec![1.0]).into_data();
1007        assert!(
1008            dec.decode_batch(batch_buf(&array), 0).unwrap().is_none(),
1009            "the oldest schema must be evicted beyond the cap"
1010        );
1011        assert_eq!(
1012            dec.decode_batch(batch_buf(&array), 1).unwrap().unwrap(),
1013            array,
1014            "schemas within the cap must be retained"
1015        );
1016    }
1017
1018    /// A failed (truncated) batch decode must soft-reset the persistent decoder
1019    /// so the next valid batch is not misinterpreted against the truncated message's
1020    /// stale state — and, because the schema is retained, that next batch
1021    /// re-primes locally and decodes (instant recovery, not drop-until-refresh).
1022    #[test]
1023    fn decode_batch_resets_on_error_then_reprimes() {
1024        let mut dec = InputDecoder::new();
1025        dec.set_schema(
1026            7,
1027            Buffer::from_vec(encode_schema_message(&DataType::Float32).unwrap()),
1028        )
1029        .unwrap();
1030
1031        // A valid batch decodes.
1032        let good = Float32Array::from(vec![4.0, 5.0]).into_data();
1033        assert_eq!(
1034            dec.decode_batch(batch_buf(&good), 7).unwrap().unwrap(),
1035            good
1036        );
1037
1038        // A truncated batch errors and soft-resets the live decoder. Cut the
1039        // buffer in half so the record-batch header survives but its declared
1040        // body is incomplete (a tail loss) — dropping only the trailing 8-byte
1041        // EOS marker would leave a still-valid batch.
1042        let dropped = Float32Array::from(vec![6.0, 7.0, 8.0, 9.0]).into_data();
1043        let mut truncated = batch_bytes(&dropped);
1044        truncated.truncate(truncated.len() / 2);
1045        assert!(dec.decode_batch(Buffer::from_vec(truncated), 7).is_err());
1046
1047        // The next valid batch (same hash) re-primes from the retained schema and
1048        // decodes correctly — not dropped, not corrupted.
1049        let after = Float32Array::from(vec![10.0, 11.0]).into_data();
1050        assert_eq!(
1051            dec.decode_batch(batch_buf(&after), 7).unwrap().unwrap(),
1052            after,
1053            "after a failed batch the decoder must re-prime from the retained schema"
1054        );
1055
1056        // `reset()` forgets the schema too; batches then drop until re-installed.
1057        dec.reset();
1058        let final_batch = Float32Array::from(vec![12.0]).into_data();
1059        assert!(
1060            dec.decode_batch(batch_buf(&final_batch), 7)
1061                .unwrap()
1062                .is_none(),
1063            "after a full reset the decoder must drop until a schema is re-installed"
1064        );
1065    }
1066
1067    /// Schema-once batch path for a *fallback* (dictionary) type. Dictionary
1068    /// arrays are not fast-path eligible, so the full stream comes from the
1069    /// official writer as `[schema][dictionary batch][record batch][EOS]`. The
1070    /// schema-once optimization still applies to small messages (node/mod.rs),
1071    /// shipping `batch_slice` = `[dictionary batch][record batch]` against a
1072    /// decoder primed from the schema subtopic — i.e. a *replacement* dictionary
1073    /// on every message. The `InputDecoder` tests otherwise use only `Float32`; a
1074    /// failure here is silent intermittent input loss, so it needs an explicit
1075    /// oracle.
1076    #[test]
1077    fn input_decoder_dictionary_fallback_batch_sequence() {
1078        use arrow::array::DictionaryArray;
1079        use arrow::datatypes::Int32Type;
1080
1081        // Build a Dictionary<Int32, Utf8> from words, distinct values in
1082        // first-seen order (deterministic, no iterator-trait ambiguity).
1083        fn dict(words: &[&str]) -> ArrayData {
1084            let mut values: Vec<&str> = Vec::new();
1085            let mut keys: Vec<i32> = Vec::new();
1086            for w in words {
1087                let idx = values.iter().position(|v| v == w).unwrap_or_else(|| {
1088                    values.push(*w);
1089                    values.len() - 1
1090                });
1091                keys.push(idx as i32);
1092            }
1093            DictionaryArray::<Int32Type>::try_new(
1094                Int32Array::from(keys),
1095                Arc::new(StringArray::from(values)),
1096            )
1097            .unwrap()
1098            .into_data()
1099        }
1100
1101        let first = dict(&["a", "b", "a", "c", "b"]);
1102        // Confirm this really exercises the fallback (not the fast path).
1103        assert!(
1104            ipc_fast_path_len(&first).is_none(),
1105            "dictionary must route to the official-writer fallback"
1106        );
1107
1108        // Prime from the schema block of the dictionary stream — exactly the
1109        // bytes the producer publishes on the `@schema` subtopic.
1110        let mut dec = InputDecoder::new();
1111        let full0 = encode_ipc_to_vec(&first).unwrap();
1112        let block = schema_block_len(&full0).unwrap();
1113        dec.set_schema(1, Buffer::from(&full0[..block])).unwrap();
1114
1115        // Every message (including the first) ships only the schema-less batch
1116        // slice (which for a dictionary type carries a replacement dictionary
1117        // batch + record batch) against the primed decoder. Each must decode to
1118        // its own input.
1119        let slice0 = batch_slice(&full0).expect("fallback stream is a valid IPC stream");
1120        assert_eq!(
1121            dec.decode_batch(Buffer::from(slice0), 1)
1122                .unwrap()
1123                .expect("first batch decodes against the primed decoder"),
1124            first
1125        );
1126        for words in [
1127            ["x", "y", "x", "z"].as_slice(),
1128            ["b", "b"].as_slice(),
1129            ["new", "values", "entirely"].as_slice(),
1130        ] {
1131            let arr = dict(words);
1132            let full = encode_ipc_to_vec(&arr).unwrap();
1133            let slice = batch_slice(&full).expect("fallback stream is a valid IPC stream");
1134            let got = dec
1135                .decode_batch(Buffer::from(slice), 1)
1136                .unwrap()
1137                .expect("batch must decode against the primed decoder");
1138            assert_eq!(
1139                got, arr,
1140                "replacement-dictionary batch must decode correctly"
1141            );
1142        }
1143    }
1144
1145    /// The W3 schema-once contract: prime one persistent `StreamDecoder` with a
1146    /// single schema message, then decode a *sequence* of schema-less batch
1147    /// messages against it. (A fresh full stream per message can't do this — its
1148    /// EOS terminates the decoder.)
1149    #[test]
1150    fn schema_primed_decoder_decodes_batch_sequence() {
1151        let schema = encode_schema_message(&DataType::Float32).unwrap();
1152        let mut decoder = StreamDecoder::new();
1153
1154        // Prime: feeding the schema message yields no batch.
1155        let mut sbuf = Buffer::from_vec(schema);
1156        while !sbuf.is_empty() {
1157            assert!(
1158                decoder.decode(&mut sbuf).unwrap().is_none(),
1159                "schema message must not yield a batch"
1160            );
1161        }
1162
1163        for vals in [vec![1.0f32, 2.0, 3.0], vec![4.0, 5.0], vec![6.0]] {
1164            let array = Float32Array::from(vals).into_data();
1165            let len = batch_fast_path_len(&array).unwrap();
1166            let mut buf = vec![0u8; len];
1167            encode_batch_into(&array, &mut buf).unwrap();
1168
1169            let mut bbuf = Buffer::from_vec(buf);
1170            let mut got = None;
1171            while !bbuf.is_empty() {
1172                if let Some(b) = decoder.decode(&mut bbuf).unwrap() {
1173                    got = Some(b);
1174                    break;
1175                }
1176            }
1177            assert_eq!(
1178                got.expect("batch message must decode against the primed decoder")
1179                    .column(0)
1180                    .to_data(),
1181                array
1182            );
1183        }
1184    }
1185
1186    /// `encode_uint8_ipc_header` + caller-filled data region produces a valid
1187    /// IPC stream identical to `encode_ipc_into` of the equivalent UInt8Array —
1188    /// the zero-copy "construct in place" path for `send_output_raw`/Python.
1189    #[test]
1190    fn uint8_ipc_header_constructs_in_place() {
1191        for data_len in [0usize, 1, 7, 8, 9, 1000] {
1192            let bytes: Vec<u8> = (0..data_len).map(|i| (i % 251) as u8).collect();
1193            let total = uint8_ipc_len(data_len).unwrap();
1194            let mut dst = vec![0u8; total];
1195            let offset = encode_uint8_ipc_header(&mut dst, data_len).unwrap();
1196            // The caller writes the data in place (no copy in real use).
1197            dst[offset..offset + data_len].copy_from_slice(&bytes);
1198
1199            // Decodes to the original bytes as a no-null UInt8Array.
1200            let arr = arrow::array::make_array(read_official(&dst));
1201            let u8 = arr.as_any().downcast_ref::<UInt8Array>().unwrap();
1202            assert_eq!(u8.values(), bytes.as_slice(), "len {data_len}");
1203            assert_eq!(u8.null_count(), 0);
1204
1205            // Byte-identical to encoding the equivalent array directly.
1206            let array = UInt8Array::from(bytes).into_data();
1207            assert_eq!(dst, fast_encode(&array), "len {data_len}");
1208        }
1209    }
1210
1211    #[test]
1212    fn roundtrip_primitive_no_nulls() {
1213        let array = Float32Array::from((0..1000).map(|i| i as f32).collect::<Vec<_>>()).into_data();
1214        assert_fast_roundtrip(&array);
1215    }
1216
1217    #[test]
1218    fn roundtrip_primitive_with_nulls() {
1219        let array = UInt64Array::from(vec![Some(1), None, Some(3), None, Some(5)]).into_data();
1220        assert_fast_roundtrip(&array);
1221    }
1222
1223    #[test]
1224    fn roundtrip_empty_primitive() {
1225        let array = Int32Array::from(Vec::<i32>::new()).into_data();
1226        assert_fast_roundtrip(&array);
1227    }
1228
1229    /// A 0-row array sent through the `send_output_raw` UInt8-header path decodes
1230    /// back to a 0-row `UInt8` array via the full-stream zero-copy decoder.
1231    /// (The full stream keeps its end-of-stream marker, so this path was never
1232    /// broken — this pins it as a baseline alongside the schema-once regression
1233    /// below.)
1234    #[test]
1235    fn uint8_header_zero_len_full_stream_roundtrip() {
1236        let len = uint8_ipc_len(0).unwrap();
1237        let mut buf = vec![0u8; len];
1238        let off = encode_uint8_ipc_header(&mut buf, 0).unwrap();
1239        // Empty data region: it sits immediately before the 8-byte EOS marker.
1240        assert_eq!(off, len - PREFIX_LEN);
1241        assert_eq!(read_official(&buf).len(), 0);
1242        let (buffer, _, _) = aligned_buffer(&buf);
1243        let decoded = decode_arrow_ipc_zero_copy(buffer).unwrap();
1244        assert_eq!(decoded.data_type(), &DataType::UInt8);
1245        assert_eq!(decoded.len(), 0);
1246    }
1247
1248    /// The schema-once receive path must decode a 0-row (empty) batch, not drop
1249    /// it. A schema-less batch carries no trailing end-of-stream marker, and
1250    /// arrow's `StreamDecoder` only emits a zero-length body on the poll *after*
1251    /// the header — so `decode_one_batch` must poll once more once its input
1252    /// drains. Regression guard for the empty-array drop reported on PR #2366.
1253    #[test]
1254    fn schema_once_zero_len_batch_roundtrip() {
1255        let schema = || Buffer::from_vec(encode_schema_message(&DataType::UInt8).unwrap());
1256        let batch = |vals: &[u8]| {
1257            let a = UInt8Array::from(vals.to_vec()).into_data();
1258            let len = batch_fast_path_len(&a).unwrap();
1259            let mut b = vec![0u8; len];
1260            encode_batch_into(&a, &mut b).unwrap();
1261            Buffer::from_vec(b)
1262        };
1263
1264        let mut dec = InputDecoder::new();
1265        dec.set_schema(1, schema()).unwrap();
1266
1267        // The empty batch decodes to a 0-row UInt8 array, not `Ok(None)`/error.
1268        let decoded = dec
1269            .decode_batch(batch(&[]), 1)
1270            .unwrap()
1271            .expect("0-row batch must decode, not drop");
1272        assert_eq!(decoded.data_type(), &DataType::UInt8);
1273        assert_eq!(decoded.len(), 0);
1274
1275        // The persistent decoder stays usable across mixed empty/non-empty
1276        // batches (flushing the empty body must not wedge or terminate it).
1277        for vals in [vec![1u8, 2, 3], vec![], vec![9u8], vec![]] {
1278            let d = dec.decode_batch(batch(&vals), 1).unwrap().unwrap();
1279            assert_eq!(d.len(), vals.len());
1280            assert_eq!(d.data_type(), &DataType::UInt8);
1281        }
1282    }
1283
1284    /// Same regression, but exercising the exact production producer path: the
1285    /// full stream is built by the UInt8 header encoder, the schema is extracted
1286    /// via [`schema_block_len`] and the schema-less batch via [`batch_slice`]
1287    /// (as `zenoh_publish` does), then decoded by the per-input [`InputDecoder`].
1288    #[test]
1289    fn schema_once_zero_len_via_batch_slice_roundtrip() {
1290        let total = uint8_ipc_len(0).unwrap();
1291        let mut full = vec![0u8; total];
1292        encode_uint8_ipc_header(&mut full, 0).unwrap();
1293
1294        let sblock = schema_block_len(&full).unwrap();
1295        let schema = Buffer::from(&full[..sblock]);
1296        let batch = batch_slice(&full).expect("batch slice of a valid stream");
1297
1298        let mut dec = InputDecoder::new();
1299        dec.set_schema(7, schema).unwrap();
1300        let decoded = dec
1301            .decode_batch(Buffer::from(batch), 7)
1302            .unwrap()
1303            .expect("0-row batch via batch_slice must decode, not drop");
1304        assert_eq!(decoded.data_type(), &DataType::UInt8);
1305        assert_eq!(decoded.len(), 0);
1306    }
1307
1308    /// The daemon's `dora topic` debug path rebuilds a schema-once batch into a
1309    /// full self-describing stream by concatenating the cached `@schema` block
1310    /// with the schema-less batch. That only works because
1311    /// `schema_block ++ batch_slice` is byte-identical to the original stream —
1312    /// lock that invariant (and that the result still decodes) here.
1313    #[test]
1314    fn schema_block_plus_batch_slice_reconstructs_full_stream() {
1315        let array = Int32Array::from(vec![1, 2, 3]).into_data();
1316        let full = fast_encode(&array);
1317        let sblock = schema_block_len(&full).unwrap();
1318        let schema = &full[..sblock];
1319        let batch = batch_slice(&full).expect("batch slice of a valid stream");
1320        let rebuilt = [schema, batch].concat();
1321        assert_eq!(
1322            rebuilt, full,
1323            "schema_block ++ batch_slice must equal the original stream"
1324        );
1325        assert_eq!(read_official(&rebuilt), array);
1326    }
1327
1328    #[test]
1329    fn roundtrip_boolean() {
1330        let array =
1331            BooleanArray::from(vec![true, false, true, true, false, false, true]).into_data();
1332        assert_fast_roundtrip(&array);
1333    }
1334
1335    #[test]
1336    fn roundtrip_utf8() {
1337        let array =
1338            StringArray::from(vec![Some("hello"), None, Some(""), Some("world!")]).into_data();
1339        assert_fast_roundtrip(&array);
1340    }
1341
1342    #[test]
1343    fn roundtrip_large_utf8_64bit_offsets() {
1344        let array = LargeStringArray::from(vec!["a", "bb", "ccc"]).into_data();
1345        assert_fast_roundtrip(&array);
1346    }
1347
1348    #[test]
1349    fn roundtrip_fixed_size_binary() {
1350        let values = vec![vec![1u8, 2, 3], vec![4, 5, 6], vec![7, 8, 9]];
1351        let array = FixedSizeBinaryArray::try_from_iter(values.into_iter())
1352            .unwrap()
1353            .into_data();
1354        assert_fast_roundtrip(&array);
1355    }
1356
1357    /// `FixedSizeList` exercises the unique fast-path child-slicing branch
1358    /// (`n = len * value_size`, then `child.slice(0, n)`). No other test covers
1359    /// it — `roundtrip_fixed_size_binary` is `FixedSizeBinary`, a leaf type.
1360    #[test]
1361    fn roundtrip_fixed_size_list() {
1362        use arrow::array::FixedSizeListArray;
1363        let values = Int32Array::from((0..12).collect::<Vec<_>>());
1364        let field = Arc::new(Field::new("item", DataType::Int32, true));
1365        let array = FixedSizeListArray::try_new(field, 3, Arc::new(values), None)
1366            .unwrap()
1367            .into_data();
1368        assert_fast_roundtrip(&array);
1369    }
1370
1371    /// `FixedSizeList` with a list-level validity bitmap (null lists), so the
1372    /// child-slicing branch runs alongside the parent's null buffer.
1373    #[test]
1374    fn roundtrip_fixed_size_list_with_nulls() {
1375        use arrow::array::FixedSizeListArray;
1376        use arrow::buffer::NullBuffer;
1377        let values = Int32Array::from((0..12).collect::<Vec<_>>());
1378        let field = Arc::new(Field::new("item", DataType::Int32, true));
1379        let nulls = NullBuffer::from(vec![true, false, true, true]);
1380        let array = FixedSizeListArray::try_new(field, 3, Arc::new(values), Some(nulls))
1381            .unwrap()
1382            .into_data();
1383        assert_fast_roundtrip(&array);
1384    }
1385
1386    /// `Decimal128` is fast-path (buffer copied verbatim) but otherwise untested.
1387    #[test]
1388    fn roundtrip_decimal128() {
1389        use arrow::array::Decimal128Array;
1390        let array = Decimal128Array::from(vec![Some(12_345i128), None, Some(-9_876), Some(0)])
1391            .with_precision_and_scale(20, 4)
1392            .unwrap()
1393            .into_data();
1394        assert_fast_roundtrip(&array);
1395    }
1396
1397    /// A temporal type (`Timestamp`) — also fast-path-verbatim, also untested.
1398    #[test]
1399    fn roundtrip_timestamp_temporal() {
1400        use arrow::array::TimestampMicrosecondArray;
1401        let array =
1402            TimestampMicrosecondArray::from(vec![Some(1_000_000i64), None, Some(2_500_000)])
1403                .into_data();
1404        assert_fast_roundtrip(&array);
1405    }
1406
1407    #[test]
1408    fn roundtrip_struct_with_multilevel_nulls() {
1409        let array = StructArray::from(vec![
1410            (
1411                Arc::new(Field::new("a", DataType::UInt64, true)),
1412                Arc::new(UInt64Array::from(vec![Some(1), None, Some(3)])) as ArrayRef,
1413            ),
1414            (
1415                Arc::new(Field::new("b", DataType::Utf8, true)),
1416                Arc::new(StringArray::from(vec![Some("x"), Some("yy"), None])) as ArrayRef,
1417            ),
1418        ])
1419        .into_data();
1420        assert_fast_roundtrip(&array);
1421    }
1422
1423    /// Regression: a Struct whose child `ArrayData` is *longer* than the struct
1424    /// (constructible via the low-level builder) must be sliced to the struct's
1425    /// len, otherwise the child field node declares the wrong row count and the
1426    /// stream is un-decodable.
1427    #[test]
1428    fn roundtrip_struct_with_oversized_child() {
1429        use arrow_schema::Fields;
1430        let child = Int32Array::from(vec![10, 20, 30]).into_data(); // len 3
1431        let fields: Fields = vec![Field::new("v", DataType::Int32, false)].into();
1432        let struct_data = ArrayData::builder(DataType::Struct(fields))
1433            .len(2) // shorter than the child
1434            .add_child_data(child)
1435            .build()
1436            .unwrap();
1437
1438        assert!(
1439            ipc_fast_path_len(&struct_data).is_some(),
1440            "struct with an oversized child should stay on the fast path"
1441        );
1442        let decoded = read_official(&fast_encode(&struct_data));
1443        assert_eq!(decoded.len(), 2);
1444        let arr = arrow::array::make_array(decoded);
1445        let sa = arr.as_any().downcast_ref::<StructArray>().unwrap();
1446        let col = sa.column(0).as_any().downcast_ref::<Int32Array>().unwrap();
1447        assert_eq!(
1448            col.values(),
1449            &[10, 20],
1450            "child must be truncated to the struct's len"
1451        );
1452    }
1453
1454    #[test]
1455    fn roundtrip_list_of_primitive() {
1456        let data = vec![
1457            Some(vec![Some(0), Some(1), Some(2)]),
1458            None,
1459            Some(vec![Some(3), None, Some(5)]),
1460            Some(vec![]),
1461        ];
1462        let array =
1463            ListArray::from_iter_primitive::<arrow::datatypes::Int32Type, _, _>(data).into_data();
1464        assert_fast_roundtrip(&array);
1465    }
1466
1467    #[test]
1468    fn roundtrip_nullarray_zero_and_n() {
1469        assert_fast_roundtrip(&NullArray::new(0).into_data());
1470        assert_fast_roundtrip(&NullArray::new(7).into_data());
1471    }
1472
1473    /// Headline zero-copy proof: the fast-path body is 64-aligned, so the strict
1474    /// decoder accepts it without a realigning copy AND the decoded data buffer
1475    /// aliases the input allocation.
1476    #[test]
1477    fn fast_path_decodes_zero_copy() {
1478        let array = UInt64Array::from((0..50_000u64).collect::<Vec<_>>()).into_data();
1479        let encoded = fast_encode(&array);
1480
1481        // (1) strict decoder: errors if any buffer needs realigning.
1482        {
1483            let (mut buffer, _, _) = aligned_buffer(&encoded);
1484            let mut decoder = StreamDecoder::new().with_require_alignment(true);
1485            let mut got = None;
1486            while !buffer.is_empty() {
1487                if let Some(b) = decoder
1488                    .decode(&mut buffer)
1489                    .expect("aligned fast-path stream must decode without realignment")
1490                {
1491                    got = Some(b);
1492                    break;
1493                }
1494            }
1495            assert_eq!(got.unwrap().column(0).to_data(), array);
1496        }
1497
1498        // (2) pointer aliasing: the decoded data buffer lies inside the input.
1499        {
1500            let (buffer, base, len) = aligned_buffer(&encoded);
1501            let decoded = decode_arrow_ipc_zero_copy(buffer).unwrap();
1502            let ptr = decoded.buffers()[0].as_ptr() as usize;
1503            assert!(
1504                ptr >= base && ptr < base + len,
1505                "decoded data buffer at {ptr:#x} is outside input [{base:#x}, {:#x}) — a copy happened",
1506                base + len
1507            );
1508        }
1509    }
1510
1511    #[test]
1512    fn fast_path_len_matches_official_decode() {
1513        // The whole stream is consumed by the official reader (no trailing junk,
1514        // no truncation): `ipc_fast_path_len` is exact.
1515        let array = Float32Array::from(vec![1.0, 2.0, 3.0, 4.0]).into_data();
1516        let encoded = fast_encode(&array);
1517        let mut reader = StreamReader::try_new(Cursor::new(&encoded[..]), None).unwrap();
1518        let _ = reader.next().unwrap().unwrap();
1519        assert!(reader.next().is_none(), "exactly one batch, fully consumed");
1520    }
1521
1522    /// Sliced (non-zero offset) arrays are routed to the fallback, which must
1523    /// still round-trip the *logical* slice (not the parent).
1524    #[test]
1525    fn sliced_array_routes_to_fallback_and_roundtrips() {
1526        let array = UInt64Array::from(vec![10, 20, 30, 40, 50])
1527            .into_data()
1528            .slice(2, 2); // offset 2 -> not fast-path
1529        assert_eq!(array.offset(), 2);
1530        assert!(ipc_fast_path_len(&array).is_none());
1531
1532        let encoded = encode_ipc_to_vec(&array).unwrap();
1533        let decoded = read_official(&encoded);
1534        assert_eq!(array.len(), decoded.len());
1535        let dec = arrow::array::make_array(decoded);
1536        let dec = dec.as_any().downcast_ref::<UInt64Array>().unwrap();
1537        assert_eq!(dec.values(), &[30, 40]);
1538    }
1539
1540    /// A `*View` type must route to the fallback (validate the classifier).
1541    #[test]
1542    fn view_type_routes_to_fallback() {
1543        use arrow::array::StringViewArray;
1544        let array = StringViewArray::from(vec!["a", "bb", "ccc"]).into_data();
1545        assert!(
1546            ipc_fast_path_len(&array).is_none(),
1547            "Utf8View is not fast-path eligible"
1548        );
1549        let encoded = encode_ipc_to_vec(&array).unwrap();
1550        let decoded = read_official(&encoded);
1551        assert_eq!(array, decoded);
1552    }
1553
1554    /// Deterministic fuzz: many shapes through the bidirectional assertion.
1555    #[test]
1556    fn fuzz_roundtrip_many_shapes() {
1557        // Simple LCG so the matrix is varied but reproducible (no rng/time).
1558        let mut state: u64 = 0x1234_5678_9abc_def0;
1559        let mut next = || {
1560            state = state
1561                .wrapping_mul(6364136223846793005)
1562                .wrapping_add(1442695040888963407);
1563            state
1564        };
1565
1566        for _ in 0..200 {
1567            let len = (next() % 64) as usize;
1568            let kind = next() % 6;
1569            let array: ArrayData = match kind {
1570                0 => UInt8Array::from(
1571                    (0..len)
1572                        .map(|i| {
1573                            if next().is_multiple_of(4) {
1574                                None
1575                            } else {
1576                                Some((i as u8).wrapping_add(1))
1577                            }
1578                        })
1579                        .collect::<Vec<_>>(),
1580                )
1581                .into_data(),
1582                1 => Float32Array::from((0..len).map(|i| i as f32 * 0.5).collect::<Vec<_>>())
1583                    .into_data(),
1584                2 => BooleanArray::from(
1585                    (0..len)
1586                        .map(|i| (i + next() as usize).is_multiple_of(2))
1587                        .collect::<Vec<_>>(),
1588                )
1589                .into_data(),
1590                3 => StringArray::from(
1591                    (0..len)
1592                        .map(|i| {
1593                            if next().is_multiple_of(5) {
1594                                None
1595                            } else {
1596                                Some("x".repeat(i % 7))
1597                            }
1598                        })
1599                        .collect::<Vec<_>>(),
1600                )
1601                .into_data(),
1602                4 => Int32Array::from(
1603                    (0..len)
1604                        .map(|i| {
1605                            if next().is_multiple_of(3) {
1606                                None
1607                            } else {
1608                                Some(i as i32 - 10)
1609                            }
1610                        })
1611                        .collect::<Vec<_>>(),
1612                )
1613                .into_data(),
1614                _ => StructArray::from(vec![(
1615                    Arc::new(Field::new("v", DataType::Int32, true)),
1616                    Arc::new(Int32Array::from(
1617                        (0..len).map(|i| Some(i as i32)).collect::<Vec<_>>(),
1618                    )) as ArrayRef,
1619                )])
1620                .into_data(),
1621            };
1622            assert_fast_roundtrip(&array);
1623        }
1624    }
1625}