nd_arrow_array/compute/
batch.rs

1//! Kernels on [`NdRecordBatch`]: materialize, select by position or by
2//! coordinate value, interpolate and resample.
3//!
4//! A coordinate is a 1-D column with the same name as its dimension, like xarray.
5
6use std::cmp::Ordering;
7use std::sync::Arc;
8
9use arrow::array::{Array, ArrayRef, AsArray, make_comparator};
10use arrow::compute::{SortOptions, cast as arrow_cast};
11use arrow::datatypes::{DataType, Field, Float64Type, Schema};
12use arrow::record_batch::{RecordBatch, RecordBatchOptions};
13
14use crate::compute::interp::interp;
15use crate::compute::reduce::Reduction;
16use crate::compute::resample::{Frequency, resample};
17use crate::compute::search::partition_point;
18use crate::compute::shape::{Selector, Slice, broadcast_to, isel, squeeze};
19use crate::dimensions::{Dimension, Dimensions};
20use crate::error::{Result, nd_err};
21use crate::{NdArrowArray, NdRecordBatch};
22
23impl NdRecordBatch {
24    /// Broadcast every column onto the target grid and return a flat
25    /// [`RecordBatch`] with one row per grid cell, in C order.
26    pub fn materialize(&self) -> Result<RecordBatch> {
27        let columns = self
28            .columns()
29            .iter()
30            .map(|c| Ok(broadcast_to(c, self.target())?.values().clone()))
31            .collect::<Result<Vec<ArrayRef>>>()?;
32        let options = RecordBatchOptions::new().with_row_count(Some(self.num_rows()));
33        Ok(RecordBatch::try_new_with_options(
34            self.schema().clone(),
35            columns,
36            &options,
37        )?)
38    }
39
40    /// The coordinate of `dim`: a column named `dim` with dimensions `(dim)`.
41    pub fn coordinate(&self, dim: &str) -> Option<&NdArrowArray> {
42        self.column_by_name(dim)
43            .filter(|c| c.rank() == 1 && c.dims().as_slice()[0].name() == dim)
44    }
45
46    /// Select by position on target dimensions, like xarray `Dataset.isel`.
47    /// The batch kernels keep the global attributes.
48    ///
49    /// Each column gets the selections of its own dimensions. A column with
50    /// size 1 on a selected dimension keeps broadcasting, and an index
51    /// selection removes that axis.
52    pub fn isel(&self, selections: &[(&str, Selector)]) -> Result<Self> {
53        let mut target: Vec<Dimension> = self.target().iter().cloned().collect();
54        let mut resolved: Vec<(&str, Selector, usize)> = Vec::new();
55        for (name, selector) in selections {
56            let axis = self.target().axis(name)?;
57            let size = target[axis].size();
58            let count = match selector {
59                Selector::Slice(slice) => slice.resolve(size)?.1,
60                Selector::Index(_) => 0,
61                Selector::Take(indices) => indices.len(),
62            };
63            resolved.push((name, selector.clone(), count));
64            target[axis] = Dimension::new(*name, count);
65        }
66        let dropped: Vec<&str> = resolved
67            .iter()
68            .filter(|(_, s, _)| matches!(s, Selector::Index(_)))
69            .map(|(n, _, _)| *n)
70            .collect();
71        target.retain(|d| !dropped.contains(&d.name()));
72
73        let columns = self
74            .columns()
75            .iter()
76            .map(|column| {
77                let mut own = Vec::new();
78                let mut squeezed = Vec::new();
79                for (name, selector, count) in &resolved {
80                    let Some(size) = column.dims().size_of(name) else {
81                        continue;
82                    };
83                    let full = self.target().size_of(name) == Some(size);
84                    if full {
85                        own.push((*name, selector.clone()));
86                    } else if matches!(selector, Selector::Index(_)) {
87                        squeezed.push(*name);
88                    } else if *count == 0 {
89                        own.push((*name, Selector::Take(Vec::new())));
90                    }
91                }
92                let selected = isel(column, &own)?;
93                if squeezed.is_empty() {
94                    Ok(selected)
95                } else {
96                    squeeze(&selected, Some(&squeezed))
97                }
98            })
99            .collect::<Result<Vec<_>>>()?;
100        NdRecordBatch::try_new(self.schema().clone(), columns, Dimensions::try_new(target)?)
101            .map(|b| b.with_attributes(self.attributes().clone()))
102    }
103
104    /// Select the positions whose coordinate lies in `[start, end]`, like
105    /// xarray `sel` with a label slice. An open end is `None`.
106    ///
107    /// The coordinate must be sorted, ascending or descending, without nulls.
108    pub fn sel_range(
109        &self,
110        dim: &str,
111        start: Option<&ArrayRef>,
112        end: Option<&ArrayRef>,
113    ) -> Result<Self> {
114        let coordinate = self.require_coordinate(dim)?;
115        let values = coordinate.values();
116        let n = values.len();
117        let options = SortOptions::default();
118        let self_compare = make_comparator(values.as_ref(), values.as_ref(), options)?;
119        let ascending = (1..n).all(|i| self_compare(i - 1, i) != Ordering::Greater);
120        let descending = (1..n).all(|i| self_compare(i - 1, i) != Ordering::Less);
121        if values.null_count() > 0 || !(ascending || descending) {
122            return Err(nd_err!(
123                InvalidArgument,
124                "coordinate '{dim}' must be sorted and without nulls"
125            ));
126        }
127        let bound = |value: Option<&ArrayRef>| -> Result<Option<ArrayRef>> {
128            value
129                .map(|v| Ok(arrow_cast(v.slice(0, 1).as_ref(), values.data_type())?))
130                .transpose()
131        };
132        let (start, end) = (bound(start)?, bound(end)?);
133        let position = |value: &Option<ArrayRef>,
134                        pred: fn(Ordering) -> bool,
135                        default: usize|
136         -> Result<usize> {
137            Ok(match value {
138                None => default,
139                Some(v) => {
140                    let compare = make_comparator(values.as_ref(), v.as_ref(), options)?;
141                    partition_point(n, |i| pred(compare(i, 0)))
142                }
143            })
144        };
145        let (first, stop) = if ascending {
146            (
147                position(&start, |o| o == Ordering::Less, 0)?,
148                position(&end, |o| o != Ordering::Greater, n)?,
149            )
150        } else {
151            (
152                position(&end, |o| o == Ordering::Greater, 0)?,
153                position(&start, |o| o != Ordering::Less, n)?,
154            )
155        };
156        let stop = stop.max(first);
157        self.isel(&[(
158            dim,
159            Selector::Slice(Slice::range(first as isize, stop as isize)),
160        )])
161    }
162
163    /// Select the position whose coordinate is nearest to `value`, like xarray
164    /// `sel(method="nearest")`. The dimension is removed. Ties pick the first.
165    pub fn sel_nearest(&self, dim: &str, value: &ArrayRef) -> Result<Self> {
166        let coordinate = self.require_coordinate(dim)?;
167        let coords = arrow_cast(coordinate.values().as_ref(), &DataType::Float64)?;
168        let target = arrow_cast(value.slice(0, 1).as_ref(), &DataType::Float64)?;
169        let target = target.as_primitive::<Float64Type>();
170        if target.is_null(0) {
171            return Err(nd_err!(InvalidArgument, "nearest selection value is null"));
172        }
173        let t = target.value(0);
174        let best = coords
175            .as_primitive::<Float64Type>()
176            .iter()
177            .enumerate()
178            .filter_map(|(i, c)| c.map(|c| (i, (c - t).abs())))
179            .min_by(|a, b| a.1.total_cmp(&b.1))
180            .ok_or_else(|| nd_err!(InvalidArgument, "coordinate '{dim}' has no values"))?;
181        self.isel(&[(dim, Selector::Index(best.0 as isize))])
182    }
183
184    /// Interpolate every column along `dim` onto `new_coords`, see [`interp`].
185    ///
186    /// Columns that span `dim` must be numeric and become floating point.
187    /// The coordinate becomes `new_coords`. Other columns stay unchanged.
188    pub fn interp(&self, dim: &str, new_coords: &NdArrowArray) -> Result<Self> {
189        let coordinate = self.require_coordinate(dim)?.clone();
190        if new_coords.rank() != 1 || new_coords.dims().as_slice()[0].name() != dim {
191            return Err(nd_err!(
192                InvalidArgument,
193                "new coordinates must be 1-D along '{dim}'"
194            ));
195        }
196        let size = self.target().size_of(dim).unwrap_or(0);
197        let mut fields: Vec<Field> = Vec::with_capacity(self.num_columns());
198        let mut columns = Vec::with_capacity(self.num_columns());
199        for (field, column) in self.schema().fields().iter().zip(self.columns()) {
200            let new_column = if field.name() == dim {
201                new_coords.clone()
202            } else if column.dims().size_of(dim) == Some(size) && size > 1 {
203                interp(column, dim, &coordinate, new_coords)?
204            } else {
205                column.clone()
206            };
207            fields.push(
208                field
209                    .as_ref()
210                    .clone()
211                    .with_data_type(new_column.data_type().clone()),
212            );
213            columns.push(new_column);
214        }
215        let target = self.target().iter().map(|d| {
216            if d.name() == dim {
217                Dimension::new(dim, new_coords.len())
218            } else {
219                d.clone()
220            }
221        });
222        let schema = Schema::new_with_metadata(fields, self.schema().metadata().clone());
223        NdRecordBatch::try_new(Arc::new(schema), columns, Dimensions::try_new(target)?)
224            .map(|b| b.with_attributes(self.attributes().clone()))
225    }
226
227    /// Resample every column along `dim` with the time coordinate of `dim`,
228    /// see [`resample`].
229    ///
230    /// Columns that span `dim` are reduced with `op`. The coordinate becomes
231    /// the bin starts. Other columns stay unchanged.
232    pub fn resample(&self, dim: &str, freq: Frequency, op: Reduction) -> Result<Self> {
233        let coordinate = self.require_coordinate(dim)?.clone();
234        let size = self.target().size_of(dim).unwrap_or(0);
235        let (_, labels) = resample(&coordinate, dim, &coordinate, freq, Reduction::Count)?;
236        let mut fields: Vec<Field> = Vec::with_capacity(self.num_columns());
237        let mut columns = Vec::with_capacity(self.num_columns());
238        for (field, column) in self.schema().fields().iter().zip(self.columns()) {
239            let new_column = if field.name() == dim {
240                labels.clone()
241            } else if column.dims().size_of(dim) == Some(size) {
242                resample(column, dim, &coordinate, freq, op)?.0
243            } else {
244                column.clone()
245            };
246            fields.push(
247                field
248                    .as_ref()
249                    .clone()
250                    .with_data_type(new_column.data_type().clone()),
251            );
252            columns.push(new_column);
253        }
254        let target = self.target().iter().map(|d| {
255            if d.name() == dim {
256                Dimension::new(dim, labels.len())
257            } else {
258                d.clone()
259            }
260        });
261        let schema = Schema::new_with_metadata(fields, self.schema().metadata().clone());
262        NdRecordBatch::try_new(Arc::new(schema), columns, Dimensions::try_new(target)?)
263            .map(|b| b.with_attributes(self.attributes().clone()))
264    }
265
266    fn require_coordinate(&self, dim: &str) -> Result<&NdArrowArray> {
267        self.target().axis(dim)?;
268        self.coordinate(dim).ok_or_else(|| {
269            nd_err!(
270                InvalidArgument,
271                "dimension '{dim}' has no coordinate column"
272            )
273        })
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use arrow::array::{Float64Array, Int32Array, TimestampMicrosecondArray};
280    use arrow::datatypes::{Float64Type, Int32Type, TimeUnit, TimestampMicrosecondType};
281
282    use super::*;
283    use crate::{Attributes, IntoScalar};
284
285    fn batch() -> NdRecordBatch {
286        let schema = Arc::new(Schema::new(vec![
287            Field::new("time", DataType::Int32, true),
288            Field::new("lat", DataType::Float64, true),
289            Field::new("sst", DataType::Float64, true),
290            Field::new("flag", DataType::Int32, true),
291        ]));
292        let d = |spec: &[(&str, usize)]| Dimensions::try_new(spec.iter().copied()).unwrap();
293        let time =
294            NdArrowArray::try_new(Arc::new(Int32Array::from(vec![10, 20])), d(&[("time", 2)]))
295                .unwrap();
296        let lat = NdArrowArray::try_new(
297            Arc::new(Float64Array::from(vec![30.0, 20.0, 10.0])),
298            d(&[("lat", 3)]),
299        )
300        .unwrap();
301        let sst = NdArrowArray::try_new(
302            Arc::new(Float64Array::from(vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0])),
303            d(&[("time", 2), ("lat", 3)]),
304        )
305        .unwrap();
306        let flag =
307            NdArrowArray::try_new(Arc::new(Int32Array::from(vec![7])), d(&[("time", 1)])).unwrap();
308        NdRecordBatch::try_new_inferred(schema, vec![time, lat, sst, flag]).unwrap()
309    }
310
311    #[test]
312    fn resample_reduces_columns_on_the_time_axis() {
313        let day = 86_400_000_000i64;
314        let d = |spec: &[(&str, usize)]| Dimensions::try_new(spec.iter().copied()).unwrap();
315        let schema = Arc::new(Schema::new(vec![
316            Field::new(
317                "time",
318                DataType::Timestamp(TimeUnit::Microsecond, None),
319                true,
320            ),
321            Field::new("lat", DataType::Float64, true),
322            Field::new("sst", DataType::Float64, true),
323        ]));
324        let time = NdArrowArray::try_new(
325            Arc::new(TimestampMicrosecondArray::from(vec![0, day / 2, 3 * day])),
326            d(&[("time", 3)]),
327        )
328        .unwrap();
329        let lat = NdArrowArray::try_new(
330            Arc::new(Float64Array::from(vec![10.0, 20.0])),
331            d(&[("lat", 2)]),
332        )
333        .unwrap();
334        let sst = NdArrowArray::try_new(
335            Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0])),
336            d(&[("time", 3), ("lat", 2)]),
337        )
338        .unwrap();
339        let b = NdRecordBatch::try_new_inferred(schema, vec![time, lat, sst])
340            .unwrap()
341            .with_attributes(Attributes::new().with("title", "test"));
342        let freq = Frequency::parse("2D").unwrap();
343        let r = b.resample("time", freq, Reduction::Mean).unwrap();
344        assert_eq!(r.target(), &d(&[("time", 2), ("lat", 2)]));
345        let time: Vec<i64> = r
346            .column(0)
347            .values()
348            .as_primitive::<TimestampMicrosecondType>()
349            .values()
350            .to_vec();
351        assert_eq!(time, vec![0, 2 * day]);
352        assert_eq!(r.column(1), b.column(1));
353        let sst: Vec<f64> = r
354            .column(2)
355            .values()
356            .as_primitive::<Float64Type>()
357            .values()
358            .to_vec();
359        assert_eq!(sst, vec![2.0, 3.0, 5.0, 6.0]);
360        assert_eq!(r.attributes().get_str("title"), Some("test"));
361    }
362
363    #[test]
364    fn materialize_broadcasts_columns() {
365        let rb = batch().materialize().unwrap();
366        assert_eq!(rb.num_rows(), 6);
367        let time: Vec<i32> = rb.column(0).as_primitive::<Int32Type>().values().to_vec();
368        assert_eq!(time, vec![10, 10, 10, 20, 20, 20]);
369        let lat: Vec<f64> = rb.column(1).as_primitive::<Float64Type>().values().to_vec();
370        assert_eq!(lat, vec![30.0, 20.0, 10.0, 30.0, 20.0, 10.0]);
371        let flag: Vec<i32> = rb.column(3).as_primitive::<Int32Type>().values().to_vec();
372        assert_eq!(flag, vec![7; 6]);
373    }
374
375    #[test]
376    fn isel_updates_target_and_columns() {
377        let b = batch()
378            .isel(&[
379                ("time", Selector::Index(1)),
380                ("lat", Slice::range(0, 2).into()),
381            ])
382            .unwrap();
383        assert_eq!(b.target(), &Dimensions::try_new([("lat", 2)]).unwrap());
384        assert!(b.column(0).is_scalar());
385        assert!(
386            b.column(3).is_scalar(),
387            "the broadcast axis of flag is squeezed"
388        );
389        let sst: Vec<f64> = b
390            .column(2)
391            .values()
392            .as_primitive::<Float64Type>()
393            .values()
394            .to_vec();
395        assert_eq!(sst, vec![3.0, 4.0]);
396    }
397
398    #[test]
399    fn sel_range_on_descending_coordinate() {
400        let b = batch()
401            .sel_range(
402                "lat",
403                Some(&15.0f64.into_scalar()),
404                Some(&30.0f64.into_scalar()),
405            )
406            .unwrap();
407        let lat: Vec<f64> = b
408            .column(1)
409            .values()
410            .as_primitive::<Float64Type>()
411            .values()
412            .to_vec();
413        assert_eq!(lat, vec![30.0, 20.0]);
414        let t = batch()
415            .sel_range("time", Some(&15i32.into_scalar()), None)
416            .unwrap();
417        assert_eq!(t.target().size_of("time"), Some(1));
418    }
419
420    #[test]
421    fn sel_nearest_drops_dim() {
422        let b = batch().sel_nearest("lat", &22.0f64.into_scalar()).unwrap();
423        assert!(!b.target().contains("lat"));
424        let sst: Vec<f64> = b
425            .column(2)
426            .values()
427            .as_primitive::<Float64Type>()
428            .values()
429            .to_vec();
430        assert_eq!(sst, vec![1.0, 4.0]);
431    }
432
433    #[test]
434    fn kernels_keep_global_attributes() {
435        let b = batch().with_attributes(crate::Attributes::new().with("title", "test"));
436        let kept = [
437            b.isel(&[("lat", Selector::Index(0))]).unwrap(),
438            b.sel_nearest("lat", &22.0f64.into_scalar()).unwrap(),
439            b.sel_range("time", None, Some(&10i32.into_scalar()))
440                .unwrap(),
441        ];
442        for k in kept {
443            assert_eq!(k.attributes().get_str("title"), Some("test"));
444        }
445    }
446
447    #[test]
448    fn interp_on_time() {
449        let new_time = NdArrowArray::try_new(
450            Arc::new(Int32Array::from(vec![15])),
451            Dimensions::try_new([("time", 1)]).unwrap(),
452        )
453        .unwrap();
454        let b = batch().interp("time", &new_time).unwrap();
455        let sst: Vec<f64> = b
456            .column(2)
457            .values()
458            .as_primitive::<Float64Type>()
459            .values()
460            .to_vec();
461        assert_eq!(sst, vec![1.5, 2.5, 3.5]);
462        assert_eq!(b.schema().field(0).data_type(), &DataType::Int32);
463    }
464}
nd_arrow_array/compute/batch.rs

nd_arrow_array/compute/
batch.rs
