nd_arrow_array/
column.rs

1//! An Arrow column of nd arrays: the storage array of [`NdArrayType`].
2
3use std::sync::Arc;
4
5use arrow::array::{
6    Array, ArrayRef, ListArray, StringArray, StringBuilder, StructArray, UInt32Array,
7    UInt32Builder, new_empty_array, new_null_array,
8};
9use arrow::buffer::{NullBuffer, OffsetBuffer};
10use arrow::compute::concat;
11use arrow::datatypes::{DataType, Field, FieldRef};
12
13use crate::array::NdArrowArray;
14use crate::attributes::Attributes;
15use crate::dimensions::{Dimension, Dimensions};
16use crate::error::{NdError, Result, nd_err};
17use crate::extension::NdArrayType;
18
19/// A validated nd column. Each row holds one [`NdArrowArray`] or a null.
20///
21/// Construction checks every row once, so [`Self::value`] is infallible and
22/// slices the storage without a copy.
23#[derive(Debug, Clone)]
24pub struct NdColumn {
25    storage: StructArray,
26    extension: NdArrayType,
27}
28
29impl NdColumn {
30    /// Wrap and validate a storage array.
31    pub fn try_new(storage: StructArray) -> Result<Self> {
32        let extension = NdArrayType::try_from_storage_type(storage.data_type())?;
33        validate_rows(&storage)?;
34        // An `attrs` struct without fields is the same as no `attrs` field.
35        let storage = if !extension.has_attributes() && storage.num_columns() == 4 {
36            let (fields, mut children, nulls) = storage.into_parts();
37            children.truncate(3);
38            StructArray::try_new(fields.iter().take(3).cloned().collect(), children, nulls)?
39        } else {
40            storage
41        };
42        Ok(Self { storage, extension })
43    }
44
45    /// Wrap and validate a storage array given as an [`ArrayRef`].
46    pub fn try_from_array(array: &dyn Array) -> Result<Self> {
47        let storage = array
48            .as_any()
49            .downcast_ref::<StructArray>()
50            .ok_or_else(|| {
51                nd_err!(
52                    InvalidEncoding,
53                    "nd column must be a Struct array, got {}",
54                    array.data_type()
55                )
56            })?;
57        Self::try_new(storage.clone())
58    }
59
60    /// Encode `arrays` as one row each. `arrays` must not be empty.
61    pub fn try_from_arrays<'a>(arrays: impl IntoIterator<Item = &'a NdArrowArray>) -> Result<Self> {
62        let mut builder = NdColumnBuilder::new();
63        for array in arrays {
64            builder.append(array)?;
65        }
66        builder.finish()
67    }
68
69    /// A column of `len` null rows.
70    pub fn new_null(extension: NdArrayType, len: usize) -> Self {
71        let storage = new_null_array(&extension.storage_type(), len)
72            .as_any()
73            .downcast_ref::<StructArray>()
74            .expect("storage type is a struct")
75            .clone();
76        Self { storage, extension }
77    }
78
79    /// Wrap storage that matches `extension` and holds valid rows.
80    pub(crate) fn from_parts_unchecked(storage: StructArray, extension: NdArrayType) -> Self {
81        debug_assert_eq!(storage.data_type(), &extension.storage_type());
82        Self { storage, extension }
83    }
84
85    /// The extension type of the column.
86    pub fn extension_type(&self) -> &NdArrayType {
87        &self.extension
88    }
89
90    /// The Arrow type of the array values.
91    pub fn value_type(&self) -> &DataType {
92        self.extension.value_type()
93    }
94
95    /// A field named `name` with the storage type and the extension tag.
96    pub fn to_field(&self, name: impl Into<String>) -> Field {
97        self.extension.to_field(name, true)
98    }
99
100    /// The number of rows.
101    pub fn len(&self) -> usize {
102        self.storage.len()
103    }
104
105    /// True when the column has no rows.
106    pub fn is_empty(&self) -> bool {
107        self.storage.is_empty()
108    }
109
110    /// True when `row` is null.
111    pub fn is_null(&self, row: usize) -> bool {
112        self.storage.is_null(row)
113    }
114
115    /// The number of null rows.
116    pub fn null_count(&self) -> usize {
117        self.storage.null_count()
118    }
119
120    /// The nd array at `row`, or `None` for a null row.
121    pub fn get(&self, row: usize) -> Option<NdArrowArray> {
122        if self.storage.is_null(row) {
123            return None;
124        }
125        let values = self.values_list().value(row);
126        let dims = self.dimensions_at(row);
127        let attrs = self
128            .attrs_struct()
129            .map(|attrs| Attributes::from_struct_row(attrs, row))
130            .unwrap_or_default();
131        Some(NdArrowArray::new_unchecked(values, dims, attrs))
132    }
133
134    /// The nd array at `row`. A null row gives a rank-0 null scalar,
135    /// which broadcasts to an all-null array.
136    pub fn value(&self, row: usize) -> NdArrowArray {
137        self.get(row)
138            .unwrap_or_else(|| NdArrowArray::new_null_scalar(self.value_type()))
139    }
140
141    /// The dimensions at `row`. A null row gives scalar dimensions.
142    pub fn dimensions_at(&self, row: usize) -> Dimensions {
143        if self.storage.is_null(row) {
144            return Dimensions::scalar();
145        }
146        let (names, sizes) = self.dims_at(row);
147        let dims: Vec<Dimension> = names
148            .iter()
149            .zip(sizes.values())
150            .map(|(name, &size)| Dimension::new(name.expect("validated"), size as usize))
151            .collect();
152        Dimensions::try_new(dims).expect("validated")
153    }
154
155    /// The array of each row, `None` for null rows.
156    pub fn iter(&self) -> impl ExactSizeIterator<Item = Option<NdArrowArray>> + '_ {
157        (0..self.len()).map(|row| self.get(row))
158    }
159
160    /// A zero-copy slice of `len` rows from `offset`.
161    pub fn slice(&self, offset: usize, len: usize) -> Self {
162        Self {
163            storage: self.storage.slice(offset, len),
164            extension: self.extension.clone(),
165        }
166    }
167
168    /// The flat values of all rows (the child of the `values` list).
169    pub fn values_list(&self) -> &ListArray {
170        list_column(&self.storage, 0)
171    }
172
173    /// The `attrs` struct, or `None` when the column has no attributes.
174    pub fn attrs_struct(&self) -> Option<&StructArray> {
175        self.extension.has_attributes().then(|| {
176            self.storage
177                .column(3)
178                .as_any()
179                .downcast_ref()
180                .expect("validated")
181        })
182    }
183
184    /// The storage `Struct` array.
185    pub fn storage(&self) -> &StructArray {
186        &self.storage
187    }
188
189    /// The storage `Struct` array, by value.
190    pub fn into_storage(self) -> StructArray {
191        self.storage
192    }
193
194    /// The storage as an [`ArrayRef`].
195    pub fn into_array_ref(self) -> ArrayRef {
196        Arc::new(self.storage)
197    }
198
199    fn dims_at(&self, row: usize) -> (StringArray, UInt32Array) {
200        let names = list_column(&self.storage, 1).value(row);
201        let sizes = list_column(&self.storage, 2).value(row);
202        (
203            names
204                .as_any()
205                .downcast_ref::<StringArray>()
206                .unwrap()
207                .clone(),
208            sizes
209                .as_any()
210                .downcast_ref::<UInt32Array>()
211                .unwrap()
212                .clone(),
213        )
214    }
215}
216
217impl From<NdColumn> for ArrayRef {
218    fn from(column: NdColumn) -> Self {
219        column.into_array_ref()
220    }
221}
222
223impl TryFrom<StructArray> for NdColumn {
224    type Error = NdError;
225
226    fn try_from(storage: StructArray) -> Result<Self> {
227        Self::try_new(storage)
228    }
229}
230
231fn list_column(storage: &StructArray, index: usize) -> &ListArray {
232    storage
233        .column(index)
234        .as_any()
235        .downcast_ref()
236        .expect("validated")
237}
238
239/// Check each non-null row: list validity, dimension names and sizes, value count.
240fn validate_rows(storage: &StructArray) -> Result<()> {
241    let values = list_column(storage, 0);
242    let names = list_column(storage, 1);
243    let sizes = list_column(storage, 2);
244    let name_values = names
245        .values()
246        .as_any()
247        .downcast_ref::<StringArray>()
248        .ok_or_else(|| nd_err!(InvalidEncoding, "'dim_names' items must be Utf8"))?;
249    let size_values = sizes
250        .values()
251        .as_any()
252        .downcast_ref::<UInt32Array>()
253        .ok_or_else(|| nd_err!(InvalidEncoding, "'dim_sizes' items must be UInt32"))?;
254
255    for row in 0..storage.len() {
256        if storage.is_null(row) {
257            continue;
258        }
259        if values.is_null(row) || names.is_null(row) || sizes.is_null(row) {
260            return Err(nd_err!(
261                InvalidEncoding,
262                "row {row}: a non-null nd row has a null list"
263            ));
264        }
265        let name_range = list_range(names, row);
266        let size_range = list_range(sizes, row);
267        if name_range.len() != size_range.len() {
268            return Err(nd_err!(
269                InvalidEncoding,
270                "row {row}: {} dimension names but {} dimension sizes",
271                name_range.len(),
272                size_range.len()
273            ));
274        }
275        let mut dims = Vec::with_capacity(name_range.len());
276        for (n, s) in name_range.zip(size_range) {
277            if name_values.is_null(n) || size_values.is_null(s) {
278                return Err(nd_err!(
279                    InvalidEncoding,
280                    "row {row}: null dimension name or size"
281                ));
282            }
283            dims.push(Dimension::new(
284                name_values.value(n),
285                size_values.value(s) as usize,
286            ));
287        }
288        let dims = Dimensions::try_new(dims)?;
289        let expected = dims.checked_num_elements();
290        let actual = values.value_length(row) as usize;
291        if expected != Some(actual) {
292            return Err(nd_err!(
293                ShapeMismatch,
294                "row {row}: {actual} values do not match dimensions {dims}"
295            ));
296        }
297    }
298    Ok(())
299}
300
301fn list_range(list: &ListArray, row: usize) -> std::ops::Range<usize> {
302    let offsets = list.value_offsets();
303    offsets[row] as usize..offsets[row + 1] as usize
304}
305
306/// Collects nd arrays and encodes them into an [`NdColumn`].
307///
308/// The `attrs` struct of the column holds the union of all attribute keys.
309/// A row without a key gets a null for it. The same key with two different
310/// types is an error.
311#[derive(Debug, Default)]
312pub struct NdColumnBuilder {
313    value_field: Option<FieldRef>,
314    rows: Vec<Option<NdArrowArray>>,
315}
316
317impl NdColumnBuilder {
318    /// A builder that takes the value type from the first array.
319    pub fn new() -> Self {
320        Self::default()
321    }
322
323    /// A builder with a known value type, so that all-null and empty columns work.
324    pub fn with_value_type(value_type: DataType) -> Self {
325        Self::with_value_field(Arc::new(Field::new(
326            crate::extension::ITEM_FIELD,
327            value_type,
328            true,
329        )))
330    }
331
332    /// A builder with a custom item field for the values list.
333    pub fn with_value_field(value_field: FieldRef) -> Self {
334        Self {
335            value_field: Some(value_field),
336            rows: Vec::new(),
337        }
338    }
339
340    /// The number of rows appended.
341    pub fn len(&self) -> usize {
342        self.rows.len()
343    }
344
345    /// True when no rows were appended.
346    pub fn is_empty(&self) -> bool {
347        self.rows.is_empty()
348    }
349
350    /// Append one array as a row. Its value type must match the column.
351    pub fn append(&mut self, array: &NdArrowArray) -> Result<()> {
352        match &self.value_field {
353            Some(field) if field.data_type() != array.data_type() => {
354                return Err(nd_err!(
355                    SchemaMismatch,
356                    "nd array has value type {} but the column has {}",
357                    array.data_type(),
358                    field.data_type()
359                ));
360            }
361            Some(_) => {}
362            None => {
363                self.value_field = Some(Arc::new(Field::new(
364                    crate::extension::ITEM_FIELD,
365                    array.data_type().clone(),
366                    true,
367                )));
368            }
369        }
370        self.rows.push(Some(array.clone()));
371        Ok(())
372    }
373
374    /// Append a null row.
375    pub fn append_null(&mut self) {
376        self.rows.push(None);
377    }
378
379    /// Append an array, or a null row for `None`.
380    pub fn append_option(&mut self, array: Option<&NdArrowArray>) -> Result<()> {
381        match array {
382            Some(array) => self.append(array),
383            None => {
384                self.append_null();
385                Ok(())
386            }
387        }
388    }
389
390    /// Encode the rows. Fails when no value type is known.
391    pub fn finish(self) -> Result<NdColumn> {
392        let value_field = self.value_field.ok_or_else(|| {
393            nd_err!(
394                SchemaMismatch,
395                "cannot infer the value type of an nd column without non-null rows"
396            )
397        })?;
398        let rows = self.rows;
399        let present: Vec<&NdArrowArray> = rows.iter().flatten().collect();
400
401        let attr_fields = Attributes::union_fields(present.iter().map(|a| a.attributes()))?;
402        let extension = NdArrayType::with_value_field(value_field.clone(), attr_fields.clone());
403
404        let values = concat_or_empty(
405            &present
406                .iter()
407                .map(|a| a.values().as_ref())
408                .collect::<Vec<_>>(),
409            value_field.data_type(),
410        )?;
411        let values_list = ListArray::try_new(
412            value_field,
413            offsets_from_lengths(rows.iter().map(|r| r.as_ref().map_or(0, NdArrowArray::len)))?,
414            values,
415            None,
416        )?;
417
418        let rank_offsets = offsets_from_lengths(
419            rows.iter()
420                .map(|r| r.as_ref().map_or(0, NdArrowArray::rank)),
421        )?;
422        let total_rank = *rank_offsets.last().unwrap() as usize;
423        let mut names = StringBuilder::with_capacity(total_rank, total_rank * 8);
424        let mut sizes = UInt32Builder::with_capacity(total_rank);
425        for dim in present.iter().flat_map(|a| a.dims().iter()) {
426            let size = u32::try_from(dim.size()).map_err(|_| {
427                nd_err!(
428                    InvalidDimensions,
429                    "dimension '{}' of size {} does not fit in UInt32",
430                    dim.name(),
431                    dim.size()
432                )
433            })?;
434            names.append_value(dim.name());
435            sizes.append_value(size);
436        }
437        let storage_fields = extension.storage_fields();
438        let item_field = |index: usize| match storage_fields[index].data_type() {
439            DataType::List(item) => item.clone(),
440            _ => unreachable!(),
441        };
442        let names_list = ListArray::try_new(
443            item_field(1),
444            rank_offsets.clone(),
445            Arc::new(names.finish()),
446            None,
447        )?;
448        let sizes_list =
449            ListArray::try_new(item_field(2), rank_offsets, Arc::new(sizes.finish()), None)?;
450
451        let mut children: Vec<ArrayRef> = vec![
452            Arc::new(values_list),
453            Arc::new(names_list),
454            Arc::new(sizes_list),
455        ];
456        if extension.has_attributes() {
457            let attrs: Vec<Option<&Attributes>> = rows
458                .iter()
459                .map(|r| r.as_ref().map(NdArrowArray::attributes))
460                .collect();
461            children.push(Arc::new(Attributes::encode_rows(&attrs, &attr_fields)?));
462        }
463
464        let nulls = rows
465            .iter()
466            .any(Option::is_none)
467            .then(|| NullBuffer::from_iter(rows.iter().map(Option::is_some)));
468
469        let storage = StructArray::try_new(storage_fields, children, nulls)?;
470        Ok(NdColumn { storage, extension })
471    }
472}
473
474fn offsets_from_lengths(lengths: impl Iterator<Item = usize>) -> Result<OffsetBuffer<i32>> {
475    let lengths: Vec<usize> = lengths.collect();
476    let total: usize = lengths.iter().sum();
477    if i32::try_from(total).is_err() {
478        return Err(nd_err!(
479            InvalidEncoding,
480            "{total} list items exceed the i32 offset range"
481        ));
482    }
483    Ok(OffsetBuffer::from_lengths(lengths))
484}
485
486fn concat_or_empty(arrays: &[&dyn Array], data_type: &DataType) -> Result<ArrayRef> {
487    Ok(match arrays {
488        [] => new_empty_array(data_type),
489        [single] => single.slice(0, single.len()),
490        many => concat(many)?,
491    })
492}
493
494#[cfg(test)]
495mod tests {
496    use arrow::array::{AsArray, Float64Array, Int32Array};
497    use arrow::datatypes::Fields;
498    use arrow::datatypes::{Float64Type, Int32Type};
499
500    use super::*;
501
502    fn nd(values: Vec<i32>, dims: &[(&str, usize)]) -> NdArrowArray {
503        NdArrowArray::try_new(
504            Arc::new(Int32Array::from(values)),
505            Dimensions::try_new(dims.iter().copied()).unwrap(),
506        )
507        .unwrap()
508    }
509
510    #[test]
511    fn rows_round_trip_with_attribute_union() {
512        let a = nd(vec![1, 2, 3, 4, 5, 6], &[("time", 2), ("lat", 3)])
513            .with_attributes(Attributes::new().with("units", "K"));
514        let b = nd(vec![7, 8], &[("lat", 2)])
515            .with_attributes(Attributes::new().with("scale_factor", 2.0f64));
516        let c = nd(vec![9], &[]);
517
518        let column = NdColumn::try_from_arrays([&a, &b, &c]).unwrap();
519        assert_eq!(column.len(), 3);
520        assert_eq!(
521            column.extension_type().attribute_fields().len(),
522            2,
523            "units and scale_factor"
524        );
525
526        let decoded: Vec<NdArrowArray> = column.iter().map(Option::unwrap).collect();
527        assert_eq!(decoded, vec![a, b, c]);
528        assert_eq!(
529            decoded[1]
530                .attributes()
531                .get_primitive::<Float64Type>("scale_factor"),
532            Some(2.0)
533        );
534        assert!(!decoded[1].attributes().contains("units"));
535    }
536
537    #[test]
538    fn null_rows_decode_to_null_scalars() {
539        let mut builder = NdColumnBuilder::new();
540        builder.append_null();
541        builder.append(&nd(vec![1, 2], &[("x", 2)])).unwrap();
542        let column = builder.finish().unwrap();
543
544        assert_eq!(column.null_count(), 1);
545        assert!(column.get(0).is_none());
546        let null = column.value(0);
547        assert!(null.is_scalar());
548        assert_eq!(null.values().null_count(), 1);
549        assert_eq!(
550            column
551                .value(1)
552                .values()
553                .as_primitive::<Int32Type>()
554                .values(),
555            &[1, 2]
556        );
557    }
558
559    #[test]
560    fn empty_column_needs_a_value_type() {
561        assert!(NdColumnBuilder::new().finish().is_err());
562        let column = NdColumnBuilder::with_value_type(DataType::Float64)
563            .finish()
564            .unwrap();
565        assert!(column.is_empty());
566        assert_eq!(column.value_type(), &DataType::Float64);
567    }
568
569    #[test]
570    fn value_type_mismatch_rejected() {
571        let mut builder = NdColumnBuilder::new();
572        builder.append(&nd(vec![1], &[])).unwrap();
573        let float = NdArrowArray::try_new_scalar(Arc::new(Float64Array::from(vec![1.0]))).unwrap();
574        assert!(builder.append(&float).is_err());
575    }
576
577    #[test]
578    fn conflicting_attribute_types_rejected() {
579        let a = nd(vec![1], &[]).with_attributes(Attributes::new().with("units", "K"));
580        let b = nd(vec![1], &[]).with_attributes(Attributes::new().with("units", 1i32));
581        assert!(NdColumn::try_from_arrays([&a, &b]).is_err());
582    }
583
584    #[test]
585    fn attrs_field_is_left_out_without_attributes() {
586        let column = NdColumn::try_from_arrays([&nd(vec![1, 2], &[("x", 2)])]).unwrap();
587        assert_eq!(column.storage().num_columns(), 3);
588        assert!(column.attrs_struct().is_none());
589
590        let (fields, mut children, nulls) = column.clone().into_storage().into_parts();
591        let mut fields: Vec<FieldRef> = fields.iter().cloned().collect();
592        fields.push(Arc::new(Field::new(
593            crate::extension::ATTRS_FIELD,
594            DataType::Struct(Fields::empty()),
595            false,
596        )));
597        children.push(Arc::new(StructArray::new_empty_fields(1, None)));
598        let with_empty = StructArray::new(fields.into(), children, nulls);
599        let normalized = NdColumn::try_new(with_empty).unwrap();
600        assert_eq!(normalized.storage().num_columns(), 3);
601        assert_eq!(normalized.value(0), column.value(0));
602    }
603
604    #[test]
605    fn slice_keeps_rows() {
606        let arrays = [nd(vec![1], &[]), nd(vec![2, 3], &[("x", 2)])];
607        let column = NdColumn::try_from_arrays(&arrays).unwrap().slice(1, 1);
608        assert_eq!(column.len(), 1);
609        assert_eq!(column.value(0), arrays[1]);
610    }
611
612    #[test]
613    fn invalid_storage_rejected() {
614        let column = NdColumn::try_from_arrays([&nd(vec![1, 2], &[("x", 2)])]).unwrap();
615        let (fields, mut children, nulls) = column.into_storage().into_parts();
616        children[2] = Arc::new(ListArray::new(
617            Arc::new(Field::new("item", DataType::UInt32, false)),
618            OffsetBuffer::from_lengths([1]),
619            Arc::new(UInt32Array::from(vec![3])),
620            None,
621        ));
622        let broken = StructArray::new(fields, children, nulls);
623        let err = NdColumn::try_new(broken).unwrap_err();
624        assert!(err.to_string().contains("do not match"), "{err}");
625    }
626}
nd_arrow_array/column.rs

nd_arrow_array/
column.rs
