1use 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#[derive(Default, Debug)]
43pub struct SerializeOptions {
44 pub offset: usize,
47 pub include_padding: bool,
49}
50
51impl ArrayRef {
52 pub fn serialize(
63 &self,
64 ctx: &ArrayContext,
65 session: &VortexSession,
66 options: &SerializeOptions,
67 ) -> VortexResult<Vec<ByteBuffer>> {
68 let array_buffers = self
70 .depth_first_traversal()
71 .flat_map(|f| f.buffers())
72 .collect::<Vec<_>>();
73
74 let mut buffers = vec![];
76 let mut fb_buffers = Vec::with_capacity(buffers.capacity());
77
78 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 let zeros = ByteBuffer::zeroed(*max_alignment);
88
89 buffers.push(ByteBuffer::zeroed_aligned(0, max_alignment));
92
93 let mut pos = options.offset;
95
96 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 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 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
161pub 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 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 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
254pub trait ArrayChildren {
257 fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult<ArrayRef>;
259
260 fn len(&self) -> usize;
262
263 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#[derive(Clone)]
289pub struct SerializedArray {
290 flatbuffer: FlatBuffer,
292 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 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 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 pub fn encoding_id(&self) -> u16 {
408 self.flatbuffer().encoding()
409 }
410
411 pub fn metadata(&self) -> &[u8] {
413 self.flatbuffer()
414 .metadata()
415 .map(|metadata| metadata.bytes())
416 .unwrap_or(&[])
417 }
418
419 pub fn nchildren(&self) -> usize {
421 self.flatbuffer()
422 .children()
423 .map_or(0, |children| children.len())
424 }
425
426 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 pub fn nbuffers(&self) -> usize {
444 self.flatbuffer()
445 .buffers()
446 .map_or(0, |buffers| buffers.len())
447 }
448
449 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 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 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 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 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 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 fn flatbuffer(&self) -> fba::ArrayNode<'_> {
567 unsafe { fba::ArrayNode::follow(self.flatbuffer.as_ref(), self.flatbuffer_loc) }
568 }
569
570 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 pub fn from_flatbuffer_and_segment(
584 array_tree: ByteBuffer,
585 segment: BufferHandle,
586 ) -> VortexResult<Self> {
587 Self::from_flatbuffer_and_segment_with_overrides(array_tree, segment, &HashMap::new())
589 }
590
591 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 let segment = segment.ensure_aligned(Alignment::none())?;
605
606 let (fb_buffer, flatbuffer_loc) = Self::validate_array_tree(array_tree)?;
609 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 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 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 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 if value.len() < 4 {
690 vortex_bail!("SerializedArray buffer is too short");
691 }
692
693 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 #[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 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 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 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 #[test]
773 fn from_flatbuffer_and_segment_rejects_excessive_buffer_alignment() -> VortexResult<()> {
774 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}