nd_arrow_array/
attributes.rs

1//! Per-array attributes (CF/netCDF style metadata such as `units`).
2//!
3//! Each attribute is a named Arrow value of length 1, so an attribute can be
4//! any Arrow type: a string, a number, a list (`valid_range`), and so on. In the
5//! encoding, the attributes of a column become the `attrs` struct: one child
6//! field per attribute key.
7
8use std::sync::Arc;
9
10use arrow::array::{
11    Array, ArrayRef, BooleanArray, ListArray, PrimitiveArray, StringArray, StructArray,
12    new_empty_array, new_null_array,
13};
14use arrow::buffer::OffsetBuffer;
15use arrow::compute::concat;
16use arrow::datatypes::{
17    ArrowPrimitiveType, DataType, Field, FieldRef, Fields, Float32Type, Float64Type, Int8Type,
18    Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
19};
20
21use crate::error::{Result, nd_err};
22
23/// An ordered set of named, length-1 Arrow values.
24#[derive(Debug, Clone, Default)]
25pub struct Attributes {
26    entries: Vec<(Arc<str>, ArrayRef)>,
27}
28
29impl Attributes {
30    /// Empty attributes.
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    /// The number of attributes.
36    pub fn len(&self) -> usize {
37        self.entries.len()
38    }
39
40    /// True when there are no attributes.
41    pub fn is_empty(&self) -> bool {
42        self.entries.is_empty()
43    }
44
45    /// The attribute value as a length-1 array.
46    pub fn get(&self, name: &str) -> Option<&ArrayRef> {
47        self.entries
48            .iter()
49            .find(|(key, _)| key.as_ref() == name)
50            .map(|(_, value)| value)
51    }
52
53    /// True when attribute `name` exists.
54    pub fn contains(&self, name: &str) -> bool {
55        self.get(name).is_some()
56    }
57
58    /// Convenience access for a string attribute.
59    pub fn get_str(&self, name: &str) -> Option<&str> {
60        let value = self.get(name)?.as_any().downcast_ref::<StringArray>()?;
61        value.is_valid(0).then(|| value.value(0))
62    }
63
64    /// Convenience access for a primitive attribute of type `T`.
65    pub fn get_primitive<T: ArrowPrimitiveType>(&self, name: &str) -> Option<T::Native> {
66        let value = self
67            .get(name)?
68            .as_any()
69            .downcast_ref::<PrimitiveArray<T>>()?;
70        value.is_valid(0).then(|| value.value(0))
71    }
72
73    /// Insert or replace an attribute.
74    pub fn insert(&mut self, name: impl Into<Arc<str>>, value: impl IntoScalar) {
75        self.insert_unchecked(name.into(), value.into_scalar());
76    }
77
78    /// Insert or replace an attribute from an Arrow array of length 1.
79    pub fn try_insert_array(&mut self, name: impl Into<Arc<str>>, value: ArrayRef) -> Result<()> {
80        let name = name.into();
81        if value.len() != 1 {
82            return Err(nd_err!(
83                InvalidAttribute,
84                "attribute '{name}' must have length 1, got {}",
85                value.len()
86            ));
87        }
88        self.insert_unchecked(name, value);
89        Ok(())
90    }
91
92    /// Builder form of [`Self::insert`].
93    pub fn with(mut self, name: impl Into<Arc<str>>, value: impl IntoScalar) -> Self {
94        self.insert(name, value);
95        self
96    }
97
98    /// Remove attribute `name` and return its value.
99    pub fn remove(&mut self, name: &str) -> Option<ArrayRef> {
100        let index = self
101            .entries
102            .iter()
103            .position(|(key, _)| key.as_ref() == name)?;
104        Some(self.entries.remove(index).1)
105    }
106
107    /// The attributes in insertion order.
108    pub fn iter(&self) -> impl Iterator<Item = (&str, &ArrayRef)> {
109        self.entries
110            .iter()
111            .map(|(key, value)| (key.as_ref(), value))
112    }
113
114    /// The `attrs` struct fields for these attributes. All fields are nullable.
115    pub fn fields(&self) -> Fields {
116        self.entries
117            .iter()
118            .map(|(key, value)| attribute_field(key, value.data_type()))
119            .collect()
120    }
121
122    /// Read row `row` of an `attrs` struct. A null child value is an absent attribute.
123    pub fn from_struct_row(attrs: &StructArray, row: usize) -> Self {
124        let entries = attrs
125            .fields()
126            .iter()
127            .zip(attrs.columns())
128            .filter(|(_, column)| column.is_valid(row))
129            .map(|(field, column)| (Arc::from(field.name().as_str()), column.slice(row, 1)))
130            .collect();
131        Self { entries }
132    }
133
134    /// The struct fields for the union of the keys in `all`. One key with two
135    /// types is an error.
136    pub fn union_fields<'a>(all: impl IntoIterator<Item = &'a Attributes>) -> Result<Fields> {
137        let mut fields: Vec<FieldRef> = Vec::new();
138        for attrs in all {
139            for (name, value) in attrs.iter() {
140                match fields.iter().find(|f| f.name() == name) {
141                    Some(existing) if existing.data_type() != value.data_type() => {
142                        return Err(nd_err!(
143                            InvalidAttribute,
144                            "attribute '{name}' has types {} and {}",
145                            existing.data_type(),
146                            value.data_type()
147                        ));
148                    }
149                    Some(_) => {}
150                    None => fields.push(attribute_field(name, value.data_type())),
151                }
152            }
153        }
154        Ok(fields.into())
155    }
156
157    /// Encode one struct row per entry of `rows`, with `fields` from
158    /// [`Self::union_fields`]. A missing key or a `None` row gives nulls.
159    pub fn encode_rows(rows: &[Option<&Attributes>], fields: &Fields) -> Result<StructArray> {
160        let columns = fields
161            .iter()
162            .map(|field| {
163                let null = new_null_array(field.data_type(), 1);
164                let parts: Vec<&dyn Array> = rows
165                    .iter()
166                    .map(|row| {
167                        row.and_then(|a| a.get(field.name()))
168                            .map_or(null.as_ref(), |v| v.as_ref())
169                    })
170                    .collect();
171                Ok(if parts.is_empty() {
172                    new_empty_array(field.data_type())
173                } else {
174                    concat(&parts)?
175                })
176            })
177            .collect::<Result<Vec<_>>>()?;
178        if fields.is_empty() {
179            return Ok(StructArray::new_empty_fields(rows.len(), None));
180        }
181        Ok(StructArray::try_new(fields.clone(), columns, None)?)
182    }
183
184    fn insert_unchecked(&mut self, name: Arc<str>, value: ArrayRef) {
185        match self.entries.iter_mut().find(|(key, _)| *key == name) {
186            Some(entry) => entry.1 = value,
187            None => self.entries.push((name, value)),
188        }
189    }
190}
191
192impl PartialEq for Attributes {
193    /// Equal when both hold the same keys with equal values, in any order.
194    fn eq(&self, other: &Self) -> bool {
195        self.len() == other.len()
196            && self
197                .iter()
198                .all(|(key, value)| other.get(key).is_some_and(|v| v == value))
199    }
200}
201
202pub(crate) fn attribute_field(name: &str, data_type: &DataType) -> FieldRef {
203    Arc::new(Field::new(name, data_type.clone(), true))
204}
205
206/// Conversion of a Rust value into a length-1 Arrow array (a scalar).
207pub trait IntoScalar {
208    /// The value as an Arrow array of length 1.
209    fn into_scalar(self) -> ArrayRef;
210}
211
212impl IntoScalar for &str {
213    fn into_scalar(self) -> ArrayRef {
214        Arc::new(StringArray::from(vec![self]))
215    }
216}
217
218impl IntoScalar for String {
219    fn into_scalar(self) -> ArrayRef {
220        self.as_str().into_scalar()
221    }
222}
223
224impl IntoScalar for bool {
225    fn into_scalar(self) -> ArrayRef {
226        Arc::new(BooleanArray::from(vec![self]))
227    }
228}
229
230impl IntoScalar for Vec<&str> {
231    fn into_scalar(self) -> ArrayRef {
232        list_of(Arc::new(StringArray::from(self)))
233    }
234}
235
236impl IntoScalar for Vec<String> {
237    fn into_scalar(self) -> ArrayRef {
238        list_of(Arc::new(StringArray::from(self)))
239    }
240}
241
242macro_rules! primitive_attribute {
243    ($($native:ty => $arrow:ty),* $(,)?) => {$(
244        impl IntoScalar for $native {
245            fn into_scalar(self) -> ArrayRef {
246                Arc::new(PrimitiveArray::<$arrow>::from(vec![self]))
247            }
248        }
249
250        impl IntoScalar for Vec<$native> {
251            fn into_scalar(self) -> ArrayRef {
252                list_of(Arc::new(PrimitiveArray::<$arrow>::from(self)))
253            }
254        }
255    )*};
256}
257
258primitive_attribute!(
259    i8 => Int8Type,
260    i16 => Int16Type,
261    i32 => Int32Type,
262    i64 => Int64Type,
263    u8 => UInt8Type,
264    u16 => UInt16Type,
265    u32 => UInt32Type,
266    u64 => UInt64Type,
267    f32 => Float32Type,
268    f64 => Float64Type,
269);
270
271fn list_of(values: ArrayRef) -> ArrayRef {
272    let field = Arc::new(Field::new("item", values.data_type().clone(), true));
273    Arc::new(ListArray::new(
274        field,
275        OffsetBuffer::from_lengths([values.len()]),
276        values,
277        None,
278    ))
279}
280
281#[cfg(test)]
282mod tests {
283    use arrow::array::AsArray;
284
285    use super::*;
286
287    #[test]
288    fn insert_replaces_and_keeps_order() {
289        let mut attrs = Attributes::new()
290            .with("units", "degC")
291            .with("scale_factor", 0.5f64)
292            .with("valid_range", vec![-5.0f32, 40.0]);
293        attrs.insert("units", "K");
294
295        assert_eq!(attrs.len(), 3);
296        assert_eq!(attrs.get_str("units"), Some("K"));
297        assert_eq!(
298            attrs.get_primitive::<Float64Type>("scale_factor"),
299            Some(0.5)
300        );
301        let names: Vec<&str> = attrs.iter().map(|(k, _)| k).collect();
302        assert_eq!(names, vec!["units", "scale_factor", "valid_range"]);
303
304        let range = attrs.get("valid_range").unwrap().as_list::<i32>().value(0);
305        assert_eq!(range.as_primitive::<Float32Type>().values(), &[-5.0, 40.0]);
306    }
307
308    #[test]
309    fn array_attribute_must_have_length_one() {
310        let mut attrs = Attributes::new();
311        let two: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
312        assert!(attrs.try_insert_array("bad", two).is_err());
313    }
314
315    #[test]
316    fn equality_ignores_order() {
317        let a = Attributes::new().with("a", 1i32).with("b", "x");
318        let b = Attributes::new().with("b", "x").with("a", 1i32);
319        assert_eq!(a, b);
320        assert_ne!(a, Attributes::new().with("a", 2i32).with("b", "x"));
321    }
322}
nd_arrow_array/attributes.rs

nd_arrow_array/
attributes.rs
