nd_arrow_array/
batch.rs

1//! Grid-shaped record batches: columns on subsets of one shared target grid.
2
3use std::sync::Arc;
4
5use arrow::array::{ArrayRef, StructArray};
6use arrow::datatypes::{Field, Schema, SchemaRef};
7use arrow::record_batch::{RecordBatch, RecordBatchOptions};
8
9use crate::array::NdArrowArray;
10use crate::attributes::Attributes;
11use crate::column::NdColumn;
12use crate::dimensions::{Dimension, Dimensions};
13use crate::error::{Result, nd_err};
14use crate::extension::{
15    GLOBAL_ATTRIBUTES_FIELD, global_attributes_field, is_global_attributes_field, is_nd_field,
16    logical_schema, with_base_metadata,
17};
18
19/// A record batch whose columns are [`NdArrowArray`]s over a shared target grid.
20///
21/// Each column can use a subset of the target dimensions: a scalar attribute,
22/// a 1-D coordinate, a full-rank data variable. Columns stay un-broadcast.
23/// The schema is the logical schema: one field per column, typed by its values.
24/// Global attributes describe the whole batch, like xarray `Dataset.attrs`.
25#[derive(Debug, Clone)]
26pub struct NdRecordBatch {
27    schema: SchemaRef,
28    columns: Vec<NdArrowArray>,
29    target: Dimensions,
30    attributes: Attributes,
31}
32
33impl NdRecordBatch {
34    /// Create a batch. The schema must match the column types, and every column
35    /// must broadcast onto `target`.
36    pub fn try_new(
37        schema: SchemaRef,
38        columns: Vec<NdArrowArray>,
39        target: Dimensions,
40    ) -> Result<Self> {
41        if schema.fields().len() != columns.len() {
42            return Err(nd_err!(
43                SchemaMismatch,
44                "nd batch has {} columns but the schema declares {} fields",
45                columns.len(),
46                schema.fields().len()
47            ));
48        }
49        for (field, column) in schema.fields().iter().zip(&columns) {
50            if field.data_type() != column.data_type() {
51                return Err(nd_err!(
52                    SchemaMismatch,
53                    "column '{}' has type {} but the schema declares {}",
54                    field.name(),
55                    column.data_type(),
56                    field.data_type()
57                ));
58            }
59            column.dims().check_broadcast_to(&target)?;
60        }
61        Ok(Self {
62            schema,
63            columns,
64            target,
65            attributes: Attributes::new(),
66        })
67    }
68
69    /// Replace the global attributes.
70    pub fn with_attributes(mut self, attributes: Attributes) -> Self {
71        self.attributes = attributes;
72        self
73    }
74
75    /// The global attributes of the batch.
76    pub fn attributes(&self) -> &Attributes {
77        &self.attributes
78    }
79
80    /// The global attributes, for changes in place.
81    pub fn attributes_mut(&mut self) -> &mut Attributes {
82        &mut self.attributes
83    }
84
85    /// A batch with the target grid inferred by [`infer_target`].
86    pub fn try_new_inferred(schema: SchemaRef, columns: Vec<NdArrowArray>) -> Result<Self> {
87        let target = infer_target(&columns)?;
88        Self::try_new(schema, columns, target)
89    }
90
91    /// The logical schema: one field per column, typed by its values.
92    pub fn schema(&self) -> &SchemaRef {
93        &self.schema
94    }
95
96    /// All columns, in schema order.
97    pub fn columns(&self) -> &[NdArrowArray] {
98        &self.columns
99    }
100
101    /// The column at `index`. Panics when `index` is out of bounds.
102    pub fn column(&self, index: usize) -> &NdArrowArray {
103        &self.columns[index]
104    }
105
106    /// The column named `name`, if present.
107    pub fn column_by_name(&self, name: &str) -> Option<&NdArrowArray> {
108        let index = self.schema.index_of(name).ok()?;
109        Some(&self.columns[index])
110    }
111
112    /// The number of columns.
113    pub fn num_columns(&self) -> usize {
114        self.columns.len()
115    }
116
117    /// The shared target grid.
118    pub fn target(&self) -> &Dimensions {
119        &self.target
120    }
121
122    /// Rows after a broadcast onto the full target grid.
123    pub fn num_rows(&self) -> usize {
124        self.target.num_elements()
125    }
126
127    /// Encode as a single-row [`RecordBatch`] of nd columns.
128    pub fn to_record_batch(&self) -> Result<RecordBatch> {
129        encode_batches(std::slice::from_ref(self))
130    }
131
132    /// Decode every row of an nd-encoded [`RecordBatch`]: one nd batch per row.
133    pub fn try_from_record_batch(batch: &RecordBatch) -> Result<Vec<Self>> {
134        let decoder = Decoder::try_new(batch)?;
135        (0..batch.num_rows()).map(|row| decoder.row(row)).collect()
136    }
137
138    /// Decode row `row` of an nd-encoded [`RecordBatch`].
139    pub fn try_from_record_batch_row(batch: &RecordBatch, row: usize) -> Result<Self> {
140        if row >= batch.num_rows() {
141            return Err(nd_err!(
142                InvalidEncoding,
143                "row {row} is out of bounds for a batch of {} rows",
144                batch.num_rows()
145            ));
146        }
147        Decoder::try_new(batch)?.row(row)
148    }
149}
150
151/// Encode several nd batches with one shared logical schema, one row each.
152///
153/// The `attrs` struct of each column holds the union of the attribute keys of all batches.
154/// When a batch has global attributes, a last column [`GLOBAL_ATTRIBUTES_FIELD`]
155/// holds them, also as the union of keys.
156pub fn encode_batches(batches: &[NdRecordBatch]) -> Result<RecordBatch> {
157    let Some(first) = batches.first() else {
158        return Err(nd_err!(SchemaMismatch, "no nd batches to encode"));
159    };
160    let schema = first.schema();
161    if let Some(other) = batches
162        .iter()
163        .find(|b| b.schema().fields() != schema.fields())
164    {
165        return Err(nd_err!(
166            SchemaMismatch,
167            "nd batches disagree on the schema: {} vs {}",
168            schema,
169            other.schema()
170        ));
171    }
172
173    let mut fields: Vec<Field> = Vec::with_capacity(schema.fields().len());
174    let mut columns: Vec<ArrayRef> = Vec::with_capacity(schema.fields().len());
175    for (index, logical) in schema.fields().iter().enumerate() {
176        let column = NdColumn::try_from_arrays(batches.iter().map(|b| b.column(index)))?;
177        fields.push(with_base_metadata(
178            column.to_field(logical.name()),
179            logical.metadata(),
180        ));
181        columns.push(column.into_array_ref());
182    }
183
184    if batches.iter().any(|b| !b.attributes().is_empty()) {
185        if schema.field_with_name(GLOBAL_ATTRIBUTES_FIELD).is_ok() {
186            return Err(nd_err!(
187                SchemaMismatch,
188                "column name '{GLOBAL_ATTRIBUTES_FIELD}' is reserved for global attributes"
189            ));
190        }
191        let global = Attributes::union_fields(batches.iter().map(NdRecordBatch::attributes))?;
192        let rows: Vec<Option<&Attributes>> = batches.iter().map(|b| Some(b.attributes())).collect();
193        columns.push(Arc::new(Attributes::encode_rows(&rows, &global)?));
194        fields.push(global_attributes_field(global));
195    }
196
197    let schema = Schema::new_with_metadata(fields, schema.metadata().clone());
198    let options = RecordBatchOptions::new().with_row_count(Some(batches.len()));
199    Ok(RecordBatch::try_new_with_options(
200        Arc::new(schema),
201        columns,
202        &options,
203    )?)
204}
205
206struct Decoder {
207    schema: SchemaRef,
208    columns: Vec<NdColumn>,
209    global: Option<StructArray>,
210}
211
212impl Decoder {
213    fn try_new(batch: &RecordBatch) -> Result<Self> {
214        let mut nd_fields = Vec::new();
215        let mut columns = Vec::new();
216        let mut global = None;
217        for (field, column) in batch.schema().fields().iter().zip(batch.columns()) {
218            if is_global_attributes_field(field) {
219                if global.is_some() {
220                    return Err(nd_err!(
221                        InvalidEncoding,
222                        "more than one global attributes column"
223                    ));
224                }
225                global = column.as_any().downcast_ref::<StructArray>().cloned();
226            } else if is_nd_field(field) {
227                nd_fields.push(field.clone());
228                columns.push(NdColumn::try_from_array(column.as_ref())?);
229            } else {
230                return Err(nd_err!(
231                    InvalidEncoding,
232                    "column '{}' is not an nd column",
233                    field.name()
234                ));
235            }
236        }
237        let encoded = Schema::new_with_metadata(nd_fields, batch.schema().metadata().clone());
238        let schema = Arc::new(logical_schema(&encoded)?);
239        Ok(Self {
240            schema,
241            columns,
242            global,
243        })
244    }
245
246    fn row(&self, row: usize) -> Result<NdRecordBatch> {
247        let columns: Vec<NdArrowArray> = self.columns.iter().map(|c| c.value(row)).collect();
248        let attributes = self
249            .global
250            .as_ref()
251            .map(|g| Attributes::from_struct_row(g, row))
252            .unwrap_or_default();
253        Ok(
254            NdRecordBatch::try_new_inferred(self.schema.clone(), columns)?
255                .with_attributes(attributes),
256        )
257    }
258}
259
260/// Infer the target grid from columns. The highest-rank column sets the axis
261/// order. Axes that only lower-rank columns carry are appended. A size-1 axis
262/// takes a larger size from another column.
263pub fn infer_target(columns: &[NdArrowArray]) -> Result<Dimensions> {
264    let mut order: Vec<Dimension> = Vec::new();
265    if let Some(widest) = columns.iter().max_by_key(|c| c.rank()) {
266        order.extend(widest.dims().iter().cloned());
267    }
268    for dim in columns.iter().flat_map(|c| c.dims().iter()) {
269        match order.iter_mut().find(|d| d.name() == dim.name()) {
270            Some(existing) if existing.size() == 1 => *existing = dim.clone(),
271            Some(_) => {}
272            None => order.push(dim.clone()),
273        }
274    }
275    Dimensions::try_new(order)
276}
277
278#[cfg(test)]
279mod tests {
280    use crate::extension::is_global_attributes_field;
281    use arrow::array::{Float64Array, Int32Array};
282    use arrow::datatypes::DataType;
283
284    use super::*;
285
286    fn dims(spec: &[(&str, usize)]) -> Dimensions {
287        Dimensions::try_new(spec.iter().copied()).unwrap()
288    }
289
290    fn grid_batch() -> NdRecordBatch {
291        let schema = Arc::new(Schema::new(vec![
292            Field::new("time", DataType::Int32, true),
293            Field::new("lat", DataType::Int32, true),
294            Field::new("sst", DataType::Float64, true),
295        ]));
296        let time =
297            NdArrowArray::try_new(Arc::new(Int32Array::from(vec![7, 8])), dims(&[("time", 2)]))
298                .unwrap();
299        let lat = NdArrowArray::try_new(
300            Arc::new(Int32Array::from(vec![10, 20, 30])),
301            dims(&[("lat", 3)]),
302        )
303        .unwrap()
304        .with_attributes(Attributes::new().with("units", "degrees_north"));
305        let sst = NdArrowArray::try_new(
306            Arc::new(Float64Array::from(vec![0.0, 0.1, 0.2, 1.0, 1.1, 1.2])),
307            dims(&[("time", 2), ("lat", 3)]),
308        )
309        .unwrap()
310        .with_attributes(
311            Attributes::new()
312                .with("units", "K")
313                .with("_FillValue", -999.0),
314        );
315        NdRecordBatch::try_new(
316            schema,
317            vec![time, lat, sst],
318            dims(&[("time", 2), ("lat", 3)]),
319        )
320        .unwrap()
321    }
322
323    #[test]
324    fn record_batch_round_trip() {
325        let nd = grid_batch();
326        let encoded = nd.to_record_batch().unwrap();
327        assert_eq!(encoded.num_rows(), 1);
328        assert!(encoded.schema().fields().iter().all(|f| is_nd_field(f)));
329
330        let decoded = NdRecordBatch::try_from_record_batch_row(&encoded, 0).unwrap();
331        assert_eq!(decoded.target(), nd.target());
332        assert_eq!(decoded.schema().fields(), nd.schema().fields());
333        assert_eq!(decoded.columns(), nd.columns());
334    }
335
336    #[test]
337    fn every_row_of_a_concatenated_batch_decodes() {
338        let other = NdRecordBatch::try_new_inferred(
339            grid_batch().schema().clone(),
340            vec![
341                NdArrowArray::try_new(Arc::new(Int32Array::from(vec![1])), dims(&[("time", 1)]))
342                    .unwrap(),
343                NdArrowArray::try_new(Arc::new(Int32Array::from(vec![5, 6])), dims(&[("lat", 2)]))
344                    .unwrap(),
345                NdArrowArray::try_new(
346                    Arc::new(Float64Array::from(vec![3.0, 4.0])),
347                    dims(&[("time", 1), ("lat", 2)]),
348                )
349                .unwrap(),
350            ],
351        )
352        .unwrap();
353
354        let encoded = encode_batches(&[grid_batch(), other.clone()]).unwrap();
355        let decoded = NdRecordBatch::try_from_record_batch(&encoded).unwrap();
356        let rows: Vec<usize> = decoded.iter().map(NdRecordBatch::num_rows).collect();
357        assert_eq!(rows, vec![6, 2]);
358        assert_eq!(decoded[1].columns(), other.columns());
359    }
360
361    #[test]
362    fn global_attributes_round_trip_per_row() {
363        let a = grid_batch().with_attributes(
364            Attributes::new()
365                .with("title", "first")
366                .with("version", 2i32),
367        );
368        let b = grid_batch().with_attributes(Attributes::new().with("title", "second"));
369        let encoded = encode_batches(&[a.clone(), b.clone()]).unwrap();
370        let last = encoded.schema().fields().last().unwrap().clone();
371        assert_eq!(last.name(), GLOBAL_ATTRIBUTES_FIELD);
372        assert!(is_global_attributes_field(&last));
373
374        let decoded = NdRecordBatch::try_from_record_batch(&encoded).unwrap();
375        assert_eq!(decoded[0].attributes(), a.attributes());
376        assert_eq!(decoded[1].attributes(), b.attributes());
377        assert_eq!(decoded[0].schema().fields(), a.schema().fields());
378
379        // Concatenated encoded batches keep each row's own attributes.
380        let merged =
381            arrow::compute::concat_batches(&encoded.schema(), &[encoded.clone(), encoded]).unwrap();
382        let again = NdRecordBatch::try_from_record_batch(&merged).unwrap();
383        assert_eq!(again[3].attributes().get_str("title"), Some("second"));
384    }
385
386    #[test]
387    fn no_global_column_without_attributes() {
388        let encoded = grid_batch().to_record_batch().unwrap();
389        assert!(encoded.schema().fields().iter().all(|f| is_nd_field(f)));
390    }
391
392    #[test]
393    fn reserved_column_name_rejected() {
394        let schema = Arc::new(Schema::new(vec![Field::new(
395            GLOBAL_ATTRIBUTES_FIELD,
396            DataType::Int32,
397            true,
398        )]));
399        let column =
400            NdArrowArray::try_new(Arc::new(Int32Array::from(vec![1])), dims(&[("x", 1)])).unwrap();
401        let b = NdRecordBatch::try_new_inferred(schema, vec![column])
402            .unwrap()
403            .with_attributes(Attributes::new().with("title", "x"));
404        assert!(b.to_record_batch().is_err());
405    }
406
407    #[test]
408    fn target_appends_axes_of_narrow_columns() {
409        let sst = NdArrowArray::try_new(
410            Arc::new(Float64Array::from(vec![0.0; 6])),
411            dims(&[("time", 2), ("lat", 3)]),
412        )
413        .unwrap();
414        let depth = NdArrowArray::try_new(
415            Arc::new(Int32Array::from(vec![5, 10])),
416            dims(&[("depth", 2)]),
417        )
418        .unwrap();
419        let target = infer_target(&[sst, depth]).unwrap();
420        assert_eq!(target, dims(&[("time", 2), ("lat", 3), ("depth", 2)]));
421    }
422
423    #[test]
424    fn schema_mismatch_rejected() {
425        let schema = Arc::new(Schema::new(vec![Field::new("t", DataType::Float64, true)]));
426        let t = NdArrowArray::try_new(Arc::new(Int32Array::from(vec![7, 8])), dims(&[("t", 2)]))
427            .unwrap();
428        assert!(NdRecordBatch::try_new(schema, vec![t], dims(&[("t", 2)])).is_err());
429    }
430
431    #[test]
432    fn incompatible_column_dims_rejected() {
433        let schema = Arc::new(Schema::new(vec![Field::new("d", DataType::Int32, true)]));
434        let d = NdArrowArray::try_new(Arc::new(Int32Array::from(vec![1, 2])), dims(&[("d", 2)]))
435            .unwrap();
436        assert!(NdRecordBatch::try_new(schema, vec![d], dims(&[("t", 2)])).is_err());
437    }
438}
nd_arrow_array/batch.rs

nd_arrow_array/
batch.rs
