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