nd_arrow_array/
array.rs

1//! A single N-dimensional Arrow array.
2
3use arrow::array::{Array, ArrayRef, new_null_array};
4use arrow::datatypes::DataType;
5
6use crate::attributes::Attributes;
7use crate::dimensions::Dimensions;
8use crate::error::{Result, nd_err};
9
10/// A flat Arrow array plus named C-order dimensions and attributes.
11///
12/// Invariant: `dims.num_elements() == values.len()`. Validity travels in the
13/// null buffer of the values. There is no separate mask or fill value, so
14/// readers map fill values to nulls when they build the values.
15#[derive(Debug, Clone)]
16pub struct NdArrowArray {
17    values: ArrayRef,
18    dims: Dimensions,
19    attrs: Attributes,
20}
21
22impl PartialEq for NdArrowArray {
23    fn eq(&self, other: &Self) -> bool {
24        self.dims == other.dims
25            && self.values.as_ref() == other.values.as_ref()
26            && self.attrs == other.attrs
27    }
28}
29
30impl NdArrowArray {
31    /// Create an array. The length of `values` must equal the number of elements of `dims`.
32    pub fn try_new(values: ArrayRef, dims: Dimensions) -> Result<Self> {
33        if dims.checked_num_elements() != Some(values.len()) {
34            return Err(nd_err!(
35                ShapeMismatch,
36                "values length {} does not match dimensions {dims} ({} elements)",
37                values.len(),
38                dims.num_elements()
39            ));
40        }
41        Ok(Self::new_unchecked(values, dims, Attributes::new()))
42    }
43
44    /// A rank-0 array. `value` must have length 1.
45    pub fn try_new_scalar(value: ArrayRef) -> Result<Self> {
46        Self::try_new(value, Dimensions::scalar())
47    }
48
49    /// A rank-0 array holding one null of `data_type`.
50    pub fn new_null_scalar(data_type: &DataType) -> Self {
51        Self::new_unchecked(
52            new_null_array(data_type, 1),
53            Dimensions::scalar(),
54            Attributes::new(),
55        )
56    }
57
58    pub(crate) fn new_unchecked(values: ArrayRef, dims: Dimensions, attrs: Attributes) -> Self {
59        debug_assert_eq!(dims.num_elements(), values.len());
60        Self {
61            values,
62            dims,
63            attrs,
64        }
65    }
66
67    /// Replace the attributes.
68    pub fn with_attributes(mut self, attrs: Attributes) -> Self {
69        self.attrs = attrs;
70        self
71    }
72
73    /// The flat values in C order.
74    pub fn values(&self) -> &ArrayRef {
75        &self.values
76    }
77
78    /// The named dimensions.
79    pub fn dims(&self) -> &Dimensions {
80        &self.dims
81    }
82
83    /// The attributes.
84    pub fn attributes(&self) -> &Attributes {
85        &self.attrs
86    }
87
88    /// The attributes, for changes in place.
89    pub fn attributes_mut(&mut self) -> &mut Attributes {
90        &mut self.attrs
91    }
92
93    /// The Arrow type of the values.
94    pub fn data_type(&self) -> &DataType {
95        self.values.data_type()
96    }
97
98    /// The size of each axis.
99    pub fn shape(&self) -> Vec<usize> {
100        self.dims.shape()
101    }
102
103    /// The number of axes.
104    pub fn rank(&self) -> usize {
105        self.dims.rank()
106    }
107
108    /// Number of logical elements, equal to the length of the values.
109    pub fn len(&self) -> usize {
110        self.values.len()
111    }
112
113    /// True when the array has no elements.
114    pub fn is_empty(&self) -> bool {
115        self.values.is_empty()
116    }
117
118    /// True for rank 0.
119    pub fn is_scalar(&self) -> bool {
120        self.dims.is_scalar()
121    }
122
123    /// Replace the dimensions without a copy of the values.
124    /// The new dimensions must hold the same number of elements.
125    pub fn reshape(self, dims: Dimensions) -> Result<Self> {
126        let attrs = self.attrs;
127        Ok(Self::try_new(self.values, dims)?.with_attributes(attrs))
128    }
129
130    /// Split into values, dimensions and attributes.
131    pub fn into_parts(self) -> (ArrayRef, Dimensions, Attributes) {
132        (self.values, self.dims, self.attrs)
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use std::sync::Arc;
139
140    use arrow::array::Int32Array;
141
142    use super::*;
143
144    #[test]
145    fn len_mismatch_rejected() {
146        let dims = Dimensions::try_new([("x", 2)]).unwrap();
147        let result = NdArrowArray::try_new(Arc::new(Int32Array::from(vec![1, 2, 3])), dims);
148        assert!(result.is_err());
149    }
150
151    #[test]
152    fn reshape_keeps_values_and_attributes() {
153        let a = NdArrowArray::try_new(
154            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])),
155            Dimensions::try_new([("x", 6)]).unwrap(),
156        )
157        .unwrap()
158        .with_attributes(Attributes::new().with("units", "m"));
159
160        let b = a
161            .clone()
162            .reshape(Dimensions::try_new([("y", 2), ("x", 3)]).unwrap())
163            .unwrap();
164        assert_eq!(b.shape(), vec![2, 3]);
165        assert_eq!(b.values(), a.values());
166        assert_eq!(b.attributes().get_str("units"), Some("m"));
167        assert!(a.reshape(Dimensions::try_new([("x", 4)]).unwrap()).is_err());
168    }
169}
nd_arrow_array/array.rs

nd_arrow_array/
array.rs
