nd_arrow_array/
dimensions.rs

1//! Named dimensions for N-dimensional Arrow arrays.
2//!
3//! A [`Dimensions`] value describes the logical shape of an
4//! [`NdArrowArray`](crate::NdArrowArray) in C-order (row-major, first dimension
5//! outermost). Dimension names drive xarray-style broadcasting: axes align by
6//! name, never by position.
7
8use std::fmt;
9use std::sync::Arc;
10
11use crate::error::{Result, nd_err};
12
13/// A single named axis with a fixed size.
14#[derive(Debug, Clone, PartialEq, Eq, Hash)]
15pub struct Dimension {
16    name: Arc<str>,
17    size: usize,
18}
19
20impl Dimension {
21    /// A dimension named `name` with `size` positions.
22    pub fn new(name: impl Into<Arc<str>>, size: usize) -> Self {
23        Self {
24            name: name.into(),
25            size,
26        }
27    }
28
29    /// The name.
30    pub fn name(&self) -> &str {
31        &self.name
32    }
33
34    /// The number of positions.
35    pub fn size(&self) -> usize {
36        self.size
37    }
38}
39
40impl fmt::Display for Dimension {
41    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42        write!(f, "{}={}", self.name, self.size)
43    }
44}
45
46impl<S: Into<Arc<str>>> From<(S, usize)> for Dimension {
47    fn from((name, size): (S, usize)) -> Self {
48        Self::new(name, size)
49    }
50}
51
52/// An ordered list of named axes, C-order. Cheap to clone.
53///
54/// Rank 0 (no axes) describes a scalar: one logical element.
55#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
56pub struct Dimensions {
57    dims: Arc<[Dimension]>,
58}
59
60impl Dimensions {
61    /// Create dimensions, validating that axis names are non-empty and unique.
62    pub fn try_new(dims: impl IntoIterator<Item = impl Into<Dimension>>) -> Result<Self> {
63        let dims: Vec<Dimension> = dims.into_iter().map(Into::into).collect();
64        for (i, dim) in dims.iter().enumerate() {
65            if dim.name().is_empty() {
66                return Err(nd_err!(
67                    InvalidDimensions,
68                    "dimension at axis {i} has an empty name"
69                ));
70            }
71            if dims[..i].iter().any(|d| d.name() == dim.name()) {
72                return Err(nd_err!(
73                    InvalidDimensions,
74                    "duplicate dimension name '{}'",
75                    dim.name()
76                ));
77            }
78        }
79        Ok(Self { dims: dims.into() })
80    }
81
82    /// Rank-0 dimensions describing a scalar.
83    pub fn scalar() -> Self {
84        Self::default()
85    }
86
87    /// The number of axes.
88    pub fn rank(&self) -> usize {
89        self.dims.len()
90    }
91
92    /// True for rank 0.
93    pub fn is_scalar(&self) -> bool {
94        self.dims.is_empty()
95    }
96
97    /// The axes, outermost first.
98    pub fn iter(&self) -> std::slice::Iter<'_, Dimension> {
99        self.dims.iter()
100    }
101
102    /// The axes as a slice.
103    pub fn as_slice(&self) -> &[Dimension] {
104        &self.dims
105    }
106
107    /// The dimension at `axis`, if present.
108    pub fn get(&self, axis: usize) -> Option<&Dimension> {
109        self.dims.get(axis)
110    }
111
112    /// The size of each axis.
113    pub fn shape(&self) -> Vec<usize> {
114        self.dims.iter().map(Dimension::size).collect()
115    }
116
117    /// The name of each axis.
118    pub fn names(&self) -> Vec<&str> {
119        self.dims.iter().map(Dimension::name).collect()
120    }
121
122    /// Total number of logical elements (1 for a scalar, 0 if any axis is empty).
123    pub fn num_elements(&self) -> usize {
124        self.dims.iter().map(Dimension::size).product()
125    }
126
127    /// Like [`Self::num_elements`], but `None` on `usize` overflow.
128    pub fn checked_num_elements(&self) -> Option<usize> {
129        self.dims
130            .iter()
131            .try_fold(1usize, |acc, d| acc.checked_mul(d.size()))
132    }
133
134    /// The axis of `name`, if present.
135    pub fn position(&self, name: &str) -> Option<usize> {
136        self.dims.iter().position(|d| d.name() == name)
137    }
138
139    /// True when dimension `name` exists.
140    pub fn contains(&self, name: &str) -> bool {
141        self.position(name).is_some()
142    }
143
144    /// The axis of `name`, or an error that names the dimensions.
145    pub fn axis(&self, name: &str) -> Result<usize> {
146        self.position(name)
147            .ok_or_else(|| nd_err!(InvalidArgument, "dimension '{name}' is not in {self}"))
148    }
149
150    /// The size of dimension `name`, if present.
151    pub fn size_of(&self, name: &str) -> Option<usize> {
152        self.position(name).map(|axis| self.dims[axis].size())
153    }
154
155    /// The broadcast result of several dimension sets, xarray style.
156    ///
157    /// Axes align by name. The order follows the first appearance of each name.
158    /// Sizes must be equal, or one of them must be 1.
159    pub fn broadcast(all: &[&Dimensions]) -> Result<Dimensions> {
160        let mut out: Vec<Dimension> = Vec::new();
161        for dims in all {
162            for dim in dims.iter() {
163                match out.iter_mut().find(|d| d.name() == dim.name()) {
164                    None => out.push(dim.clone()),
165                    Some(existing) if existing.size() == dim.size() || dim.size() == 1 => {}
166                    Some(existing) if existing.size() == 1 => *existing = dim.clone(),
167                    Some(existing) => {
168                        return Err(nd_err!(
169                            ShapeMismatch,
170                            "dimension '{}' has incompatible sizes {} and {}",
171                            dim.name(),
172                            existing.size(),
173                            dim.size()
174                        ));
175                    }
176                }
177            }
178        }
179        Dimensions::try_new(out)
180    }
181
182    /// Fails when an array of these dimensions cannot broadcast onto `target`:
183    /// each axis must be in `target`, with the same size or with size 1.
184    pub fn check_broadcast_to(&self, target: &Dimensions) -> Result<()> {
185        for dim in self.iter() {
186            let Some(axis) = target.position(dim.name()) else {
187                return Err(nd_err!(
188                    ShapeMismatch,
189                    "cannot broadcast {self} to {target}: dimension '{}' is missing from the target",
190                    dim.name()
191                ));
192            };
193            let target_size = target.dims[axis].size();
194            if dim.size() != target_size && dim.size() != 1 {
195                return Err(nd_err!(
196                    ShapeMismatch,
197                    "cannot broadcast {self} to {target}: dimension '{}' has size {} but the target expects {target_size}",
198                    dim.name(),
199                    dim.size(),
200                ));
201            }
202        }
203        Ok(())
204    }
205
206    /// C-order (row-major) strides, in elements.
207    pub fn c_strides(&self) -> Vec<usize> {
208        let mut strides = vec![0usize; self.rank()];
209        let mut acc = 1usize;
210        for axis in (0..self.rank()).rev() {
211            strides[axis] = acc;
212            acc = acc.saturating_mul(self.dims[axis].size());
213        }
214        strides
215    }
216}
217
218impl<'a> IntoIterator for &'a Dimensions {
219    type Item = &'a Dimension;
220    type IntoIter = std::slice::Iter<'a, Dimension>;
221
222    fn into_iter(self) -> Self::IntoIter {
223        self.iter()
224    }
225}
226
227impl fmt::Display for Dimensions {
228    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
229        write!(f, "(")?;
230        for (i, dim) in self.dims.iter().enumerate() {
231            if i > 0 {
232                write!(f, ", ")?;
233            }
234            write!(f, "{dim}")?;
235        }
236        write!(f, ")")
237    }
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243
244    #[test]
245    fn strides_are_c_order() {
246        let d = Dimensions::try_new([("time", 2), ("lat", 3), ("lon", 4)]).unwrap();
247        assert_eq!(d.c_strides(), vec![12, 4, 1]);
248        assert_eq!(d.num_elements(), 24);
249        assert_eq!(d.names(), vec!["time", "lat", "lon"]);
250    }
251
252    #[test]
253    fn scalar_has_one_element() {
254        let d = Dimensions::scalar();
255        assert_eq!(d.rank(), 0);
256        assert_eq!(d.num_elements(), 1);
257        assert!(d.c_strides().is_empty());
258    }
259
260    #[test]
261    fn invalid_names_rejected() {
262        assert!(Dimensions::try_new([("x", 2), ("x", 3)]).is_err());
263        assert!(Dimensions::try_new([("", 2)]).is_err());
264    }
265
266    #[test]
267    fn checked_num_elements_detects_overflow() {
268        let d = Dimensions::try_new([("a", usize::MAX), ("b", 2)]).unwrap();
269        assert_eq!(d.checked_num_elements(), None);
270    }
271
272    #[test]
273    fn broadcast_check_follows_the_rule() {
274        let target = Dimensions::try_new([("time", 2), ("lat", 3)]).unwrap();
275        let ok = |spec: &[(&str, usize)]| {
276            Dimensions::try_new(spec.iter().copied())
277                .unwrap()
278                .check_broadcast_to(&target)
279        };
280        assert!(ok(&[("lat", 3)]).is_ok());
281        assert!(ok(&[("time", 1)]).is_ok());
282        assert!(Dimensions::scalar().check_broadcast_to(&target).is_ok());
283        let missing = ok(&[("depth", 2)]).unwrap_err();
284        assert!(
285            missing.to_string().contains("'depth' is missing"),
286            "{missing}"
287        );
288        let size = ok(&[("lat", 2)]).unwrap_err();
289        assert!(size.to_string().contains("has size 2"), "{size}");
290    }
291}
nd_arrow_array/dimensions.rs

nd_arrow_array/
dimensions.rs
