nd_arrow_array/compute/
window.rs

1//! Moving windows along one dimension, like `xarray.DataArray.rolling`.
2
3use std::sync::Arc;
4
5use arrow::array::{Array, AsArray, Float64Array};
6use arrow::compute::cast as arrow_cast;
7use arrow::datatypes::{DataType, Float64Type};
8
9use crate::NdArrowArray;
10use crate::error::{Result, nd_err};
11
12/// The statistic of each window.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum RollingOp {
15    /// Sum of the window.
16    Sum,
17    /// Mean of the window.
18    Mean,
19    /// Smallest value of the window.
20    Min,
21    /// Largest value of the window.
22    Max,
23    /// Number of non-null values in the window.
24    Count,
25}
26
27/// Apply `op` over a moving window of `window` elements along `dim`.
28///
29/// Like xarray, the window ends at each element, or is centered on it when
30/// `center` is true. A window with fewer than `min_periods` non-null values
31/// gives null (`min_periods` of `None` means `window`). The result has the
32/// shape of `a` and type `Float64` (`Int64` for `Count`).
33pub fn rolling(
34    a: &NdArrowArray,
35    dim: &str,
36    window: usize,
37    op: RollingOp,
38    min_periods: Option<usize>,
39    center: bool,
40) -> Result<NdArrowArray> {
41    if window == 0 {
42        return Err(nd_err!(
43            InvalidArgument,
44            "rolling window must be at least 1"
45        ));
46    }
47    let min_periods = min_periods.unwrap_or(window).max(1);
48    let axis = a.dims().axis(dim)?;
49    let shape = a.shape();
50    let n = shape[axis];
51    let outer: usize = shape[..axis].iter().product();
52    let inner: usize = shape[axis + 1..].iter().product();
53    let values = arrow_cast(a.values().as_ref(), &DataType::Float64)?;
54    let values = values.as_primitive::<Float64Type>();
55    let get = |p: usize| values.is_valid(p).then(|| values.value(p));
56
57    // Window of element i covers [i + 1 - window + shift, i + shift].
58    let shift = if center { window / 2 } else { 0 };
59    let mut out = vec![None; a.len()];
60    let mut lane = Vec::with_capacity(n);
61    for o in 0..outer {
62        for j in 0..inner {
63            lane.clear();
64            lane.extend((0..n).map(|i| get((o * n + i) * inner + j)));
65            for i in 0..n {
66                let end = (i + shift).min(n.saturating_sub(1));
67                let start = (i + shift + 1).saturating_sub(window);
68                let items = lane[start..=end].iter().flatten();
69                let valid = items.clone().count();
70                let p = (o * n + i) * inner + j;
71                out[p] = match op {
72                    RollingOp::Count => Some(valid as f64),
73                    _ if valid < min_periods => None,
74                    RollingOp::Sum => Some(items.sum()),
75                    RollingOp::Mean => Some(items.sum::<f64>() / valid as f64),
76                    RollingOp::Min => items.copied().reduce(f64::min),
77                    RollingOp::Max => items.copied().reduce(f64::max),
78                };
79            }
80        }
81    }
82    let result: Arc<dyn Array> = if op == RollingOp::Count {
83        arrow_cast(&Float64Array::from(out), &DataType::Int64)?
84    } else {
85        Arc::new(Float64Array::from(out))
86    };
87    NdArrowArray::try_new(result, a.dims().clone())
88}
89
90#[cfg(test)]
91mod tests {
92    use arrow::array::{ArrayRef, Int32Array};
93
94    use super::*;
95    use crate::Dimensions;
96
97    fn lane(values: Vec<Option<i32>>) -> NdArrowArray {
98        let n = values.len();
99        NdArrowArray::try_new(
100            Arc::new(Int32Array::from(values)) as ArrayRef,
101            Dimensions::try_new([("t", n)]).unwrap(),
102        )
103        .unwrap()
104    }
105
106    fn f64s(a: &NdArrowArray) -> Vec<Option<f64>> {
107        a.values().as_primitive::<Float64Type>().iter().collect()
108    }
109
110    #[test]
111    fn trailing_mean_matches_xarray() {
112        let a = lane(vec![Some(1), Some(2), Some(3), Some(4)]);
113        let m = rolling(&a, "t", 2, RollingOp::Mean, None, false).unwrap();
114        assert_eq!(f64s(&m), vec![None, Some(1.5), Some(2.5), Some(3.5)]);
115    }
116
117    #[test]
118    fn centered_sum_with_min_periods() {
119        let a = lane(vec![Some(1), None, Some(3), Some(4)]);
120        let s = rolling(&a, "t", 3, RollingOp::Sum, Some(1), true).unwrap();
121        assert_eq!(f64s(&s), vec![Some(1.0), Some(4.0), Some(7.0), Some(7.0)]);
122        let c = rolling(&a, "t", 3, RollingOp::Count, None, true).unwrap();
123        assert_eq!(c.data_type(), &DataType::Int64);
124    }
125}
nd_arrow_array/compute/window.rs

nd_arrow_array/compute/
window.rs
