Skip to main content

vortex_array/
serde.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::borrow::Cow;
5use std::fmt::Debug;
6use std::fmt::Formatter;
7use std::iter;
8use std::sync::Arc;
9
10use flatbuffers::FlatBufferBuilder;
11use flatbuffers::Follow;
12use flatbuffers::WIPOffset;
13use flatbuffers::root;
14use vortex_buffer::Alignment;
15use vortex_buffer::ByteBuffer;
16use vortex_error::VortexError;
17use vortex_error::VortexExpect;
18use vortex_error::VortexResult;
19use vortex_error::vortex_bail;
20use vortex_error::vortex_err;
21use vortex_error::vortex_panic;
22use vortex_flatbuffers::FlatBuffer;
23use vortex_flatbuffers::WriteFlatBuffer;
24use vortex_flatbuffers::array as fba;
25use vortex_flatbuffers::array::Compression;
26use vortex_session::VortexSession;
27use vortex_session::registry::ReadContext;
28use vortex_utils::aliases::hash_map::HashMap;
29
30use crate::ArrayContext;
31use crate::ArrayRef;
32use crate::ArraySlots;
33use crate::array::ArrayId;
34use crate::array::new_foreign_array;
35use crate::buffer::BufferHandle;
36use crate::dtype::DType;
37use crate::dtype::TryFromBytes;
38use crate::session::ArraySessionExt;
39use crate::stats::StatsSet;
40
41/// Options for serializing an array.
42#[derive(Default, Debug)]
43pub struct SerializeOptions {
44    /// The starting position within an external stream or file. This offset is used to compute
45    /// appropriate padding to enable zero-copy reads.
46    pub offset: usize,
47    /// Whether to include sufficient zero-copy padding.
48    pub include_padding: bool,
49}
50
51impl ArrayRef {
52    /// Serialize the array into a sequence of byte buffers that should be written contiguously.
53    /// This function returns a vec to avoid copying data buffers.
54    ///
55    /// Optionally, padding can be included to guarantee buffer alignment and ensure zero-copy
56    /// reads within the context of an external file or stream. In this case, the alignment of
57    /// the first byte buffer should be respected when writing the buffers to the stream or file.
58    ///
59    /// The format of this blob is a sequence of data buffers, possible with prefixed padding,
60    /// followed by a flatbuffer containing an [`fba::Array`] message, and ending with a
61    /// little-endian u32 describing the length of the flatbuffer message.
62    pub fn serialize(
63        &self,
64        ctx: &ArrayContext,
65        session: &VortexSession,
66        options: &SerializeOptions,
67    ) -> VortexResult<Vec<ByteBuffer>> {
68        // Collect all array buffers
69        let array_buffers = self
70            .depth_first_traversal()
71            .flat_map(|f| f.buffers())
72            .collect::<Vec<_>>();
73
74        // Allocate result buffers, including a possible padding buffer for each.
75        let mut buffers = vec![];
76        let mut fb_buffers = Vec::with_capacity(buffers.capacity());
77
78        // If we're including padding, we need to find the maximum required buffer alignment.
79        let max_alignment = array_buffers
80            .iter()
81            .map(|buf| buf.alignment())
82            .chain(iter::once(FlatBuffer::alignment()))
83            .max()
84            .unwrap_or_else(FlatBuffer::alignment);
85
86        // Create a shared buffer of zeros we can use for padding
87        let zeros = ByteBuffer::zeroed(*max_alignment);
88
89        // We push an empty buffer with the maximum alignment, so then subsequent buffers
90        // will be aligned. For subsequent buffers, we always push a 1-byte alignment.
91        buffers.push(ByteBuffer::zeroed_aligned(0, max_alignment));
92
93        // Keep track of where we are in the "file" to calculate padding.
94        let mut pos = options.offset;
95
96        // Push all the array buffers with padding as necessary.
97        for buffer in array_buffers {
98            let padding = if options.include_padding {
99                let padding = pos.next_multiple_of(*buffer.alignment()) - pos;
100                if padding > 0 {
101                    pos += padding;
102                    buffers.push(zeros.slice(0..padding));
103                }
104                padding
105            } else {
106                0
107            };
108
109            fb_buffers.push(fba::Buffer::new(
110                u16::try_from(padding).vortex_expect("padding fits into u16"),
111                buffer.alignment().exponent(),
112                Compression::None,
113                u32::try_from(buffer.len())
114                    .map_err(|_| vortex_err!("All buffers must fit into u32 for serialization"))?,
115            ));
116
117            pos += buffer.len();
118            buffers.push(buffer.aligned(Alignment::none()));
119        }
120
121        // Set up the flatbuffer builder
122        let mut fbb = FlatBufferBuilder::new();
123
124        let root = ArrayNodeFlatBuffer::try_new(ctx, session, self)?;
125        let fb_root = root.try_write_flatbuffer(&mut fbb)?;
126
127        let fb_buffers = fbb.create_vector(&fb_buffers);
128        let fb_array = fba::Array::create(
129            &mut fbb,
130            &fba::ArrayArgs {
131                root: Some(fb_root),
132                buffers: Some(fb_buffers),
133            },
134        );
135        fbb.finish_minimal(fb_array);
136        let (fb_vec, fb_start) = fbb.collapse();
137        let fb_end = fb_vec.len();
138        let fb_buffer = ByteBuffer::from(fb_vec).slice(fb_start..fb_end);
139        let fb_length = fb_buffer.len();
140
141        if options.include_padding {
142            let padding = pos.next_multiple_of(*FlatBuffer::alignment()) - pos;
143            if padding > 0 {
144                buffers.push(zeros.slice(0..padding));
145            }
146        }
147        buffers.push(fb_buffer);
148
149        // Finally, we write down the u32 length for the flatbuffer.
150        buffers.push(ByteBuffer::from(
151            u32::try_from(fb_length)
152                .map_err(|_| vortex_err!("Array metadata flatbuffer must fit into u32 for serialization. Array encoding tree is too large."))?
153                .to_le_bytes()
154                .to_vec(),
155        ));
156
157        Ok(buffers)
158    }
159}
160
161/// A utility struct for creating an [`fba::ArrayNode`] flatbuffer.
162pub struct ArrayNodeFlatBuffer<'a> {
163    ctx: &'a ArrayContext,
164    session: &'a VortexSession,
165    array: &'a ArrayRef,
166    buffer_idx: u16,
167}
168
169impl<'a> ArrayNodeFlatBuffer<'a> {
170    pub fn try_new(
171        ctx: &'a ArrayContext,
172        session: &'a VortexSession,
173        array: &'a ArrayRef,
174    ) -> VortexResult<Self> {
175        let n_buffers_recursive = array.nbuffers_recursive();
176        if n_buffers_recursive > u16::MAX as usize {
177            vortex_bail!(
178                "Array and all descendent arrays can have at most u16::MAX buffers: {}",
179                n_buffers_recursive
180            );
181        };
182        Ok(Self {
183            ctx,
184            session,
185            array,
186            buffer_idx: 0,
187        })
188    }
189
190    pub fn try_write_flatbuffer<'fb>(
191        &self,
192        fbb: &mut FlatBufferBuilder<'fb>,
193    ) -> VortexResult<WIPOffset<fba::ArrayNode<'fb>>> {
194        let encoding_idx = self.ctx.intern(&self.array.encoding_id()).ok_or_else(|| {
195            vortex_err!(
196                "Array encoding {} not permitted by ctx",
197                self.array.encoding_id()
198            )
199        })?;
200
201        let metadata_bytes = self.session.array_serialize(self.array)?.ok_or_else(|| {
202            vortex_err!(
203                "Array {} does not support serialization",
204                self.array.encoding_id()
205            )
206        })?;
207        let metadata = Some(fbb.create_vector(metadata_bytes.as_slice()));
208
209        // Assign buffer indices for all child arrays.
210        let nbuffers = u16::try_from(self.array.nbuffers())
211            .map_err(|_| vortex_err!("Array can have at most u16::MAX buffers"))?;
212        let mut child_buffer_idx = self.buffer_idx + nbuffers;
213
214        let children = self
215            .array
216            .children()
217            .iter()
218            .map(|child| {
219                // Update the number of buffers required.
220                let msg = ArrayNodeFlatBuffer {
221                    ctx: self.ctx,
222                    session: self.session,
223                    array: child,
224                    buffer_idx: child_buffer_idx,
225                }
226                .try_write_flatbuffer(fbb)?;
227
228                child_buffer_idx = u16::try_from(child.nbuffers_recursive())
229                    .ok()
230                    .and_then(|nbuffers| nbuffers.checked_add(child_buffer_idx))
231                    .ok_or_else(|| vortex_err!("Too many buffers (u16) for Array"))?;
232
233                Ok(msg)
234            })
235            .collect::<VortexResult<Vec<_>>>()?;
236        let children = Some(fbb.create_vector(&children));
237
238        let buffers = Some(fbb.create_vector_from_iter((0..nbuffers).map(|i| i + self.buffer_idx)));
239        let stats = Some(self.array.statistics().write_flatbuffer(fbb)?);
240
241        Ok(fba::ArrayNode::create(
242            fbb,
243            &fba::ArrayNodeArgs {
244                encoding: encoding_idx,
245                metadata,
246                children,
247                buffers,
248                stats,
249            },
250        ))
251    }
252}
253
254/// To minimize the serialized form, arrays do not persist their own dtype and length. Instead,
255/// parent arrays pass this information down during deserialization.
256pub trait ArrayChildren {
257    /// Returns the nth child of the array with the given dtype and length.
258    fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult<ArrayRef>;
259
260    /// The number of children.
261    fn len(&self) -> usize;
262
263    /// Returns true if there are no children.
264    fn is_empty(&self) -> bool {
265        self.len() == 0
266    }
267}
268
269impl<T: AsRef<[ArrayRef]>> ArrayChildren for T {
270    fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult<ArrayRef> {
271        let array = self.as_ref()[index].clone();
272        assert_eq!(array.len(), len);
273        assert_eq!(array.dtype(), dtype);
274        Ok(array)
275    }
276
277    fn len(&self) -> usize {
278        self.as_ref().len()
279    }
280}
281
282/// [`SerializedArray`] represents a parsed but not-yet-decoded deserialized array.
283/// It contains all the information from the serialized form, without anything extra. i.e.
284/// it is missing a [`DType`] and `len`, and the `encoding_id` is not yet resolved to a concrete
285/// vtable.
286///
287/// An [`SerializedArray`] can be fully decoded into an [`ArrayRef`] using the `decode` function.
288#[derive(Clone)]
289pub struct SerializedArray {
290    // Typed as fb::ArrayNode
291    flatbuffer: FlatBuffer,
292    // The location of the current fb::ArrayNode
293    flatbuffer_loc: usize,
294    buffers: Arc<[BufferHandle]>,
295}
296
297impl Debug for SerializedArray {
298    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
299        f.debug_struct("SerializedArray")
300            .field("encoding_id", &self.encoding_id())
301            .field("children", &(0..self.nchildren()).map(|i| self.child(i)))
302            .field(
303                "buffers",
304                &(0..self.nbuffers()).map(|i| self.buffer(i).ok()),
305            )
306            .field("metadata", &self.metadata())
307            .finish()
308    }
309}
310
311impl SerializedArray {
312    /// Decode an [`SerializedArray`] into an [`ArrayRef`].
313    pub fn decode(
314        &self,
315        dtype: &DType,
316        len: usize,
317        ctx: &ReadContext,
318        session: &VortexSession,
319    ) -> VortexResult<ArrayRef> {
320        let encoding_idx = self.flatbuffer().encoding();
321        let encoding_id = ctx
322            .resolve(encoding_idx)
323            .ok_or_else(|| vortex_err!("Unknown encoding index: {}", encoding_idx))?;
324        let Some(plugin) = session.arrays().registry().get(&encoding_id) else {
325            if session.allows_unknown() {
326                return self.decode_foreign(encoding_id, dtype, len, ctx);
327            }
328            vortex_bail!("Unknown encoding: {}", encoding_id);
329        };
330
331        let children = SerializedArrayChildren {
332            ser: self,
333            ctx,
334            session,
335        };
336
337        let buffers = self.collect_buffers()?;
338
339        let decoded =
340            plugin.deserialize(dtype, len, self.metadata(), &buffers, &children, session)?;
341
342        assert_eq!(
343            decoded.len(),
344            len,
345            "Array decoded from {} has incorrect length {}, expected {}",
346            encoding_id,
347            decoded.len(),
348            len
349        );
350        assert_eq!(
351            decoded.dtype(),
352            dtype,
353            "Array decoded from {} has incorrect dtype {}, expected {}",
354            encoding_id,
355            decoded.dtype(),
356            dtype,
357        );
358
359        assert!(
360            plugin.is_supported_encoding(&decoded.encoding_id()),
361            "Array decoded from {} has incorrect encoding {}",
362            encoding_id,
363            decoded.encoding_id(),
364        );
365
366        // Populate statistics from the serialized array.
367        if let Some(stats) = self.flatbuffer().stats() {
368            decoded
369                .statistics()
370                .set_iter(StatsSet::from_flatbuffer(&stats, dtype, session)?.into_iter());
371        }
372
373        Ok(decoded)
374    }
375
376    fn decode_foreign(
377        &self,
378        encoding_id: ArrayId,
379        dtype: &DType,
380        len: usize,
381        ctx: &ReadContext,
382    ) -> VortexResult<ArrayRef> {
383        let children = (0..self.nchildren())
384            .map(|idx| {
385                let child = self.child(idx);
386                let child_encoding_idx = child.flatbuffer().encoding();
387                let child_encoding_id = ctx
388                    .resolve(child_encoding_idx)
389                    .ok_or_else(|| vortex_err!("Unknown encoding index: {}", child_encoding_idx))?;
390                child
391                    .decode_foreign(child_encoding_id, dtype, len, ctx)
392                    .map(Some)
393            })
394            .collect::<VortexResult<ArraySlots>>()?;
395
396        new_foreign_array(
397            encoding_id,
398            dtype.clone(),
399            len,
400            self.metadata().to_vec(),
401            self.collect_buffers()?.into_owned(),
402            children,
403        )
404    }
405
406    /// Returns the array encoding.
407    pub fn encoding_id(&self) -> u16 {
408        self.flatbuffer().encoding()
409    }
410
411    /// Returns the array metadata bytes.
412    pub fn metadata(&self) -> &[u8] {
413        self.flatbuffer()
414            .metadata()
415            .map(|metadata| metadata.bytes())
416            .unwrap_or(&[])
417    }
418
419    /// Returns the number of children.
420    pub fn nchildren(&self) -> usize {
421        self.flatbuffer()
422            .children()
423            .map_or(0, |children| children.len())
424    }
425
426    /// Returns the nth child of the array.
427    pub fn child(&self, idx: usize) -> SerializedArray {
428        let children = self
429            .flatbuffer()
430            .children()
431            .vortex_expect("Expected array to have children");
432        if idx >= children.len() {
433            vortex_panic!(
434                "Invalid child index {} for array with {} children",
435                idx,
436                children.len()
437            );
438        }
439        self.with_root(children.get(idx))
440    }
441
442    /// Returns the number of buffers.
443    pub fn nbuffers(&self) -> usize {
444        self.flatbuffer()
445            .buffers()
446            .map_or(0, |buffers| buffers.len())
447    }
448
449    /// Returns the nth buffer of the current array.
450    pub fn buffer(&self, idx: usize) -> VortexResult<BufferHandle> {
451        let buffer_idx = self
452            .flatbuffer()
453            .buffers()
454            .ok_or_else(|| vortex_err!("Array has no buffers"))?
455            .get(idx);
456        self.buffers
457            .get(buffer_idx as usize)
458            .cloned()
459            .ok_or_else(|| {
460                vortex_err!(
461                    "Invalid buffer index {} for array with {} buffers",
462                    buffer_idx,
463                    self.nbuffers()
464                )
465            })
466    }
467
468    /// Returns all buffers for the current array node.
469    ///
470    /// If buffer indices are contiguous, returns a zero-copy borrowed slice.
471    /// Otherwise falls back to collecting each buffer individually.
472    fn collect_buffers(&self) -> VortexResult<Cow<'_, [BufferHandle]>> {
473        let Some(fb_buffers) = self.flatbuffer().buffers() else {
474            return Ok(Cow::Borrowed(&[]));
475        };
476        let count = fb_buffers.len();
477        if count == 0 {
478            return Ok(Cow::Borrowed(&[]));
479        }
480        let start = fb_buffers.get(0) as usize;
481        let contiguous = fb_buffers
482            .iter()
483            .enumerate()
484            .all(|(i, idx)| idx as usize == start + i);
485        if contiguous {
486            self.buffers.get(start..start + count).map_or_else(
487                || {
488                    vortex_bail!(
489                        "buffer indices {}..{} out of range for {} buffers",
490                        start,
491                        start + count,
492                        self.buffers.len()
493                    )
494                },
495                |slice| Ok(Cow::Borrowed(slice)),
496            )
497        } else {
498            (0..count)
499                .map(|idx| self.buffer(idx))
500                .collect::<VortexResult<Vec<_>>>()
501                .map(Cow::Owned)
502        }
503    }
504
505    /// Returns the buffer lengths as stored in the flatbuffer metadata.
506    ///
507    /// This reads the buffer descriptors from the flatbuffer, which contain the
508    /// serialized length of each buffer. This is useful for displaying buffer sizes
509    /// without needing to access the actual buffer data.
510    pub fn buffer_lengths(&self) -> Vec<usize> {
511        let fb_array = root::<fba::Array>(self.flatbuffer.as_ref())
512            .vortex_expect("SerializedArray flatbuffer must be a valid Array");
513        fb_array
514            .buffers()
515            .map(|buffers| buffers.iter().map(|b| b.length() as usize).collect())
516            .unwrap_or_default()
517    }
518
519    /// Validate and align the array tree flatbuffer, returning the aligned buffer and root location.
520    fn validate_array_tree(array_tree: impl Into<ByteBuffer>) -> VortexResult<(FlatBuffer, usize)> {
521        let fb_buffer = FlatBuffer::align_from(array_tree.into());
522        let fb_array = root::<fba::Array>(fb_buffer.as_ref())?;
523        let fb_root = fb_array
524            .root()
525            .ok_or_else(|| vortex_err!("Array must have a root node"))?;
526        let flatbuffer_loc = fb_root._tab.loc();
527        Ok((fb_buffer, flatbuffer_loc))
528    }
529
530    /// Create an [`SerializedArray`] from a pre-existing array tree flatbuffer and pre-resolved buffer
531    /// handles.
532    ///
533    /// The caller is responsible for resolving buffers from whatever source (device segments, host
534    /// overrides, or a mix). The buffers must be in the same order as the `Array.buffers` descriptor
535    /// list in the flatbuffer.
536    pub fn from_flatbuffer_with_buffers(
537        array_tree: impl Into<ByteBuffer>,
538        buffers: Vec<BufferHandle>,
539    ) -> VortexResult<Self> {
540        let (flatbuffer, flatbuffer_loc) = Self::validate_array_tree(array_tree)?;
541        Ok(SerializedArray {
542            flatbuffer,
543            flatbuffer_loc,
544            buffers: buffers.into(),
545        })
546    }
547
548    /// Create an [`SerializedArray`] from a raw array tree flatbuffer (metadata only).
549    ///
550    /// This constructor creates a `SerializedArray` with no buffer data, useful for
551    /// inspecting the metadata when the actual buffer data is not needed
552    /// (e.g., displaying buffer sizes from inlined array tree metadata).
553    ///
554    /// Note: Calling `buffer()` on the returned `SerializedArray` will fail since
555    /// no actual buffer data is available.
556    pub fn from_array_tree(array_tree: impl Into<ByteBuffer>) -> VortexResult<Self> {
557        let (flatbuffer, flatbuffer_loc) = Self::validate_array_tree(array_tree)?;
558        Ok(SerializedArray {
559            flatbuffer,
560            flatbuffer_loc,
561            buffers: Arc::new([]),
562        })
563    }
564
565    /// Returns the root ArrayNode flatbuffer.
566    fn flatbuffer(&self) -> fba::ArrayNode<'_> {
567        unsafe { fba::ArrayNode::follow(self.flatbuffer.as_ref(), self.flatbuffer_loc) }
568    }
569
570    /// Returns a new [`SerializedArray`] with the given node as the root
571    // TODO(ngates): we may want a wrapper that avoids this clone.
572    fn with_root(&self, root: fba::ArrayNode) -> Self {
573        let mut this = self.clone();
574        this.flatbuffer_loc = root._tab.loc();
575        this
576    }
577
578    /// Create an [`SerializedArray`] from a pre-existing flatbuffer (ArrayNode) and a segment containing
579    /// only the data buffers (without the flatbuffer suffix).
580    ///
581    /// This is used when the flatbuffer is stored separately in layout metadata (e.g., when
582    /// `FLAT_LAYOUT_INLINE_ARRAY_NODE` is enabled).
583    pub fn from_flatbuffer_and_segment(
584        array_tree: ByteBuffer,
585        segment: BufferHandle,
586    ) -> VortexResult<Self> {
587        // HashMap::new doesn't allocate when empty, so this has no overhead
588        Self::from_flatbuffer_and_segment_with_overrides(array_tree, segment, &HashMap::new())
589    }
590
591    /// Create an [`SerializedArray`] from a pre-existing flatbuffer (ArrayNode) and a segment,
592    /// substituting host-resident buffer overrides for specific buffer indices.
593    ///
594    /// Buffers whose index appears in `buffer_overrides` are resolved from the provided
595    /// host data instead of the segment. All other buffers are sliced from the segment
596    /// using the padding and alignment described in the flatbuffer.
597    pub fn from_flatbuffer_and_segment_with_overrides(
598        array_tree: ByteBuffer,
599        segment: BufferHandle,
600        buffer_overrides: &HashMap<u32, ByteBuffer>,
601    ) -> VortexResult<Self> {
602        // We align each buffer individually, so we remove alignment requirements on the segment
603        // for host-resident buffers. Device buffers are sliced directly.
604        let segment = segment.ensure_aligned(Alignment::none())?;
605
606        // this can't return the validated array because there is no lifetime to give it, so we
607        // need to cast it below, which is safe.
608        let (fb_buffer, flatbuffer_loc) = Self::validate_array_tree(array_tree)?;
609        // SAFETY: fb_buffer was already validated by validate_array_tree above.
610        let fb_array = unsafe { fba::root_as_array_unchecked(fb_buffer.as_ref()) };
611
612        let mut offset = 0usize;
613        let buffers = fb_array
614            .buffers()
615            .unwrap_or_default()
616            .iter()
617            .enumerate()
618            .map(|(idx, fb_buf)| {
619                let idx = u32::try_from(idx).vortex_expect("buffer count must fit in u32");
620
621                // The padding, length, and resulting offsets all come from the flatbuffer, which
622                // may be corrupt. Use checked arithmetic so malformed metadata returns a
623                // `VortexError` rather than panicking (see issue #8819).
624                let buffer_len = fb_buf.length() as usize;
625                let start = offset
626                    .checked_add(fb_buf.padding() as usize)
627                    .ok_or_else(|| {
628                        vortex_err!("Buffer {idx} offset overflows when adding its padding")
629                    })?;
630                let end = start.checked_add(buffer_len).ok_or_else(|| {
631                    vortex_err!("Buffer {idx} offset overflows when adding its length")
632                })?;
633
634                // The alignment exponent comes from the flatbuffer and may be corrupt, so validate
635                // it rather than panicking on a too-large shift (see issue #8819).
636                let alignment =
637                    Alignment::try_from_untrusted_exponent(fb_buf.alignment_exponent())?;
638                let handle = if let Some(host_data) = buffer_overrides.get(&idx) {
639                    BufferHandle::new_host(host_data.clone()).ensure_aligned(alignment)?
640                } else {
641                    // Bounds-check against the segment so an out-of-range buffer returns a
642                    // `VortexError` rather than panicking when slicing (see issue #8819).
643                    if end > segment.len() {
644                        vortex_bail!(
645                            "Buffer {idx} at offset {start} with length {buffer_len} is out of \
646                             bounds of the {}-byte segment",
647                            segment.len(),
648                        );
649                    }
650                    segment.slice(start..end).ensure_aligned(alignment)?
651                };
652
653                offset = end;
654                Ok(handle)
655            })
656            .collect::<VortexResult<Arc<[_]>>>()?;
657
658        Ok(SerializedArray {
659            flatbuffer: fb_buffer,
660            flatbuffer_loc,
661            buffers,
662        })
663    }
664}
665
666struct SerializedArrayChildren<'a> {
667    ser: &'a SerializedArray,
668    ctx: &'a ReadContext,
669    session: &'a VortexSession,
670}
671
672impl ArrayChildren for SerializedArrayChildren<'_> {
673    fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult<ArrayRef> {
674        self.ser
675            .child(index)
676            .decode(dtype, len, self.ctx, self.session)
677    }
678
679    fn len(&self) -> usize {
680        self.ser.nchildren()
681    }
682}
683
684impl TryFrom<ByteBuffer> for SerializedArray {
685    type Error = VortexError;
686
687    fn try_from(value: ByteBuffer) -> Result<Self, Self::Error> {
688        // The final 4 bytes contain the length of the flatbuffer.
689        if value.len() < 4 {
690            vortex_bail!("SerializedArray buffer is too short");
691        }
692
693        // We align each buffer individually, so we remove alignment requirements on the buffer.
694        let value = value.aligned(Alignment::none());
695
696        let fb_length = u32::try_from_le_bytes(&value.as_slice()[value.len() - 4..])? as usize;
697        if value.len() < 4 + fb_length {
698            vortex_bail!("SerializedArray buffer is too short for flatbuffer");
699        }
700
701        let fb_offset = value.len() - 4 - fb_length;
702        let array_tree = value.slice(fb_offset..fb_offset + fb_length);
703        let segment = BufferHandle::new_host(value.slice(0..fb_offset));
704
705        Self::from_flatbuffer_and_segment(array_tree, segment)
706    }
707}
708
709impl TryFrom<BufferHandle> for SerializedArray {
710    type Error = VortexError;
711
712    fn try_from(value: BufferHandle) -> Result<Self, Self::Error> {
713        Self::try_from(value.try_to_host_sync()?)
714    }
715}
716
717#[cfg(test)]
718mod tests {
719    use vortex_buffer::ByteBufferMut;
720
721    use super::*;
722    use crate::IntoArray;
723    use crate::array_session;
724    use crate::arrays::PrimitiveArray;
725
726    /// A corrupt array tree can declare a buffer that extends past the backing segment. Slicing
727    /// such a buffer must return a [`VortexError`] rather than panicking (see issue #8819).
728    #[test]
729    fn from_flatbuffer_and_segment_rejects_out_of_bounds_buffer() -> VortexResult<()> {
730        let session = array_session();
731        let array_ctx = ArrayContext::empty();
732
733        // Serialize a simple array so we have a valid array tree flatbuffer whose declared buffer
734        // lengths describe the trailing data segment.
735        let serialized = PrimitiveArray::from_iter([1i32, 2, 3, 4])
736            .into_array()
737            .serialize(&array_ctx, &session, &SerializeOptions::default())?;
738
739        let mut concat = ByteBufferMut::empty();
740        for buf in serialized {
741            concat.extend_from_slice(buf.as_ref());
742        }
743        let value = concat.freeze().aligned(Alignment::none());
744
745        // Split the blob into the trailing flatbuffer and the leading data segment, mirroring
746        // `SerializedArray::try_from`.
747        let fb_length = u32::try_from_le_bytes(&value.as_slice()[value.len() - 4..])? as usize;
748        let fb_offset = value.len() - 4 - fb_length;
749        assert!(
750            fb_offset > 0,
751            "the array must have at least one data buffer"
752        );
753        let array_tree = value.slice(fb_offset..fb_offset + fb_length);
754
755        // Truncate the data segment by one byte so the declared buffer no longer fits.
756        let truncated = BufferHandle::new_host(value.slice(0..fb_offset - 1));
757
758        let Some(err) = SerializedArray::from_flatbuffer_and_segment(array_tree, truncated).err()
759        else {
760            vortex_bail!("out-of-bounds buffer must be rejected");
761        };
762        assert!(
763            err.to_string().contains("out of bounds"),
764            "unexpected error: {err}"
765        );
766
767        Ok(())
768    }
769
770    /// A corrupt array tree can declare a buffer alignment of up to 2^63, which the copy that
771    /// satisfies it allocates as slack. It must be rejected (see issue #8819).
772    #[test]
773    fn from_flatbuffer_and_segment_rejects_excessive_buffer_alignment() -> VortexResult<()> {
774        // Padding and length fit the segment exactly, so the exponent is the only defect.
775        // `validate_array_tree` only requires a root node to be present, so an empty one will do.
776        let mut fbb = FlatBufferBuilder::new();
777        let fb_root = fba::ArrayNode::create(&mut fbb, &fba::ArrayNodeArgs::default());
778        let fb_buffers = fbb.create_vector(&[fba::Buffer::new(0, 40, Compression::None, 4)]);
779        let fb_array = fba::Array::create(
780            &mut fbb,
781            &fba::ArrayArgs {
782                root: Some(fb_root),
783                buffers: Some(fb_buffers),
784            },
785        );
786        fbb.finish_minimal(fb_array);
787        let (fb_vec, fb_start) = fbb.collapse();
788        let fb_end = fb_vec.len();
789        let array_tree = ByteBuffer::from(fb_vec).slice(fb_start..fb_end);
790
791        let segment = BufferHandle::new_host(ByteBuffer::from(vec![0u8; 4]));
792        let Some(err) = SerializedArray::from_flatbuffer_and_segment(array_tree, segment).err()
793        else {
794            vortex_bail!("excessive buffer alignment must be rejected");
795        };
796        assert!(
797            err.to_string().contains("exceeds"),
798            "unexpected error: {err}"
799        );
800
801        Ok(())
802    }
803}