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}