nd_arrow_array/compute/
resample.rs

1//! Resampling along a time dimension, like `xarray.DataArray.resample`.
2//!
3//! Each time goes into one bin of a regular [`Frequency`]. Bins are closed on
4//! the left and are labeled by their start, like the pandas defaults:
5//!
6//! - Fixed periods (day, hour, minute, second) count from midnight of the day
7//!   of the first time, like pandas `origin="start_day"`.
8//! - Calendar periods (month, year) count from the month or year of the first
9//!   time, like pandas `MS` and `YS`.
10
11use std::str::FromStr;
12use std::sync::Arc;
13
14use arrow::array::{TimestampMicrosecondArray, new_null_array};
15
16use crate::NdArrowArray;
17use crate::compute::reduce::{Reduction, reduce};
18use crate::compute::shape::{Selector, Slice, concat, expand_dims, isel};
19use crate::compute::time::{
20    MICROS_PER_DAY, MICROS_PER_SECOND, TimePeriod, civil_from_days, days_from_civil, to_micros,
21};
22use crate::error::{Result, nd_err};
23
24/// Upper limit of the number of bins, against memory exhaustion.
25pub const MAX_BINS: usize = u32::MAX as usize;
26
27/// A regular time step: `step` times a calendar `period`, for example 6 hours.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub struct Frequency {
30    step: u32,
31    period: TimePeriod,
32}
33
34impl Frequency {
35    /// A frequency of `step` periods. `step` must be at least 1.
36    pub fn new(step: u32, period: TimePeriod) -> Result<Self> {
37        if step == 0 {
38            return Err(nd_err!(
39                InvalidArgument,
40                "frequency step must be at least 1"
41            ));
42        }
43        Ok(Self { step, period })
44    }
45
46    /// Parse a pandas frequency string, for example `"6h"`, `"1D"` or `"3MS"`.
47    ///
48    /// The aliases are `YS`, `MS`, `D`, `h`, `min` and `s`, with an optional
49    /// step in front. End-anchored aliases such as `M` and `Y` are not supported.
50    pub fn parse(text: &str) -> Result<Self> {
51        let text = text.trim();
52        let split = text
53            .find(|c: char| !c.is_ascii_digit())
54            .unwrap_or(text.len());
55        let (digits, alias) = text.split_at(split);
56        let step = match digits {
57            "" => 1,
58            _ => digits
59                .parse()
60                .map_err(|_| nd_err!(InvalidArgument, "invalid frequency step in '{text}'"))?,
61        };
62        let period = match alias {
63            "YS" => TimePeriod::Year,
64            "MS" => TimePeriod::Month,
65            "D" => TimePeriod::Day,
66            "h" => TimePeriod::Hour,
67            "min" => TimePeriod::Minute,
68            "s" => TimePeriod::Second,
69            _ => {
70                return Err(nd_err!(
71                    InvalidArgument,
72                    "unsupported frequency '{text}': use YS, MS, D, h, min or s"
73                ));
74            }
75        };
76        Self::new(step, period)
77    }
78
79    /// The number of periods per bin.
80    pub fn step(&self) -> u32 {
81        self.step
82    }
83
84    /// The calendar period.
85    pub fn period(&self) -> TimePeriod {
86        self.period
87    }
88}
89
90impl FromStr for Frequency {
91    type Err = crate::NdError;
92
93    fn from_str(text: &str) -> Result<Self> {
94        Self::parse(text)
95    }
96}
97
98/// Reduce `a` per time bin along `dim`.
99///
100/// `times` is 1-D with one timestamp (or date) per position along `dim`, in
101/// any order. Null times are left out. The result replaces `dim` with one
102/// position per bin, from the bin of the first time to the bin of the last.
103/// A bin without times gives null, like xarray. The second value holds the
104/// bin starts as `Timestamp(Microsecond)` along `dim`, with the attributes of
105/// `times`.
106pub fn resample(
107    a: &NdArrowArray,
108    dim: &str,
109    times: &NdArrowArray,
110    freq: Frequency,
111    op: Reduction,
112) -> Result<(NdArrowArray, NdArrowArray)> {
113    let axis = a.dims().axis(dim)?;
114    let n = a.dims().as_slice()[axis].size();
115    if times.rank() != 1 || times.len() != n {
116        return Err(nd_err!(
117            InvalidArgument,
118            "resample times must be 1-D with {n} values along '{dim}'"
119        ));
120    }
121    let stamps = to_micros(times)?;
122    let (bin_of, labels) = assign_bins(&stamps, freq)?;
123
124    // Counting sort: the positions of bin b are positions[offsets[b]..offsets[b + 1]].
125    let mut offsets = vec![0usize; labels.len() + 1];
126    for bin in bin_of.iter().flatten() {
127        offsets[bin + 1] += 1;
128    }
129    for b in 0..labels.len() {
130        offsets[b + 1] += offsets[b];
131    }
132    let mut cursor = offsets.clone();
133    let mut positions = vec![0isize; offsets[labels.len()]];
134    for (position, bin) in bin_of.iter().enumerate() {
135        if let Some(bin) = bin {
136            positions[cursor[*bin]] = position as isize;
137            cursor[*bin] += 1;
138        }
139    }
140
141    let mut parts: Vec<Option<NdArrowArray>> = Vec::with_capacity(labels.len());
142    for b in 0..labels.len() {
143        let members = &positions[offsets[b]..offsets[b + 1]];
144        let (Some(&first), Some(&last)) = (members.first(), members.last()) else {
145            parts.push(None);
146            continue;
147        };
148        let selector = if (last - first) as usize + 1 == members.len() {
149            Selector::Slice(Slice::range(first, last + 1))
150        } else {
151            Selector::Take(members.to_vec())
152        };
153        let group = isel(a, &[(dim, selector)])?;
154        let reduced = reduce(&group, Some(&[dim]), op)?;
155        parts.push(Some(expand_dims(&reduced, dim, axis)?));
156    }
157
158    let template = parts
159        .iter()
160        .flatten()
161        .next()
162        .ok_or_else(|| nd_err!(InvalidArgument, "resample times are all null"))?;
163    let empty = NdArrowArray::try_new(
164        new_null_array(template.data_type(), template.len()),
165        template.dims().clone(),
166    )?;
167    let parts: Vec<&NdArrowArray> = parts.iter().map(|p| p.as_ref().unwrap_or(&empty)).collect();
168    let result = concat(&parts, dim)?;
169
170    let labels = NdArrowArray::try_new(
171        Arc::new(TimestampMicrosecondArray::from(labels)),
172        crate::Dimensions::try_new([(dim, parts.len())])?,
173    )?
174    .with_attributes(times.attributes().clone());
175    Ok((result, labels))
176}
177
178/// The bin of each time (`None` for a null time) and the start of each bin.
179fn assign_bins(
180    stamps: &TimestampMicrosecondArray,
181    freq: Frequency,
182) -> Result<(Vec<Option<usize>>, Vec<i64>)> {
183    let (Some(min), Some(max)) = (arrow::compute::min(stamps), arrow::compute::max(stamps)) else {
184        return Err(nd_err!(InvalidArgument, "resample times are all null"));
185    };
186    let step = i128::from(freq.step);
187    let unit = match freq.period {
188        TimePeriod::Second => Some(MICROS_PER_SECOND),
189        TimePeriod::Minute => Some(60 * MICROS_PER_SECOND),
190        TimePeriod::Hour => Some(3_600 * MICROS_PER_SECOND),
191        TimePeriod::Day => Some(MICROS_PER_DAY),
192        TimePeriod::Month | TimePeriod::Year => None,
193    };
194    let scheme = match unit {
195        Some(unit) => Bins::Fixed {
196            origin: i128::from(min.div_euclid(MICROS_PER_DAY) * MICROS_PER_DAY),
197            width: step * i128::from(unit),
198        },
199        None => {
200            let first = month_of(min);
201            let year = freq.period == TimePeriod::Year;
202            Bins::Calendar {
203                base: if year {
204                    first - first.rem_euclid(12)
205                } else {
206                    first
207                },
208                months: if year { 12 * step } else { step },
209            }
210        }
211    };
212    let index = |t: i64| scheme.index(t);
213    let label = |k: i128| scheme.label(k);
214
215    // The origin can lie bins before the first time; the output starts at its bin.
216    let first = index(min);
217    let count = usize::try_from(index(max) - first + 1)
218        .ok()
219        .filter(|&c| c <= MAX_BINS)
220        .ok_or_else(|| nd_err!(InvalidArgument, "resample needs more than {MAX_BINS} bins"))?;
221    let bins = stamps
222        .iter()
223        .map(|t| t.map(|t| (index(t) - first) as usize))
224        .collect();
225    let labels = (0..count as i128).map(|k| label(first + k)).collect();
226    Ok((bins, labels))
227}
228
229/// Bin layout: fixed-width bins from an origin, or whole months from a base month.
230enum Bins {
231    Fixed { origin: i128, width: i128 },
232    Calendar { base: i128, months: i128 },
233}
234
235impl Bins {
236    /// The bin index of time `t`, counted from the origin.
237    fn index(&self, t: i64) -> i128 {
238        match *self {
239            Bins::Fixed { origin, width } => (i128::from(t) - origin) / width,
240            Bins::Calendar { base, months } => (month_of(t) - base) / months,
241        }
242    }
243
244    /// The start of bin `k`.
245    fn label(&self, k: i128) -> i64 {
246        match *self {
247            Bins::Fixed { origin, width } => (origin + k * width) as i64,
248            Bins::Calendar { base, months } => {
249                let month = base + k * months;
250                let (year, month) = (month.div_euclid(12), month.rem_euclid(12) as u32 + 1);
251                days_from_civil(year as i64, month, 1) * MICROS_PER_DAY
252            }
253        }
254    }
255}
256
257/// Months since year 0 of time `t`.
258fn month_of(t: i64) -> i128 {
259    let (year, month, _) = civil_from_days(t.div_euclid(MICROS_PER_DAY));
260    i128::from(year) * 12 + i128::from(month) - 1
261}
262
263#[cfg(test)]
264mod tests {
265    use arrow::array::{AsArray, Float64Array, Int32Array};
266    use arrow::datatypes::{Float64Type, Int64Type, TimestampMicrosecondType};
267
268    use super::*;
269    use crate::{Attributes, Dimensions};
270
271    fn at(year: i64, month: u32, day: u32, hour: i64) -> i64 {
272        days_from_civil(year, month, day) * MICROS_PER_DAY + hour * 3_600 * MICROS_PER_SECOND
273    }
274
275    fn time_axis(stamps: Vec<Option<i64>>) -> NdArrowArray {
276        let n = stamps.len();
277        NdArrowArray::try_new(
278            Arc::new(TimestampMicrosecondArray::from(stamps)),
279            Dimensions::try_new([("time", n)]).unwrap(),
280        )
281        .unwrap()
282    }
283
284    fn f64s(a: &NdArrowArray) -> Vec<Option<f64>> {
285        a.values().as_primitive::<Float64Type>().iter().collect()
286    }
287
288    fn stamps(a: &NdArrowArray) -> Vec<i64> {
289        a.values()
290            .as_primitive::<TimestampMicrosecondType>()
291            .values()
292            .to_vec()
293    }
294
295    // Expected values come from xarray 2026.9 `DataArray.resample`.
296    fn hourly() -> (NdArrowArray, NdArrowArray) {
297        let times = time_axis(vec![
298            Some(at(2021, 1, 1, 1)),
299            Some(at(2021, 1, 1, 2)),
300            Some(at(2021, 1, 1, 5)),
301            Some(at(2021, 1, 3, 23)),
302            Some(at(2021, 1, 4, 0)),
303        ]);
304        let values = NdArrowArray::try_new(
305            Arc::new(Float64Array::from(vec![
306                Some(1.0),
307                None,
308                Some(2.0),
309                Some(4.0),
310                Some(8.0),
311            ])),
312            times.dims().clone(),
313        )
314        .unwrap();
315        (values, times)
316    }
317
318    #[test]
319    fn daily_bins_match_xarray() {
320        let (values, times) = hourly();
321        let freq = Frequency::parse("1D").unwrap();
322        let (sum, labels) = resample(&values, "time", &times, freq, Reduction::Sum).unwrap();
323        assert_eq!(f64s(&sum), vec![Some(3.0), None, Some(4.0), Some(8.0)]);
324        assert_eq!(
325            stamps(&labels),
326            vec![
327                at(2021, 1, 1, 0),
328                at(2021, 1, 2, 0),
329                at(2021, 1, 3, 0),
330                at(2021, 1, 4, 0)
331            ]
332        );
333        let (mean, _) = resample(&values, "time", &times, freq, Reduction::Mean).unwrap();
334        assert_eq!(f64s(&mean), vec![Some(1.5), None, Some(4.0), Some(8.0)]);
335        let (count, _) = resample(&values, "time", &times, freq, Reduction::Count).unwrap();
336        let count: Vec<Option<i64>> = count.values().as_primitive::<Int64Type>().iter().collect();
337        assert_eq!(count, vec![Some(2), None, Some(1), Some(1)]);
338
339        let two_days = Frequency::parse("2D").unwrap();
340        let (sum, _) = resample(&values, "time", &times, two_days, Reduction::Sum).unwrap();
341        assert_eq!(f64s(&sum), vec![Some(3.0), Some(12.0)]);
342    }
343
344    #[test]
345    fn fixed_bins_start_at_midnight_of_the_first_day() {
346        let times = time_axis(vec![Some(at(1969, 12, 31, 23)), Some(at(1970, 1, 1, 1))]);
347        let values = NdArrowArray::try_new(
348            Arc::new(Float64Array::from(vec![1.0, 2.0])),
349            times.dims().clone(),
350        )
351        .unwrap();
352        let freq = Frequency::parse("7h").unwrap();
353        let (sum, labels) = resample(&values, "time", &times, freq, Reduction::Sum).unwrap();
354        assert_eq!(f64s(&sum), vec![Some(3.0)]);
355        assert_eq!(stamps(&labels), vec![at(1969, 12, 31, 21)]);
356    }
357
358    #[test]
359    fn calendar_bins_match_xarray() {
360        let times = time_axis(vec![
361            Some(at(2020, 11, 15, 0)),
362            Some(at(2021, 2, 10, 0)),
363            Some(at(2021, 3, 1, 0)),
364            Some(at(2022, 7, 4, 0)),
365        ]);
366        let values = NdArrowArray::try_new(
367            Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])),
368            times.dims().clone(),
369        )
370        .unwrap();
371        let sum = |freq: &str| {
372            let freq = Frequency::parse(freq).unwrap();
373            let (sum, labels) = resample(&values, "time", &times, freq, Reduction::Sum).unwrap();
374            (f64s(&sum), stamps(&labels))
375        };
376        let (v, l) = sum("3MS");
377        assert_eq!(
378            v,
379            vec![Some(1.0), Some(5.0), None, None, None, None, Some(4.0)]
380        );
381        assert_eq!(l[0], at(2020, 11, 1, 0));
382        assert_eq!(l[1], at(2021, 2, 1, 0));
383        let (v, l) = sum("YS");
384        assert_eq!(v, vec![Some(1.0), Some(5.0), Some(4.0)]);
385        assert_eq!(l[2], at(2022, 1, 1, 0));
386        let (v, l) = sum("2YS");
387        assert_eq!(v, vec![Some(6.0), Some(4.0)]);
388        assert_eq!(l, vec![at(2020, 1, 1, 0), at(2022, 1, 1, 0)]);
389    }
390
391    #[test]
392    fn unsorted_and_null_times_on_a_grid() {
393        let times = time_axis(vec![Some(at(2021, 1, 2, 6)), None, Some(at(2021, 1, 1, 6))])
394            .with_attributes(Attributes::new().with("long_name", "time"));
395        let values = NdArrowArray::try_new(
396            Arc::new(Int32Array::from_iter_values(0..6)),
397            Dimensions::try_new([("y", 2), ("time", 3)]).unwrap(),
398        )
399        .unwrap();
400        let freq = Frequency::new(1, TimePeriod::Day).unwrap();
401        let (sum, labels) = resample(&values, "time", &times, freq, Reduction::Sum).unwrap();
402        assert_eq!(sum.dims().names(), vec!["y", "time"]);
403        assert_eq!(
404            sum.values().as_primitive::<Int64Type>().values(),
405            &[2, 0, 5, 3]
406        );
407        assert_eq!(labels.attributes().get_str("long_name"), Some("time"));
408    }
409
410    #[test]
411    fn invalid_input_rejected() {
412        assert!(Frequency::parse("M").is_err());
413        assert!(Frequency::parse("0D").is_err());
414        assert_eq!(
415            "15min".parse::<Frequency>().unwrap(),
416            Frequency::new(15, TimePeriod::Minute).unwrap()
417        );
418        let times = time_axis(vec![None, None]);
419        let values = NdArrowArray::try_new(
420            Arc::new(Float64Array::from(vec![1.0, 2.0])),
421            times.dims().clone(),
422        )
423        .unwrap();
424        let freq = Frequency::parse("D").unwrap();
425        assert!(resample(&values, "time", &times, freq, Reduction::Sum).is_err());
426        let seconds = time_axis(vec![Some(i64::MIN / 2), Some(i64::MAX / 2)]);
427        let freq = Frequency::parse("s").unwrap();
428        assert!(resample(&values, "time", &seconds, freq, Reduction::Sum).is_err());
429    }
430}
nd_arrow_array/compute/resample.rs

nd_arrow_array/compute/
resample.rs
