16728097825

Committed 04 Aug 2025 04:00PM UTC coverage: 48.355% (-35.1%) from 83.429%

Build # 16728097825

Build Type

Pull #4108

github

Committed by

web-flow

Commit Message

Merge 1b2d27fd8 into 649ba9576

Pull Request Pull Request #4108: perf[vortex-array]: use all_valid instead of `invalid_count() == 0`

Coverage Stats

1 of 1 new or added line in 1 file covered. (100.0%)

11596 existing lines in 378 files now uncovered.

18635 of 38538 relevant lines covered (48.35%)

151786.4 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

86.04

/vortex-array/src/serde.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::fmt::{Debug, Formatter};
use std::iter;
use std::sync::Arc;

use flatbuffers::{FlatBufferBuilder, Follow, WIPOffset, root};
use itertools::Itertools;
use vortex_buffer::{Alignment, ByteBuffer};
use vortex_dtype::{DType, TryFromBytes};
use vortex_error::{
    VortexError, VortexExpect, VortexResult, vortex_bail, vortex_err, vortex_panic,
};
use vortex_flatbuffers::array::Compression;
use vortex_flatbuffers::{
    FlatBuffer, FlatBufferRoot, ReadFlatBuffer, WriteFlatBuffer, array as fba,
};

use crate::stats::StatsSet;
use crate::{Array, ArrayContext, ArrayRef, ArrayVisitor, ArrayVisitorExt};

/// Options for serializing an array.
#[derive(Default, Debug)]
pub struct SerializeOptions {
    /// The starting position within an external stream or file. This offset is used to compute
    /// appropriate padding to enable zero-copy reads.
    pub offset: usize,
    /// Whether to include sufficient zero-copy padding.
    pub include_padding: bool,
}

impl dyn Array + '_ {
    /// Serialize the array into a sequence of byte buffers that should be written contiguously.
    /// This function returns a vec to avoid copying data buffers.
    ///
    /// Optionally, padding can be included to guarantee buffer alignment and ensure zero-copy
    /// reads within the context of an external file or stream. In this case, the alignment of
    /// the first byte buffer should be respected when writing the buffers to the stream or file.
    ///
    /// The format of this blob is a sequence of data buffers, possible with prefixed padding,
    /// followed by a flatbuffer containing an [`fba::Array`] message, and ending with a
    /// little-endian u32 describing the length of the flatbuffer message.
    pub fn serialize(
        &self,
        ctx: &ArrayContext,
        options: &SerializeOptions,
    ) -> VortexResult<Vec<ByteBuffer>> {
        // Collect all array buffers
        let array_buffers = self
            .depth_first_traversal()
            .flat_map(|f| f.buffers())
            .collect::<Vec<_>>();

        // Allocate result buffers, including a possible padding buffer for each.
        let mut buffers = vec![];
        let mut fb_buffers = Vec::with_capacity(buffers.capacity());

        // If we're including padding, we need to find the maximum required buffer alignment.
        let max_alignment = array_buffers
            .iter()
            .map(|buf| buf.alignment())
            .chain(iter::once(FlatBuffer::alignment()))
            .max()
            .unwrap_or_else(FlatBuffer::alignment);

        // Create a shared buffer of zeros we can use for padding
        let zeros = ByteBuffer::zeroed(*max_alignment);

        // We push an empty buffer with the maximum alignment, so then subsequent buffers
        // will be aligned. For subsequent buffers, we always push a 1-byte alignment.
        buffers.push(ByteBuffer::zeroed_aligned(0, max_alignment));

        // Keep track of where we are in the "file" to calculate padding.
        let mut pos = options.offset;

        // Push all the array buffers with padding as necessary.
        for buffer in array_buffers {
            let padding = if options.include_padding {
                let padding = pos.next_multiple_of(*buffer.alignment()) - pos;
                if padding > 0 {
                    pos += padding;
                    buffers.push(zeros.slice(0..padding));
                }
                padding
            } else {
                0
            };

            fb_buffers.push(fba::Buffer::new(
                u16::try_from(padding).vortex_expect("padding fits into u16"),
                buffer.alignment().exponent(),
                Compression::None,
                u32::try_from(buffer.len())
                    .map_err(|_| vortex_err!("All buffers must fit into u32 for serialization"))?,
            ));

            pos += buffer.len();
            buffers.push(buffer.aligned(Alignment::none()));
        }

        // Set up the flatbuffer builder
        let mut fbb = FlatBufferBuilder::new();
        let root = ArrayNodeFlatBuffer::try_new(ctx, self)?;
        let fb_root = root.write_flatbuffer(&mut fbb);
        let fb_buffers = fbb.create_vector(&fb_buffers);
        let fb_array = fba::Array::create(
            &mut fbb,
            &fba::ArrayArgs {
                root: Some(fb_root),
                buffers: Some(fb_buffers),
            },
        );
        fbb.finish_minimal(fb_array);
        let (fb_vec, fb_start) = fbb.collapse();
        let fb_end = fb_vec.len();
        let fb_buffer = ByteBuffer::from(fb_vec).slice(fb_start..fb_end);
        let fb_length = fb_buffer.len();

        if options.include_padding {
            let padding = pos.next_multiple_of(*FlatBuffer::alignment()) - pos;
            if padding > 0 {
                buffers.push(zeros.slice(0..padding));
            }
        }
        buffers.push(fb_buffer);

        // Finally, we write down the u32 length for the flatbuffer.
        buffers.push(ByteBuffer::from(
            u32::try_from(fb_length)
                .map_err(|_| vortex_err!("Array metadata flatbuffer must fit into u32 for serialization. Array encoding tree is too large."))?
                .to_le_bytes()
                .to_vec(),
        ));

        Ok(buffers)
    }
}

