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