vortex_array/
data.rs

1use itertools::Itertools;
2use vortex_buffer::ByteBuffer;
3use vortex_dtype::DType;
4use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
5
6use crate::serde::ArrayChildren;
7use crate::{Array, ArrayRef, EncodingRef};
8
9#[derive(Clone, Debug)]
10pub struct ArrayData {
11    encoding: EncodingRef,
12    len: usize,
13    dtype: DType,
14    metadata: Vec<u8>,
15    buffers: Vec<ByteBuffer>,
16    children: Vec<ArrayData>,
17}
18
19impl ArrayChildren for ArrayData {
20    fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult<ArrayRef> {
21        if index >= self.children.len() {
22            vortex_bail!("Index out of bounds");
23        }
24        let child = &self.children[index];
25        if child.len != len {
26            vortex_bail!(
27                "Child length mismatch. Provided {}, but actually {}",
28                len,
29                child.len
30            );
31        }
32        if child.dtype != *dtype {
33            vortex_bail!(
34                "Child dtype mismatch. Provided {:?}, but actually {:?}",
35                dtype,
36                child.dtype
37            );
38        }
39        ArrayRef::try_from(child)
40    }
41
42    fn len(&self) -> usize {
43        self.children.len()
44    }
45}
46
47impl TryFrom<&dyn Array> for ArrayData {
48    type Error = VortexError;
49
50    fn try_from(value: &dyn Array) -> Result<Self, Self::Error> {
51        Ok(ArrayData {
52            encoding: value.encoding(),
53            len: value.len(),
54            dtype: value.dtype().clone(),
55            metadata: value
56                .metadata()?
57                .ok_or_else(|| vortex_err!("Array does not support serialization"))?,
58            buffers: value.buffers().to_vec(),
59            children: value
60                .children()
61                .iter()
62                .map(|child| child.as_ref().try_into())
63                .try_collect()?,
64        })
65    }
66}
67
68impl TryFrom<&ArrayData> for ArrayRef {
69    type Error = VortexError;
70
71    fn try_from(value: &ArrayData) -> Result<Self, Self::Error> {
72        value.encoding.build(
73            &value.dtype,
74            value.len,
75            &value.metadata,
76            &value.buffers,
77            value,
78        )
79    }
80}