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