nd_arrow_array/compute/
groupby.rs

1//! Group-by reductions along one dimension, like `xarray.DataArray.groupby`.
2//!
3//! For regular time bins, see [`resample`](fn@crate::compute::resample).
4
5use std::cmp::Ordering;
6
7use arrow::array::make_comparator;
8use arrow::array::{Array, UInt32Array};
9use arrow::compute::{SortOptions, sort_to_indices};
10
11use crate::NdArrowArray;
12use crate::compute::reduce::{Reduction, reduce};
13use crate::compute::search::unique;
14use crate::compute::shape::{concat, expand_dims, take};
15use crate::error::{Result, nd_err};
16
17/// Reduce `a` per group of equal `labels` along `dim`.
18///
19/// `labels` is 1-D with one label per position along `dim`. Null labels are
20/// left out. The result replaces `dim` with one position per unique label,
21/// in sorted order, under the same dimension name. The second value holds
22/// the unique labels along `dim`.
23pub fn groupby_reduce(
24    a: &NdArrowArray,
25    dim: &str,
26    labels: &NdArrowArray,
27    op: Reduction,
28) -> Result<(NdArrowArray, NdArrowArray)> {
29    let axis = a.dims().axis(dim)?;
30    let n = a.dims().as_slice()[axis].size();
31    if labels.rank() != 1 || labels.len() != n {
32        return Err(nd_err!(
33            InvalidArgument,
34            "groupby labels must be 1-D with {n} values along '{dim}'"
35        ));
36    }
37    let groups = unique(labels, dim)?;
38    if groups.is_empty() {
39        return Err(nd_err!(InvalidArgument, "groupby labels are all null"));
40    }
41
42    let options = SortOptions {
43        descending: false,
44        nulls_first: false,
45    };
46    let order: UInt32Array = sort_to_indices(labels.values().as_ref(), Some(options), None)?;
47    let compare = make_comparator(groups.values().as_ref(), labels.values().as_ref(), options)?;
48
49    let mut parts = Vec::with_capacity(groups.len());
50    let mut cursor = 0usize;
51    for g in 0..groups.len() {
52        let mut members: Vec<isize> = Vec::new();
53        while cursor < order.len() {
54            let position = order.value(cursor) as usize;
55            if labels.values().is_null(position) || compare(g, position) != Ordering::Equal {
56                break;
57            }
58            members.push(position as isize);
59            cursor += 1;
60        }
61        members.sort_unstable();
62        let group = take(a, dim, &members)?;
63        let reduced = reduce(&group, Some(&[dim]), op)?;
64        parts.push(expand_dims(&reduced, dim, axis)?);
65    }
66    let result = concat(&parts.iter().collect::<Vec<_>>(), dim)?;
67    Ok((result, groups))
68}
69
70#[cfg(test)]
71mod tests {
72    use std::sync::Arc;
73
74    use arrow::array::{AsArray, Float64Array, StringArray};
75    use arrow::datatypes::Float64Type;
76
77    use super::*;
78    use crate::Dimensions;
79
80    #[test]
81    fn mean_per_label() {
82        let a = NdArrowArray::try_new(
83            Arc::new(Float64Array::from(vec![
84                1.0, 2.0, 3.0, 4.0, 10.0, 20.0, 30.0, 40.0,
85            ])),
86            Dimensions::try_new([("y", 2), ("t", 4)]).unwrap(),
87        )
88        .unwrap();
89        let labels = NdArrowArray::try_new(
90            Arc::new(StringArray::from(vec![
91                Some("b"),
92                Some("a"),
93                Some("b"),
94                None,
95            ])),
96            Dimensions::try_new([("t", 4)]).unwrap(),
97        )
98        .unwrap();
99        let (out, groups) = groupby_reduce(&a, "t", &labels, Reduction::Mean).unwrap();
100        assert_eq!(out.dims().names(), vec!["y", "t"]);
101        assert_eq!(out.shape(), vec![2, 2]);
102        let v: Vec<f64> = out.values().as_primitive::<Float64Type>().values().to_vec();
103        assert_eq!(v, vec![2.0, 2.0, 20.0, 20.0]);
104        let g: Vec<&str> = groups
105            .values()
106            .as_string::<i32>()
107            .iter()
108            .map(Option::unwrap)
109            .collect();
110        assert_eq!(g, vec!["a", "b"]);
111    }
112}
nd_arrow_array/compute/groupby.rs

nd_arrow_array/compute/
groupby.rs