/// A utility struct for creating an [`fba::ArrayNode`] flatbuffer.
pub struct ArrayNodeFlatBuffer<'a> {
    ctx: &'a ArrayContext,
    array: &'a dyn Array,
    buffer_idx: u16,
}

impl<'a> ArrayNodeFlatBuffer<'a> {
    pub fn try_new(ctx: &'a ArrayContext, array: &'a dyn Array) -> VortexResult<Self> {
        // Depth-first traversal of the array to ensure it supports serialization.
        for child in array.depth_first_traversal() {
            if child.metadata()?.is_none() {
                vortex_bail!(
                    "Array {} does not support serialization",
                    child.encoding_id()
                );
            }
        }
        Ok(Self {
            ctx,
            array,
            buffer_idx: 0,
        })
    }
}

impl FlatBufferRoot for ArrayNodeFlatBuffer<'_> {}

impl WriteFlatBuffer for ArrayNodeFlatBuffer<'_> {
    type Target<'t> = fba::ArrayNode<'t>;

    fn write_flatbuffer<'fb>(
        &self,
        fbb: &mut FlatBufferBuilder<'fb>,
    ) -> WIPOffset<Self::Target<'fb>> {
        let encoding = self.ctx.encoding_idx(&self.array.encoding());
        let metadata = self
            .array
            .metadata()
            // TODO(ngates): add try_write_flatbuffer
            .vortex_expect("Failed to serialize metadata")
            .vortex_expect("Validated that all arrays support serialization");
        let metadata = Some(fbb.create_vector(metadata.as_slice()));

        // Assign buffer indices for all child arrays.
        let nbuffers = u16::try_from(self.array.nbuffers())
            .vortex_expect("Array can have at most u16::MAX buffers");
        let mut child_buffer_idx = self.buffer_idx + nbuffers;

        let children = &self
            .array
            .children()
            .iter()
            .map(|child| {
                // Update the number of buffers required.
                let msg = ArrayNodeFlatBuffer {
                    ctx: self.ctx,
                    array: child,
                    buffer_idx: child_buffer_idx,
                }
                .write_flatbuffer(fbb);
                child_buffer_idx = u16::try_from(child.nbuffers_recursive())
                    .ok()
                    .and_then(|nbuffers| nbuffers.checked_add(child_buffer_idx))
                    .vortex_expect("Too many buffers (u16) for Array");
                msg
            })
            .collect::<Vec<_>>();
        let children = Some(fbb.create_vector(children));

        let buffers = Some(fbb.create_vector_from_iter((0..nbuffers).map(|i| i + self.buffer_idx)));
        let stats = Some(self.array.statistics().to_owned().write_flatbuffer(fbb));

        fba::ArrayNode::create(
            fbb,
            &fba::ArrayNodeArgs {
                encoding,
                metadata,
                children,
                buffers,
                stats,
            },
        )
    }
}

/// To minimize the serialized form, arrays do not persist their own dtype and length. Instead,
/// parent arrays pass this information down during deserialization. This trait abstracts
/// over either a serialized [`crate::serde::ArrayParts`] or the
/// in-memory [`crate::data::ArrayData`].
pub trait ArrayChildren {
    /// Returns the nth child of the array with the given dtype and length.
    fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult<ArrayRef>;

    /// The number of children.
    fn len(&self) -> usize;

