nd_arrow_array/
display.rs

1//! A NumPy-like text form of an [`NdArrowArray`].
2//!
3//! ```text
4//! <NdArrowArray Float64 (time: 2, lat: 3)>
5//! [[0.0, 0.1, 0.2],
6//!  [1.0, 1.1, 1.2]]
7//! Attributes:
8//!     units: K
9//! ```
10//!
11//! Like NumPy, an axis longer than [`EDGE_ITEMS`] * 2 shows only its first
12//! and last items when the array has more than [`SUMMARY_THRESHOLD`] values.
13
14use std::fmt;
15
16use arrow::util::display::{ArrayFormatter, FormatOptions};
17
18use crate::NdArrowArray;
19
20/// Arrays with more values than this are summarized, like NumPy's `threshold`.
21pub const SUMMARY_THRESHOLD: usize = 1000;
22/// Items shown at each end of a summarized axis, like NumPy's `edgeitems`.
23pub const EDGE_ITEMS: usize = 3;
24
25impl fmt::Display for NdArrowArray {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        let dims: Vec<String> = self
28            .dims()
29            .iter()
30            .map(|d| format!("{}: {}", d.name(), d.size()))
31            .collect();
32        writeln!(
33            f,
34            "<NdArrowArray {} ({})>",
35            self.data_type(),
36            dims.join(", ")
37        )?;
38
39        let options = FormatOptions::default().with_null("null");
40        let formatter =
41            ArrayFormatter::try_new(self.values().as_ref(), &options).map_err(|_| fmt::Error)?;
42        let value = |i: usize| formatter.value(i).to_string();
43        let summarize = self.len() > SUMMARY_THRESHOLD;
44        let shape = self.shape();
45        let strides = self.dims().c_strides();
46        if shape.is_empty() {
47            write!(f, "{}", value(0))?;
48        } else {
49            write_axis(f, &value, &shape, &strides, 0, 0, summarize)?;
50        }
51
52        if !self.attributes().is_empty() {
53            write!(f, "\nAttributes:")?;
54            for (key, attr) in self.attributes().iter() {
55                let attr_formatter =
56                    ArrayFormatter::try_new(attr.as_ref(), &options).map_err(|_| fmt::Error)?;
57                write!(f, "\n    {key}: {}", attr_formatter.value(0))?;
58            }
59        }
60        Ok(())
61    }
62}
63
64fn write_axis(
65    f: &mut fmt::Formatter<'_>,
66    value: &impl Fn(usize) -> String,
67    shape: &[usize],
68    strides: &[usize],
69    axis: usize,
70    offset: usize,
71    summarize: bool,
72) -> fmt::Result {
73    let size = shape[axis];
74    let indices: Vec<Option<usize>> = if summarize && size > 2 * EDGE_ITEMS {
75        (0..EDGE_ITEMS)
76            .map(Some)
77            .chain(std::iter::once(None))
78            .chain((size - EDGE_ITEMS..size).map(Some))
79            .collect()
80    } else {
81        (0..size).map(Some).collect()
82    };
83    let last = axis + 1 == shape.len();
84    write!(f, "[")?;
85    for (k, index) in indices.iter().enumerate() {
86        if k > 0 {
87            if last {
88                write!(f, ", ")?;
89            } else {
90                write!(f, ",{}", "\n".repeat(shape.len() - axis - 1))?;
91                write!(f, "{}", " ".repeat(axis + 1))?;
92            }
93        }
94        match index {
95            None => write!(f, "...")?,
96            Some(i) if last => write!(f, "{}", value(offset + i * strides[axis]))?,
97            Some(i) => write_axis(
98                f,
99                value,
100                shape,
101                strides,
102                axis + 1,
103                offset + i * strides[axis],
104                summarize,
105            )?,
106        }
107    }
108    write!(f, "]")
109}
110
111#[cfg(test)]
112mod tests {
113    use std::sync::Arc;
114
115    use arrow::array::{Float64Array, Int32Array};
116
117    use crate::{Attributes, Dimensions, NdArrowArray};
118
119    #[test]
120    fn prints_like_numpy() {
121        let a = NdArrowArray::try_new(
122            Arc::new(Float64Array::from(vec![
123                Some(0.0),
124                Some(0.5),
125                None,
126                Some(1.0),
127                Some(1.5),
128                Some(2.0),
129            ])),
130            Dimensions::try_new([("time", 2), ("lat", 3)]).unwrap(),
131        )
132        .unwrap()
133        .with_attributes(Attributes::new().with("units", "K"));
134        assert_eq!(
135            a.to_string(),
136            "<NdArrowArray Float64 (time: 2, lat: 3)>\n[[0.0, 0.5, null],\n [1.0, 1.5, 2.0]]\nAttributes:\n    units: K"
137        );
138    }
139
140    #[test]
141    fn summarizes_large_arrays() {
142        let a = NdArrowArray::try_new(
143            Arc::new(Int32Array::from_iter_values(0..2000)),
144            Dimensions::try_new([("x", 2000)]).unwrap(),
145        )
146        .unwrap();
147        assert!(a.to_string().ends_with("[0, 1, 2, ..., 1997, 1998, 1999]"));
148        let s = NdArrowArray::try_new_scalar(Arc::new(Int32Array::from(vec![7]))).unwrap();
149        assert_eq!(s.to_string(), "<NdArrowArray Int32 ()>\n7");
150    }
151}
nd_arrow_array/display.rs

nd_arrow_array/
display.rs
