nd_arrow_array/
ndarray_interop.rs

1//! Conversion to and from [`ndarray::ArrayD`], with the `ndarray` feature.
2//!
3//! `ndarray` arrays have no dimension names and no nulls, so a conversion
4//! from an nd array needs values without nulls, and a conversion to an nd
5//! array needs dimension names.
6
7use std::sync::Arc;
8
9use arrow::array::{Array, ArrowPrimitiveType, PrimitiveArray};
10use ndarray::{ArrayD, IxDyn};
11
12use crate::NdArrowArray;
13use crate::dimensions::Dimensions;
14use crate::error::{Result, nd_err};
15
16impl NdArrowArray {
17    /// Copy the values into an `ndarray` array of the same shape.
18    /// The values must be of type `T` and without nulls.
19    pub fn to_ndarray<T: ArrowPrimitiveType>(&self) -> Result<ArrayD<T::Native>> {
20        let values = self
21            .values()
22            .as_any()
23            .downcast_ref::<PrimitiveArray<T>>()
24            .ok_or_else(|| {
25                nd_err!(
26                    InvalidArgument,
27                    "values have type {}, not {}",
28                    self.data_type(),
29                    T::DATA_TYPE
30                )
31            })?;
32        if values.null_count() > 0 {
33            return Err(nd_err!(
34                NotSupported,
35                "ndarray has no nulls; fill the nulls first"
36            ));
37        }
38        ArrayD::from_shape_vec(IxDyn(&self.shape()), values.values().to_vec())
39            .map_err(|e| nd_err!(ShapeMismatch, "{e}"))
40    }
41
42    /// Build an nd array from an `ndarray` array in any memory layout,
43    /// with one dimension name per axis.
44    pub fn from_ndarray<T: ArrowPrimitiveType>(
45        array: &ArrayD<T::Native>,
46        names: &[&str],
47    ) -> Result<Self> {
48        if names.len() != array.ndim() {
49            return Err(nd_err!(
50                InvalidArgument,
51                "{} names for an array of rank {}",
52                names.len(),
53                array.ndim()
54            ));
55        }
56        let dims = Dimensions::try_new(names.iter().copied().zip(array.shape().iter().copied()))?;
57        let values: PrimitiveArray<T> = PrimitiveArray::from_iter_values(array.iter().copied());
58        NdArrowArray::try_new(Arc::new(values), dims)
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use arrow::datatypes::Float64Type;
65    use ndarray::{Array2, ShapeBuilder};
66
67    use super::*;
68
69    #[test]
70    fn round_trip_keeps_logical_order() {
71        let fortran = Array2::from_shape_vec((2, 3).f(), vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0])
72            .unwrap()
73            .into_dyn();
74        let nd = NdArrowArray::from_ndarray::<Float64Type>(&fortran, &["y", "x"]).unwrap();
75        assert_eq!(nd.shape(), vec![2, 3]);
76        let back = nd.to_ndarray::<Float64Type>().unwrap();
77        assert_eq!(back, fortran);
78        assert!(NdArrowArray::from_ndarray::<Float64Type>(&fortran, &["y"]).is_err());
79    }
80}
nd_arrow_array/ndarray_interop.rs

nd_arrow_array/
ndarray_interop.rs