    /// Returns true if there are no children.
    fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// [`ArrayParts`] represents a parsed but not-yet-decoded deserialized [`Array`].
/// It contains all the information from the serialized form, without anything extra. i.e.
/// it is missing a [`DType`] and `len`, and the `encoding_id` is not yet resolved to a concrete
/// vtable.
///
/// An [`ArrayParts`] can be fully decoded into an [`ArrayRef`] using the `decode` function.
#[derive(Clone)]
pub struct ArrayParts {
    // Typed as fb::ArrayNode
    flatbuffer: FlatBuffer,
    // The location of the current fb::ArrayNode
    flatbuffer_loc: usize,
    buffers: Arc<[ByteBuffer]>,
}

impl Debug for ArrayParts {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ArrayParts")
            .field("encoding_id", &self.encoding_id())
            .field("children", &(0..self.nchildren()).map(|i| self.child(i)))
            .field(
                "buffers",
                &(0..self.nbuffers()).map(|i| self.buffer(i).ok()),
            )
            .field("metadata", &self.metadata())
            .finish()
    }
}

impl ArrayParts {
    /// Decode an [`ArrayParts`] into an [`ArrayRef`].
    pub fn decode(&self, ctx: &ArrayContext, dtype: &DType, len: usize) -> VortexResult<ArrayRef> {
        let encoding_id = self.flatbuffer().encoding();
        let vtable = ctx
            .lookup_encoding(encoding_id)
            .ok_or_else(|| vortex_err!("Unknown encoding: {}", encoding_id))?;

        let buffers: Vec<_> = (0..self.nbuffers())
            .map(|idx| self.buffer(idx))
            .try_collect()?;

        let children = ArrayPartsChildren { parts: self, ctx };

        let decoded = vtable.build(dtype, len, self.metadata(), &buffers, &children)?;

        assert_eq!(
            decoded.len(),
            len,
            "Array decoded from {} has incorrect length {}, expected {}",
            vtable.id(),
            decoded.len(),
            len
        );
        assert_eq!(
            decoded.dtype(),
            dtype,
            "Array decoded from {} has incorrect dtype {}, expected {}",
            vtable.id(),
            decoded.dtype(),
            dtype,
        );
        assert_eq!(
            decoded.encoding_id(),
            vtable.id(),
            "Array decoded from {} has incorrect encoding {}",
            vtable.id(),
            decoded.encoding_id(),
        );

        // Populate statistics from the serialized array.
        if let Some(stats) = self.flatbuffer().stats() {
            let decoded_statistics = decoded.statistics();
            StatsSet::read_flatbuffer(&stats)?
                .into_iter()
                .for_each(|(stat, val)| decoded_statistics.set(stat, val));
        }

        Ok(decoded)
    }

    /// Returns the array encoding.
    pub fn encoding_id(&self) -> u16 {
        self.flatbuffer().encoding()
    }

    /// Returns the array metadata bytes.
    pub fn metadata(&self) -> &[u8] {
        self.flatbuffer()
            .metadata()
            .map(|metadata| metadata.bytes())
            .unwrap_or(&[])
    }

    /// Returns the number of children.
    pub fn nchildren(&self) -> usize {
        self.flatbuffer()
            .children()
            .map_or(0, |children| children.len())
    }

    /// Returns the nth child of the array.
    pub fn child(&self, idx: usize) -> ArrayParts {
        let children = self
            .flatbuffer()
            .children()
            .vortex_expect("Expected array to have children");
        if idx >= children.len() {
            vortex_panic!(
                "Invalid child index {} for array with {} children",
                idx,
                children.len()
            );
        }
        self.with_root(children.get(idx))
    }

    /// Returns the number of buffers.
    pub fn nbuffers(&self) -> usize {
        self.flatbuffer()
            .buffers()
            .map_or(0, |buffers| buffers.len())
    }

    /// Returns the nth buffer of the current array.
    pub fn buffer(&self, idx: usize) -> VortexResult<ByteBuffer> {
        let buffer_idx = self
            .flatbuffer()
            .buffers()
            .ok_or_else(|| vortex_err!("Array has no buffers"))?
            .get(idx);
        self.buffers
            .get(buffer_idx as usize)
            .cloned()
            .ok_or_else(|| {
                vortex_err!(
                    "Invalid buffer index {} for array with {} buffers",
                    buffer_idx,
                    self.nbuffers()
                )
            })
    }

    /// Returns the root ArrayNode flatbuffer.
    fn flatbuffer(&self) -> fba::ArrayNode<'_> {
        unsafe { fba::ArrayNode::follow(self.flatbuffer.as_ref(), self.flatbuffer_loc) }
    }

    /// Returns a new [`ArrayParts`] with the given node as the root
    // TODO(ngates): we may want a wrapper that avoids this clone.
    fn with_root(&self, root: fba::ArrayNode) -> Self {
        let mut this = self.clone();
        this.flatbuffer_loc = root._tab.loc();
        this
    }
}

