vortex_array/
data.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use 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}