struct ArrayPartsChildren<'a> {
    parts: &'a ArrayParts,
    ctx: &'a ArrayContext,
}

impl ArrayChildren for ArrayPartsChildren<'_> {
    fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult<ArrayRef> {
        self.parts.child(index).decode(self.ctx, dtype, len)
    }

    fn len(&self) -> usize {
        self.parts.nchildren()
    }
}

impl TryFrom<ByteBuffer> for ArrayParts {
    type Error = VortexError;

    fn try_from(value: ByteBuffer) -> Result<Self, Self::Error> {
        // The final 4 bytes contain the length of the flatbuffer.
        if value.len() < 4 {
            vortex_bail!("ArrayParts buffer is too short");
        }

        // We align each buffer individually, so we remove alignment requirements on the buffer.
        let value = value.aligned(Alignment::none());

        let fb_length = u32::try_from_le_bytes(&value.as_slice()[value.len() - 4..])? as usize;
        if value.len() < 4 + fb_length {
            vortex_bail!("ArrayParts buffer is too short for flatbuffer");
        }

        let fb_offset = value.len() - 4 - fb_length;
        let fb_buffer = value.slice(fb_offset..fb_offset + fb_length);
        let fb_buffer = FlatBuffer::align_from(fb_buffer);

        let fb_array = root::<fba::Array>(fb_buffer.as_ref())?;
        let fb_root = fb_array.root().vortex_expect("Array must have a root node");

        let mut offset = 0;
        let buffers: Arc<[ByteBuffer]> = fb_array
            .buffers()
            .unwrap_or_default()
            .iter()
            .map(|fb_buffer| {
                // Skip padding
                offset += fb_buffer.padding() as usize;

                let buffer_len = fb_buffer.length() as usize;

                // Extract a buffer and ensure it's aligned, copying if necessary
                let buffer = value
                    .slice(offset..(offset + buffer_len))
                    .aligned(Alignment::from_exponent(fb_buffer.alignment_exponent()));

                offset += buffer_len;
                buffer
            })
            .collect();

        Ok(ArrayParts {
            flatbuffer: fb_buffer.clone(),
            flatbuffer_loc: fb_root._tab.loc(),
            buffers,
        })
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::fmt::{Debug, Formatter};
5	use std::iter;
6	use std::sync::Arc;
7
8	use flatbuffers::{FlatBufferBuilder, Follow, WIPOffset, root};
9	use itertools::Itertools;
10	use vortex_buffer::{Alignment, ByteBuffer};
11	use vortex_dtype::{DType, TryFromBytes};
12	use vortex_error::{
13	VortexError, VortexExpect, VortexResult, vortex_bail, vortex_err, vortex_panic,
14	};
15	use vortex_flatbuffers::array::Compression;
16	use vortex_flatbuffers::{
17	FlatBuffer, FlatBufferRoot, ReadFlatBuffer, WriteFlatBuffer, array as fba,
18	};
19
20	use crate::stats::StatsSet;
21	use crate::{Array, ArrayContext, ArrayRef, ArrayVisitor, ArrayVisitorExt};
22
23	/// Options for serializing an array.
24	#[derive(Default, Debug)]
25	pub struct SerializeOptions {
26	/// The starting position within an external stream or file. This offset is used to compute
27	/// appropriate padding to enable zero-copy reads.
28	pub offset: usize,
29	/// Whether to include sufficient zero-copy padding.
30	pub include_padding: bool,
31	}
32
33	impl dyn Array + '_ {
34	/// Serialize the array into a sequence of byte buffers that should be written contiguously.
35	/// This function returns a vec to avoid copying data buffers.
36	///
37	/// Optionally, padding can be included to guarantee buffer alignment and ensure zero-copy
38	/// reads within the context of an external file or stream. In this case, the alignment of
39	/// the first byte buffer should be respected when writing the buffers to the stream or file.
40	///
41	/// The format of this blob is a sequence of data buffers, possible with prefixed padding,
42	/// followed by a flatbuffer containing an [`fba::Array`] message, and ending with a
43	/// little-endian u32 describing the length of the flatbuffer message.
44	pub fn serialize(	638✔
45	&self,	638✔
46	ctx: &ArrayContext,	638✔
47	options: &SerializeOptions,	638✔
48	) -> VortexResult<Vec<ByteBuffer>> {	638✔
49	// Collect all array buffers
50	let array_buffers = self	638✔
51	.depth_first_traversal()	638✔
52	.flat_map(\|f\| f.buffers())	2,970✔
53	.collect::<Vec<_>>();	638✔
54
55	// Allocate result buffers, including a possible padding buffer for each.
56	let mut buffers = vec![];	638✔
57	let mut fb_buffers = Vec::with_capacity(buffers.capacity());	638✔
58
59	// If we're including padding, we need to find the maximum required buffer alignment.
60	let max_alignment = array_buffers	638✔
61	.iter()	638✔
62	.map(\|buf\| buf.alignment())	2,658✔
63	.chain(iter::once(FlatBuffer::alignment()))	638✔
64	.max()	638✔
65	.unwrap_or_else(FlatBuffer::alignment);	638✔
66
67	// Create a shared buffer of zeros we can use for padding
68	let zeros = ByteBuffer::zeroed(*max_alignment);	638✔
69
70	// We push an empty buffer with the maximum alignment, so then subsequent buffers
71	// will be aligned. For subsequent buffers, we always push a 1-byte alignment.
72	buffers.push(ByteBuffer::zeroed_aligned(0, max_alignment));	638✔
73
74	// Keep track of where we are in the "file" to calculate padding.
75	let mut pos = options.offset;	638✔
76
77	// Push all the array buffers with padding as necessary.
78	for buffer in array_buffers {	3,296✔
79	let padding = if options.include_padding {	2,658✔
80	let padding = pos.next_multiple_of(*buffer.alignment()) - pos;	2,658✔
81	if padding > 0 {	2,658✔
82	pos += padding;	460✔
83	buffers.push(zeros.slice(0..padding));	460✔
84	}	2,198✔
85	padding	2,658✔
86	} else {
UNCOV 87	0	×
88	};
89
90	fb_buffers.push(fba::Buffer::new(	2,658✔
91	u16::try_from(padding).vortex_expect("padding fits into u16"),	2,658✔
92	buffer.alignment().exponent(),	2,658✔
93	Compression::None,
94	u32::try_from(buffer.len())	2,658✔
95	.map_err(\|_\| vortex_err!("All buffers must fit into u32 for serialization"))?,	2,658✔
96	));
97
98	pos += buffer.len();	2,658✔
99	buffers.push(buffer.aligned(Alignment::none()));	2,658✔
100	}
101
102	// Set up the flatbuffer builder
103	let mut fbb = FlatBufferBuilder::new();	638✔
104	let root = ArrayNodeFlatBuffer::try_new(ctx, self)?;	638✔
105	let fb_root = root.write_flatbuffer(&mut fbb);	638✔
106	let fb_buffers = fbb.create_vector(&fb_buffers);	638✔
107	let fb_array = fba::Array::create(	638✔
108	&mut fbb,	638✔
109	&fba::ArrayArgs {	638✔
110	root: Some(fb_root),	638✔
111	buffers: Some(fb_buffers),	638✔
112	},	638✔
113	);
114	fbb.finish_minimal(fb_array);	638✔
115	let (fb_vec, fb_start) = fbb.collapse();	638✔
116	let fb_end = fb_vec.len();	638✔
117	let fb_buffer = ByteBuffer::from(fb_vec).slice(fb_start..fb_end);	638✔
118	let fb_length = fb_buffer.len();	638✔
119
120	if options.include_padding {	638✔
121	let padding = pos.next_multiple_of(*FlatBuffer::alignment()) - pos;	638✔
122	if padding > 0 {	638✔
123	buffers.push(zeros.slice(0..padding));	270✔
124	}	368✔
UNCOV 125	}	×
126	buffers.push(fb_buffer);	638✔
127
128	// Finally, we write down the u32 length for the flatbuffer.
129	buffers.push(ByteBuffer::from(	638✔
130	u32::try_from(fb_length)	638✔
131	.map_err(\|_\| vortex_err!("Array metadata flatbuffer must fit into u32 for serialization. Array encoding tree is too large."))?	638✔
132	.to_le_bytes()	638✔
133	.to_vec(),	638✔
134	));
135
136	Ok(buffers)	638✔
137	}	638✔
138	}
139
140	/// A utility struct for creating an [`fba::ArrayNode`] flatbuffer.
141	pub struct ArrayNodeFlatBuffer<'a> {
142	ctx: &'a ArrayContext,
143	array: &'a dyn Array,
144	buffer_idx: u16,
145	}
146
147	impl<'a> ArrayNodeFlatBuffer<'a> {
148	pub fn try_new(ctx: &'a ArrayContext, array: &'a dyn Array) -> VortexResult<Self> {	638✔
149	// Depth-first traversal of the array to ensure it supports serialization.
150	for child in array.depth_first_traversal() {	2,970✔
151	if child.metadata()?.is_none() {	2,970✔
152	vortex_bail!(	×
153	"Array {} does not support serialization",	×
154	child.encoding_id()	×
155	);
156	}	2,970✔
157	}
158	Ok(Self {	638✔
159	ctx,	638✔
160	array,	638✔
161	buffer_idx: 0,	638✔
162	})	638✔
163	}	638✔
164	}
165
166	impl FlatBufferRoot for ArrayNodeFlatBuffer<'_> {}
167
168	impl WriteFlatBuffer for ArrayNodeFlatBuffer<'_> {
169	type Target<'t> = fba::ArrayNode<'t>;
170
171	fn write_flatbuffer<'fb>(	2,970✔
172	&self,	2,970✔
173	fbb: &mut FlatBufferBuilder<'fb>,	2,970✔
174	) -> WIPOffset<Self::Target<'fb>> {	2,970✔
175	let encoding = self.ctx.encoding_idx(&self.array.encoding());	2,970✔
176	let metadata = self	2,970✔
177	.array	2,970✔
178	.metadata()	2,970✔
179	// TODO(ngates): add try_write_flatbuffer
180	.vortex_expect("Failed to serialize metadata")	2,970✔
181	.vortex_expect("Validated that all arrays support serialization");	2,970✔
182	let metadata = Some(fbb.create_vector(metadata.as_slice()));	2,970✔
183
184	// Assign buffer indices for all child arrays.
185	let nbuffers = u16::try_from(self.array.nbuffers())	2,970✔
186	.vortex_expect("Array can have at most u16::MAX buffers");	2,970✔
187	let mut child_buffer_idx = self.buffer_idx + nbuffers;	2,970✔
188
189	let children = &self	2,970✔
190	.array	2,970✔
191	.children()	2,970✔
192	.iter()	2,970✔
193	.map(\|child\| {	2,970✔
194	// Update the number of buffers required.
195	let msg = ArrayNodeFlatBuffer {	2,332✔
196	ctx: self.ctx,	2,332✔
197	array: child,	2,332✔
198	buffer_idx: child_buffer_idx,	2,332✔
199	}	2,332✔
200	.write_flatbuffer(fbb);	2,332✔
201	child_buffer_idx = u16::try_from(child.nbuffers_recursive())	2,332✔
202	.ok()	2,332✔
203	.and_then(\|nbuffers\| nbuffers.checked_add(child_buffer_idx))	2,332✔
204	.vortex_expect("Too many buffers (u16) for Array");	2,332✔
205	msg	2,332✔
206	})	2,332✔
207	.collect::<Vec<_>>();	2,970✔
208	let children = Some(fbb.create_vector(children));	2,970✔
209
210	let buffers = Some(fbb.create_vector_from_iter((0..nbuffers).map(\|i\| i + self.buffer_idx)));	2,970✔
211	let stats = Some(self.array.statistics().to_owned().write_flatbuffer(fbb));	2,970✔
212
213	fba::ArrayNode::create(	2,970✔
214	fbb,	2,970✔
215	&fba::ArrayNodeArgs {	2,970✔
216	encoding,	2,970✔
217	metadata,	2,970✔
218	children,	2,970✔
219	buffers,	2,970✔
220	stats,	2,970✔
221	},	2,970✔
222	)
223	}	2,970✔
224	}
225
226	/// To minimize the serialized form, arrays do not persist their own dtype and length. Instead,
227	/// parent arrays pass this information down during deserialization. This trait abstracts
228	/// over either a serialized [`crate::serde::ArrayParts`] or the
229	/// in-memory [`crate::data::ArrayData`].
230	pub trait ArrayChildren {
231	/// Returns the nth child of the array with the given dtype and length.
232	fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult<ArrayRef>;
233
234	/// The number of children.
235	fn len(&self) -> usize;
236
237	/// Returns true if there are no children.
238	fn is_empty(&self) -> bool {	3,056✔
239	self.len() == 0	3,056✔
240	}	3,056✔
241	}
242
243	/// [`ArrayParts`] represents a parsed but not-yet-decoded deserialized [`Array`].
244	/// It contains all the information from the serialized form, without anything extra. i.e.
245	/// it is missing a [`DType`] and `len`, and the `encoding_id` is not yet resolved to a concrete
246	/// vtable.
247	///
248	/// An [`ArrayParts`] can be fully decoded into an [`ArrayRef`] using the `decode` function.
249	#[derive(Clone)]
250	pub struct ArrayParts {
251	// Typed as fb::ArrayNode
252	flatbuffer: FlatBuffer,
253	// The location of the current fb::ArrayNode
254	flatbuffer_loc: usize,
255	buffers: Arc<[ByteBuffer]>,
256	}
257
258	impl Debug for ArrayParts {
259	fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {	×
260	f.debug_struct("ArrayParts")	×
261	.field("encoding_id", &self.encoding_id())	×
262	.field("children", &(0..self.nchildren()).map(\|i\| self.child(i)))	×
263	.field(	×
264	"buffers",	×
265	&(0..self.nbuffers()).map(\|i\| self.buffer(i).ok()),	×
266	)
267	.field("metadata", &self.metadata())	×
268	.finish()	×
269	}	×
270	}
271
272	impl ArrayParts {
273	/// Decode an [`ArrayParts`] into an [`ArrayRef`].
274	pub fn decode(&self, ctx: &ArrayContext, dtype: &DType, len: usize) -> VortexResult<ArrayRef> {	10,720✔
275	let encoding_id = self.flatbuffer().encoding();	10,720✔
276	let vtable = ctx	10,720✔
277	.lookup_encoding(encoding_id)	10,720✔
278	.ok_or_else(\|\| vortex_err!("Unknown encoding: {}", encoding_id))?;	10,720✔
279
280	let buffers: Vec<_> = (0..self.nbuffers())	10,720✔
281	.map(\|idx\| self.buffer(idx))	10,720✔
282	.try_collect()?;	10,720✔
283
284	let children = ArrayPartsChildren { parts: self, ctx };	10,720✔
285
286	let decoded = vtable.build(dtype, len, self.metadata(), &buffers, &children)?;	10,720✔
287
288	assert_eq!(	10,720✔
289	decoded.len(),	10,720✔
290	len,
291	"Array decoded from {} has incorrect length {}, expected {}",	×
292	vtable.id(),	×
293	decoded.len(),	×
294	len
295	);
296	assert_eq!(	10,720✔
297	decoded.dtype(),	10,720✔
298	dtype,
299	"Array decoded from {} has incorrect dtype {}, expected {}",	×
300	vtable.id(),	×
301	decoded.dtype(),	×
302	dtype,
303	);
304	assert_eq!(	10,720✔
305	decoded.encoding_id(),	10,720✔
306	vtable.id(),	10,720✔
307	"Array decoded from {} has incorrect encoding {}",	×
308	vtable.id(),	×
309	decoded.encoding_id(),	×
310	);
311
312	// Populate statistics from the serialized array.
313	if let Some(stats) = self.flatbuffer().stats() {	10,720✔
314	let decoded_statistics = decoded.statistics();	10,720✔
315	StatsSet::read_flatbuffer(&stats)?	10,720✔
316	.into_iter()	10,720✔
317	.for_each(\|(stat, val)\| decoded_statistics.set(stat, val));	10,720✔
318	}	×
319
320	Ok(decoded)	10,720✔
321	}	10,720✔
322
323	/// Returns the array encoding.
324	pub fn encoding_id(&self) -> u16 {	×
325	self.flatbuffer().encoding()	×
326	}	×
327
328	/// Returns the array metadata bytes.
329	pub fn metadata(&self) -> &[u8] {	10,720✔
330	self.flatbuffer()	10,720✔
331	.metadata()	10,720✔
332	.map(\|metadata\| metadata.bytes())	10,720✔
333	.unwrap_or(&[])	10,720✔
334	}	10,720✔
335
336	/// Returns the number of children.
337	pub fn nchildren(&self) -> usize {	8,650✔
338	self.flatbuffer()	8,650✔
339	.children()	8,650✔
340	.map_or(0, \|children\| children.len())	8,650✔
341	}	8,650✔
342
343	/// Returns the nth child of the array.
344	pub fn child(&self, idx: usize) -> ArrayParts {	7,238✔
345	let children = self	7,238✔
346	.flatbuffer()	7,238✔
347	.children()	7,238✔
348	.vortex_expect("Expected array to have children");	7,238✔
349	if idx >= children.len() {	7,238✔
350	vortex_panic!(	×
351	"Invalid child index {} for array with {} children",	×
352	idx,
353	children.len()	×
354	);
355	}	7,238✔
356	self.with_root(children.get(idx))	7,238✔
357	}	7,238✔
358
359	/// Returns the number of buffers.
360	pub fn nbuffers(&self) -> usize {	10,720✔
361	self.flatbuffer()	10,720✔
362	.buffers()	10,720✔
363	.map_or(0, \|buffers\| buffers.len())	10,720✔
364	}	10,720✔
365
366	/// Returns the nth buffer of the current array.
367	pub fn buffer(&self, idx: usize) -> VortexResult<ByteBuffer> {	6,924✔
368	let buffer_idx = self	6,924✔
369	.flatbuffer()	6,924✔
370	.buffers()	6,924✔
371	.ok_or_else(\|\| vortex_err!("Array has no buffers"))?	6,924✔
372	.get(idx);	6,924✔
373	self.buffers	6,924✔
374	.get(buffer_idx as usize)	6,924✔
375	.cloned()	6,924✔
376	.ok_or_else(\|\| {	6,924✔
377	vortex_err!(	×
378	"Invalid buffer index {} for array with {} buffers",	×
379	buffer_idx,
380	self.nbuffers()	×
381	)
382	})	×
383	}	6,924✔
384
385	/// Returns the root ArrayNode flatbuffer.
386	fn flatbuffer(&self) -> fba::ArrayNode<'_> {	65,692✔
387	unsafe { fba::ArrayNode::follow(self.flatbuffer.as_ref(), self.flatbuffer_loc) }	65,692✔
388	}	65,692✔
389
390	/// Returns a new [`ArrayParts`] with the given node as the root
391	// TODO(ngates): we may want a wrapper that avoids this clone.
392	fn with_root(&self, root: fba::ArrayNode) -> Self {	7,238✔
393	let mut this = self.clone();	7,238✔
394	this.flatbuffer_loc = root._tab.loc();	7,238✔
395	this	7,238✔
396	}	7,238✔
397	}
398
399	struct ArrayPartsChildren<'a> {
400	parts: &'a ArrayParts,
401	ctx: &'a ArrayContext,
402	}
403
404	impl ArrayChildren for ArrayPartsChildren<'_> {
405	fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult<ArrayRef> {	7,238✔
406	self.parts.child(index).decode(self.ctx, dtype, len)	7,238✔
407	}	7,238✔
408
409	fn len(&self) -> usize {	8,650✔
410	self.parts.nchildren()	8,650✔
411	}	8,650✔
412	}
413
414	impl TryFrom<ByteBuffer> for ArrayParts {
415	type Error = VortexError;
416
417	fn try_from(value: ByteBuffer) -> Result<Self, Self::Error> {	3,482✔
418	// The final 4 bytes contain the length of the flatbuffer.
419	if value.len() < 4 {	3,482✔
420	vortex_bail!("ArrayParts buffer is too short");	×
421	}	3,482✔
422
423	// We align each buffer individually, so we remove alignment requirements on the buffer.
424	let value = value.aligned(Alignment::none());	3,482✔
425
426	let fb_length = u32::try_from_le_bytes(&value.as_slice()[value.len() - 4..])? as usize;	3,482✔
427	if value.len() < 4 + fb_length {	3,482✔
428	vortex_bail!("ArrayParts buffer is too short for flatbuffer");	×
429	}	3,482✔
430
431	let fb_offset = value.len() - 4 - fb_length;	3,482✔
432	let fb_buffer = value.slice(fb_offset..fb_offset + fb_length);	3,482✔
433	let fb_buffer = FlatBuffer::align_from(fb_buffer);	3,482✔
434
435	let fb_array = root::<fba::Array>(fb_buffer.as_ref())?;	3,482✔
436	let fb_root = fb_array.root().vortex_expect("Array must have a root node");	3,482✔
437
438	let mut offset = 0;	3,482✔
439	let buffers: Arc<[ByteBuffer]> = fb_array	3,482✔
440	.buffers()	3,482✔
441	.unwrap_or_default()	3,482✔
442	.iter()	3,482✔
443	.map(\|fb_buffer\| {	6,924✔
444	// Skip padding
445	offset += fb_buffer.padding() as usize;	6,924✔
446
447	let buffer_len = fb_buffer.length() as usize;	6,924✔
448
449	// Extract a buffer and ensure it's aligned, copying if necessary
450	let buffer = value	6,924✔
451	.slice(offset..(offset + buffer_len))	6,924✔
452	.aligned(Alignment::from_exponent(fb_buffer.alignment_exponent()));	6,924✔
453
454	offset += buffer_len;	6,924✔
455	buffer	6,924✔
456	})	6,924✔
457	.collect();	3,482✔
458
459	Ok(ArrayParts {	3,482✔
460	flatbuffer: fb_buffer.clone(),	3,482✔
461	flatbuffer_loc: fb_root._tab.loc(),	3,482✔
462	buffers,	3,482✔
463	})	3,482✔
464	}	3,482✔
465	}

vortex-data / vortex / 16728097825

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous