nd_arrow_array/compute/
reduce.rs

1//! Reductions along named axes, cumulative operations and differences.
2//!
3//! Nulls are skipped, like the NumPy `nan*` functions. A group without
4//! valid values gives `0` for sums, `1` for products, `false` for `any`,
5//! `true` for `all`, `0` for counts and null for all other reductions.
6//! Floats order NaN above all numbers, like Arrow, so `max` can return NaN.
7//! Use [`nan_to_null`](crate::compute::nan_to_null) to treat NaN as missing.
8
9use std::sync::Arc;
10
11use arrow::array::{
12    Array, ArrayRef, ArrowPrimitiveType, AsArray, BooleanArray, Int64Array, PrimitiveArray,
13    downcast_primitive_array,
14};
15use arrow::compute::cast as arrow_cast;
16use arrow::datatypes::{
17    ArrowNativeTypeOp, DataType, Float32Type, Float64Type, Int64Type, UInt64Type,
18};
19
20use crate::NdArrowArray;
21use crate::compute::elementwise::{ne, sub};
22use crate::compute::shape::{Selector, Slice, isel, transpose};
23use crate::dimensions::Dimensions;
24use crate::error::{Result, nd_err};
25
26/// A reduction for [`reduce`] and [`groupby_reduce`](crate::compute::groupby_reduce).
27#[derive(Debug, Clone, Copy, PartialEq)]
28pub enum Reduction {
29    /// Sum of the values.
30    Sum,
31    /// Product of the values.
32    Prod,
33    /// Arithmetic mean.
34    Mean,
35    /// Smallest value.
36    Min,
37    /// Largest value.
38    Max,
39    /// Number of non-null values.
40    Count,
41    /// Variance with `ddof` delta degrees of freedom, like `numpy.var`.
42    Var {
43        /// Delta degrees of freedom: the divisor is `n - ddof`.
44        ddof: usize,
45    },
46    /// Standard deviation with `ddof` delta degrees of freedom, like `numpy.std`.
47    Std {
48        /// Delta degrees of freedom: the divisor is `n - ddof`.
49        ddof: usize,
50    },
51    /// Middle value, with linear interpolation.
52    Median,
53    /// Quantile `q` in `[0, 1]` with linear interpolation, like `numpy.quantile`.
54    Quantile(f64),
55    /// `true` when any value is `true` or non-zero.
56    Any,
57    /// `true` when all values are `true` or non-zero.
58    All,
59    /// Index of the first minimum along the reduced axis (or the flat index for all axes).
60    ArgMin,
61    /// Index of the first maximum along the reduced axis (or the flat index for all axes).
62    ArgMax,
63}
64
65/// The values regrouped as `(outer, reduced, inner)` in C order.
66struct Blocks {
67    values: ArrayRef,
68    outer: usize,
69    reduced: usize,
70    inner: usize,
71    out_dims: Dimensions,
72}
73
74impl Blocks {
75    fn new(a: &NdArrowArray, axes: Option<&[&str]>) -> Result<Self> {
76        let dims = a.dims();
77        let mut selected: Vec<usize> = match axes {
78            None => (0..dims.rank()).collect(),
79            Some(names) => names
80                .iter()
81                .map(|name| dims.axis(name))
82                .collect::<Result<_>>()?,
83        };
84        selected.sort_unstable();
85        selected.dedup();
86        let sizes = dims.shape();
87        let product = |range: &[usize]| range.iter().map(|&i| sizes[i]).product::<usize>();
88        let kept: Vec<usize> = (0..dims.rank()).filter(|i| !selected.contains(i)).collect();
89        let out_dims = Dimensions::try_new(kept.iter().map(|&i| dims.as_slice()[i].clone()))?;
90
91        let Some((&first, &last)) = selected.first().zip(selected.last()) else {
92            return Ok(Self {
93                values: a.values().clone(),
94                outer: a.len(),
95                reduced: 1,
96                inner: 1,
97                out_dims,
98            });
99        };
100        if last - first + 1 == selected.len() {
101            return Ok(Self {
102                values: a.values().clone(),
103                outer: sizes[..first].iter().product(),
104                reduced: product(&selected),
105                inner: sizes[last + 1..].iter().product(),
106                out_dims,
107            });
108        }
109        let order: Vec<&str> = kept
110            .iter()
111            .chain(&selected)
112            .map(|&i| dims.as_slice()[i].name())
113            .collect();
114        Ok(Self {
115            values: transpose(a, &order)?.values().clone(),
116            outer: product(&kept),
117            reduced: product(&selected),
118            inner: 1,
119            out_dims,
120        })
121    }
122
123    fn groups(&self) -> usize {
124        self.outer * self.inner
125    }
126
127    /// Visit every valid position as `(group, index in group, position)`.
128    fn for_each(
129        &self,
130        nulls: Option<&arrow::buffer::NullBuffer>,
131        mut f: impl FnMut(usize, usize, usize),
132    ) {
133        let nulls = nulls.filter(|n| n.null_count() > 0);
134        for o in 0..self.outer {
135            for ri in 0..self.reduced {
136                let base = (o * self.reduced + ri) * self.inner;
137                let group = o * self.inner;
138                for j in 0..self.inner {
139                    let p = base + j;
140                    if nulls.is_some_and(|n| n.is_null(p)) {
141                        continue;
142                    }
143                    f(group + j, ri, p);
144                }
145            }
146        }
147    }
148
149    /// Each group as one contiguous slice, when groups are contiguous and there are no nulls.
150    fn contiguous_groups<'a, N>(
151        &self,
152        raw: &'a [N],
153        null_count: usize,
154    ) -> Option<std::slice::Chunks<'a, N>> {
155        (self.inner == 1 && self.reduced > 0 && null_count == 0).then(|| raw.chunks(self.reduced))
156    }
157
158    fn finish(&self, values: ArrayRef) -> Result<NdArrowArray> {
159        NdArrowArray::try_new(values, self.out_dims.clone())
160    }
161}
162
163/// Reduce `a` over the named axes. `None` reduces all axes to a scalar.
164pub fn reduce(a: &NdArrowArray, axes: Option<&[&str]>, op: Reduction) -> Result<NdArrowArray> {
165    let blocks = Blocks::new(a, axes)?;
166    let values = match op {
167        Reduction::Sum => sum_prod(&blocks, false)?,
168        Reduction::Prod => sum_prod(&blocks, true)?,
169        Reduction::Count => count_values(&blocks),
170        Reduction::Min => min_max(&blocks, false)?,
171        Reduction::Max => min_max(&blocks, true)?,
172        Reduction::ArgMin | Reduction::ArgMax => {
173            if axes.is_some_and(|names| names.len() != 1) {
174                return Err(nd_err!(
175                    InvalidArgument,
176                    "argmin and argmax reduce one axis or all axes"
177                ));
178            }
179            arg_min_max(&blocks, op == Reduction::ArgMax)?
180        }
181        Reduction::Mean => float_stat(&blocks, a.data_type(), Stat::Mean)?,
182        Reduction::Var { ddof } => float_stat(&blocks, a.data_type(), Stat::Var(ddof, false))?,
183        Reduction::Std { ddof } => float_stat(&blocks, a.data_type(), Stat::Var(ddof, true))?,
184        Reduction::Median => float_stat(&blocks, a.data_type(), Stat::Quantile(0.5))?,
185        Reduction::Quantile(q) => {
186            if !(0.0..=1.0).contains(&q) {
187                return Err(nd_err!(InvalidArgument, "quantile {q} is outside [0, 1]"));
188            }
189            float_stat(&blocks, a.data_type(), Stat::Quantile(q))?
190        }
191        Reduction::Any => any_all(&blocks, a, true)?,
192        Reduction::All => any_all(&blocks, a, false)?,
193    };
194    blocks.finish(values)
195}
196
197macro_rules! reductions {
198    ($($(#[$doc:meta])* $name:ident => $op:expr),* $(,)?) => {$(
199        $(#[$doc])*
200        pub fn $name(a: &NdArrowArray, axes: Option<&[&str]>) -> Result<NdArrowArray> {
201            reduce(a, axes, $op)
202        }
203    )*};
204}
205
206reductions!(
207    /// Sum, like `numpy.nansum`. Integers sum as `Int64` or `UInt64`, with wrapping.
208    sum => Reduction::Sum,
209    /// Product, like `numpy.nanprod`. Integers multiply as `Int64` or `UInt64`, with wrapping.
210    prod => Reduction::Prod,
211    /// Mean, like `numpy.nanmean`. Integers give `Float64`.
212    mean => Reduction::Mean,
213    /// Minimum, like `numpy.nanmin`.
214    min => Reduction::Min,
215    /// Maximum, like `numpy.nanmax`.
216    max => Reduction::Max,
217    /// Number of non-null values, like `xarray.DataArray.count`.
218    count => Reduction::Count,
219    /// Median, like `numpy.nanmedian`.
220    median => Reduction::Median,
221    /// `true` if any value is `true` (or non-zero), like `numpy.any`.
222    any => Reduction::Any,
223    /// `true` if all values are `true` (or non-zero), like `numpy.all`.
224    all => Reduction::All,
225    /// Index of the first minimum, like `numpy.nanargmin`.
226    argmin => Reduction::ArgMin,
227    /// Index of the first maximum, like `numpy.nanargmax`.
228    argmax => Reduction::ArgMax,
229);
230
231/// Variance, like `numpy.nanvar`.
232pub fn var(a: &NdArrowArray, axes: Option<&[&str]>, ddof: usize) -> Result<NdArrowArray> {
233    reduce(a, axes, Reduction::Var { ddof })
234}
235
236/// Standard deviation, like `numpy.nanstd`.
237pub fn std(a: &NdArrowArray, axes: Option<&[&str]>, ddof: usize) -> Result<NdArrowArray> {
238    reduce(a, axes, Reduction::Std { ddof })
239}
240
241/// Quantile with linear interpolation, like `numpy.nanquantile`.
242pub fn quantile(a: &NdArrowArray, axes: Option<&[&str]>, q: f64) -> Result<NdArrowArray> {
243    reduce(a, axes, Reduction::Quantile(q))
244}
245
246/// The accumulator type of sums and products, like NumPy.
247fn accumulator_type(data_type: &DataType) -> Result<DataType> {
248    Ok(match data_type {
249        DataType::Boolean
250        | DataType::Int8
251        | DataType::Int16
252        | DataType::Int32
253        | DataType::Int64 => DataType::Int64,
254        DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => {
255            DataType::UInt64
256        }
257        DataType::Float16 | DataType::Float32 => DataType::Float32,
258        DataType::Float64 => DataType::Float64,
259        other => return Err(nd_err!(NotSupported, "sum or product of {other}")),
260    })
261}
262
263fn sum_prod(blocks: &Blocks, product: bool) -> Result<ArrayRef> {
264    let to = accumulator_type(blocks.values.data_type())?;
265    let values = arrow_cast(blocks.values.as_ref(), &to)?;
266    Ok(match to {
267        DataType::Int64 => fold_sum_prod::<Int64Type>(blocks, &values, product),
268        DataType::UInt64 => fold_sum_prod::<UInt64Type>(blocks, &values, product),
269        DataType::Float32 => fold_sum_prod::<Float32Type>(blocks, &values, product),
270        _ => fold_sum_prod::<Float64Type>(blocks, &values, product),
271    })
272}
273
274fn fold_sum_prod<T: ArrowPrimitiveType>(
275    blocks: &Blocks,
276    values: &ArrayRef,
277    product: bool,
278) -> ArrayRef {
279    let array = values.as_primitive::<T>();
280    let raw = array.values();
281    let init = if product {
282        T::Native::ONE
283    } else {
284        T::Native::ZERO
285    };
286    let op = |a: T::Native, b: T::Native| {
287        if product {
288            a.mul_wrapping(b)
289        } else {
290            a.add_wrapping(b)
291        }
292    };
293    if blocks.contiguous_groups(raw, array.null_count()).is_some() {
294        let float_sum = !product && matches!(T::DATA_TYPE, DataType::Float32 | DataType::Float64);
295        let fold = |chunk: &[T::Native]| {
296            if float_sum {
297                pairwise_sum(chunk, init, &op)
298            } else {
299                lane_fold(chunk, init, op)
300            }
301        };
302        #[cfg(feature = "parallel")]
303        if raw.len() >= crate::compute::PARALLEL_MIN_LEN {
304            use rayon::prelude::*;
305            let accs: Vec<T::Native> = raw.par_chunks(blocks.reduced).map(fold).collect();
306            return Arc::new(PrimitiveArray::<T>::from_iter_values(accs));
307        }
308        let accs = raw.chunks(blocks.reduced).map(fold);
309        return Arc::new(PrimitiveArray::<T>::from_iter_values(accs));
310    }
311    let mut accs = vec![init; blocks.groups()];
312    blocks.for_each(array.nulls(), |g, _, p| accs[g] = op(accs[g], raw[p]));
313    Arc::new(PrimitiveArray::<T>::from_iter_values(accs))
314}
315
316/// NumPy's pairwise summation (`pairwise_sum` in `loops_utils.h`), so float
317/// sums over a contiguous axis round exactly like NumPy.
318fn pairwise_sum<N: Copy>(a: &[N], zero: N, add: &impl Fn(N, N) -> N) -> N {
319    const BLOCK: usize = 128;
320    let n = a.len();
321    if n < 8 {
322        a.iter().fold(zero, |acc, &v| add(acc, v))
323    } else if n <= BLOCK {
324        let mut r = [a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7]];
325        let full = n - n % 8;
326        for block in a[8..full].chunks_exact(8) {
327            for k in 0..8 {
328                r[k] = add(r[k], block[k]);
329            }
330        }
331        let mut res = add(
332            add(add(r[0], r[1]), add(r[2], r[3])),
333            add(add(r[4], r[5]), add(r[6], r[7])),
334        );
335        for &v in &a[full..] {
336            res = add(res, v);
337        }
338        res
339    } else {
340        let mut half = n / 2;
341        half -= half % 8;
342        add(
343            pairwise_sum(&a[..half], zero, add),
344            pairwise_sum(&a[half..], zero, add),
345        )
346    }
347}
348
349/// Fold with eight independent lanes, so that the compiler can vectorize.
350fn lane_fold<N: Copy>(chunk: &[N], init: N, op: impl Fn(N, N) -> N) -> N {
351    let mut lanes = [init; 8];
352    let mut exact = chunk.chunks_exact(8);
353    for block in &mut exact {
354        for k in 0..8 {
355            lanes[k] = op(lanes[k], block[k]);
356        }
357    }
358    let tail = exact.remainder().iter().fold(init, |acc, &v| op(acc, v));
359    lanes.into_iter().fold(tail, &op)
360}
361
362fn count_values(blocks: &Blocks) -> ArrayRef {
363    let mut counts = vec![0i64; blocks.groups()];
364    blocks.for_each(blocks.values.logical_nulls().as_ref(), |g, _, _| {
365        counts[g] += 1
366    });
367    Arc::new(Int64Array::from(counts))
368}
369
370fn min_max(blocks: &Blocks, is_max: bool) -> Result<ArrayRef> {
371    let values = &blocks.values;
372    downcast_primitive_array!(
373        values => Ok(fold_min_max(blocks, values, is_max)),
374        DataType::Boolean => {
375            let as_int = arrow_cast(values.as_ref(), &DataType::UInt8)?;
376            let out = fold_min_max(blocks, as_int.as_primitive::<arrow::datatypes::UInt8Type>(), is_max);
377            Ok(arrow_cast(out.as_ref(), &DataType::Boolean)?)
378        }
379        other => Err(nd_err!(NotSupported, "min or max of {other}"))
380    )
381}
382
383fn fold_min_max<T: ArrowPrimitiveType>(
384    blocks: &Blocks,
385    array: &PrimitiveArray<T>,
386    is_max: bool,
387) -> ArrayRef {
388    let raw = array.values();
389    let mut best: Vec<Option<T::Native>> = vec![None; blocks.groups()];
390    blocks.for_each(array.nulls(), |g, _, p| {
391        let v = raw[p];
392        let better = match best[g] {
393            None => true,
394            Some(b) if is_max => v.is_gt(b),
395            Some(b) => v.is_lt(b),
396        };
397        if better {
398            best[g] = Some(v);
399        }
400    });
401    Arc::new(PrimitiveArray::<T>::from_iter(best).with_data_type(array.data_type().clone()))
402}
403
404fn arg_min_max(blocks: &Blocks, is_max: bool) -> Result<ArrayRef> {
405    let values = &blocks.values;
406    downcast_primitive_array!(
407        values => Ok(fold_arg(blocks, values, is_max)),
408        other => Err(nd_err!(NotSupported, "argmin or argmax of {other}"))
409    )
410}
411
412fn fold_arg<T: ArrowPrimitiveType>(
413    blocks: &Blocks,
414    array: &PrimitiveArray<T>,
415    is_max: bool,
416) -> ArrayRef {
417    let raw = array.values();
418    let mut best: Vec<Option<(T::Native, usize)>> = vec![None; blocks.groups()];
419    blocks.for_each(array.nulls(), |g, ri, p| {
420        let v = raw[p];
421        let better = match best[g] {
422            None => true,
423            Some((b, _)) if is_max => v.is_gt(b),
424            Some((b, _)) => v.is_lt(b),
425        };
426        if better {
427            best[g] = Some((v, ri));
428        }
429    });
430    Arc::new(Int64Array::from_iter(
431        best.into_iter().map(|b| b.map(|(_, i)| i as i64)),
432    ))
433}
434
435enum Stat {
436    Mean,
437    Var(usize, bool),
438    Quantile(f64),
439}
440
441/// Statistics computed in `Float64`. The result is `Float32` for `Float32` input.
442fn float_stat(blocks: &Blocks, input: &DataType, stat: Stat) -> Result<ArrayRef> {
443    if !(input.is_numeric() || input == &DataType::Boolean) {
444        return Err(nd_err!(NotSupported, "statistics of {input}"));
445    }
446    let values = arrow_cast(blocks.values.as_ref(), &DataType::Float64)?;
447    let array = values.as_primitive::<Float64Type>();
448    let raw = array.values();
449    let groups = blocks.groups();
450    let out: Vec<Option<f64>> = match stat {
451        Stat::Mean if blocks.contiguous_groups(raw, array.null_count()).is_some() => blocks
452            .contiguous_groups(raw, 0)
453            .unwrap()
454            .map(|chunk| {
455                (!chunk.is_empty())
456                    .then(|| pairwise_sum(chunk, 0.0, &|a, b| a + b) / chunk.len() as f64)
457            })
458            .collect(),
459        Stat::Mean => {
460            let mut acc = vec![(0.0f64, 0usize); groups];
461            blocks.for_each(array.nulls(), |g, _, p| {
462                acc[g].0 += raw[p];
463                acc[g].1 += 1;
464            });
465            acc.into_iter()
466                .map(|(s, n)| (n > 0).then(|| s / n as f64))
467                .collect()
468        }
469        Stat::Var(ddof, sqrt) => {
470            // Welford update: count, mean, sum of squared deviations.
471            let mut acc = vec![(0usize, 0.0f64, 0.0f64); groups];
472            blocks.for_each(array.nulls(), |g, _, p| {
473                let (n, mean, m2) = &mut acc[g];
474                *n += 1;
475                let delta = raw[p] - *mean;
476                *mean += delta / *n as f64;
477                *m2 += delta * (raw[p] - *mean);
478            });
479            acc.into_iter()
480                .map(|(n, _, m2)| {
481                    (n > ddof).then(|| {
482                        let v = m2 / (n - ddof) as f64;
483                        if sqrt { v.sqrt() } else { v }
484                    })
485                })
486                .collect()
487        }
488        Stat::Quantile(q) => {
489            let mut groups_values: Vec<Vec<f64>> = vec![Vec::new(); groups];
490            blocks.for_each(array.nulls(), |g, _, p| groups_values[g].push(raw[p]));
491            groups_values
492                .into_iter()
493                .map(|mut v| linear_quantile(&mut v, q))
494                .collect()
495        }
496    };
497    let out: ArrayRef = Arc::new(PrimitiveArray::<Float64Type>::from_iter(out));
498    Ok(match input {
499        DataType::Float16 | DataType::Float32 => arrow_cast(out.as_ref(), &DataType::Float32)?,
500        _ => out,
501    })
502}
503
504/// NumPy's default `linear` quantile method.
505pub(crate) fn linear_quantile(values: &mut [f64], q: f64) -> Option<f64> {
506    if values.is_empty() {
507        return None;
508    }
509    values.sort_unstable_by(f64::total_cmp);
510    let position = q * (values.len() - 1) as f64;
511    let low = position.floor() as usize;
512    let high = position.ceil() as usize;
513    Some(values[low] + (values[high] - values[low]) * (position - low as f64))
514}
515
516fn any_all(blocks: &Blocks, a: &NdArrowArray, is_any: bool) -> Result<ArrayRef> {
517    let truthy = if a.data_type() == &DataType::Boolean {
518        blocks.values.clone()
519    } else {
520        let zero = NdArrowArray::try_new_scalar(arrow_cast(
521            &arrow::array::Int8Array::from(vec![0]),
522            a.data_type(),
523        )?)?;
524        let values = NdArrowArray::try_new(
525            blocks.values.clone(),
526            Dimensions::try_new([("n", blocks.values.len())])?,
527        )?;
528        ne(&values, &zero)?.values().clone()
529    };
530    let array = truthy.as_boolean();
531    let mut acc = vec![!is_any; blocks.groups()];
532    blocks.for_each(array.nulls(), |g, _, p| {
533        if array.value(p) == is_any {
534            acc[g] = is_any;
535        }
536    });
537    Ok(Arc::new(BooleanArray::from(acc)))
538}
539
540/// Cumulative sum along `dim`, like `numpy.nancumsum`. Null positions stay null.
541pub fn cumsum(a: &NdArrowArray, dim: &str) -> Result<NdArrowArray> {
542    cumulative(a, dim, false)
543}
544
545/// Cumulative product along `dim`, like `numpy.nancumprod`. Null positions stay null.
546pub fn cumprod(a: &NdArrowArray, dim: &str) -> Result<NdArrowArray> {
547    cumulative(a, dim, true)
548}
549
550fn cumulative(a: &NdArrowArray, dim: &str, product: bool) -> Result<NdArrowArray> {
551    let blocks = Blocks::new(a, Some(&[dim]))?;
552    let to = accumulator_type(a.data_type())?;
553    let values = arrow_cast(a.values().as_ref(), &to)?;
554    let out = match to {
555        DataType::Int64 => scan::<Int64Type>(&blocks, &values, product),
556        DataType::UInt64 => scan::<UInt64Type>(&blocks, &values, product),
557        DataType::Float32 => scan::<Float32Type>(&blocks, &values, product),
558        _ => scan::<Float64Type>(&blocks, &values, product),
559    };
560    NdArrowArray::try_new(out, a.dims().clone())
561}
562
563fn scan<T: ArrowPrimitiveType>(blocks: &Blocks, values: &ArrayRef, product: bool) -> ArrayRef {
564    let array = values.as_primitive::<T>();
565    let raw = array.values();
566    let init = if product {
567        T::Native::ONE
568    } else {
569        T::Native::ZERO
570    };
571    let mut accs = vec![init; blocks.groups()];
572    let mut out = vec![init; raw.len()];
573    blocks.for_each(array.nulls(), |g, _, p| {
574        accs[g] = if product {
575            accs[g].mul_wrapping(raw[p])
576        } else {
577            accs[g].add_wrapping(raw[p])
578        };
579        out[p] = accs[g];
580    });
581    Arc::new(PrimitiveArray::<T>::new(out.into(), array.nulls().cloned()))
582}
583
584/// Running minimum along `dim`, like `numpy.minimum.accumulate`. Null positions stay null.
585pub fn cummin(a: &NdArrowArray, dim: &str) -> Result<NdArrowArray> {
586    cumulative_extreme(a, dim, false)
587}
588
589/// Running maximum along `dim`, like `numpy.maximum.accumulate`. Null positions stay null.
590pub fn cummax(a: &NdArrowArray, dim: &str) -> Result<NdArrowArray> {
591    cumulative_extreme(a, dim, true)
592}
593
594fn cumulative_extreme(a: &NdArrowArray, dim: &str, is_max: bool) -> Result<NdArrowArray> {
595    let blocks = Blocks::new(a, Some(&[dim]))?;
596    let values = a.values();
597    let out = downcast_primitive_array!(
598        values => scan_extreme(&blocks, values, is_max),
599        other => return Err(nd_err!(NotSupported, "cummin or cummax of {other}"))
600    );
601    NdArrowArray::try_new(out, a.dims().clone())
602}
603
604fn scan_extreme<T: ArrowPrimitiveType>(
605    blocks: &Blocks,
606    array: &PrimitiveArray<T>,
607    is_max: bool,
608) -> ArrayRef {
609    let raw = array.values();
610    let mut best: Vec<Option<T::Native>> = vec![None; blocks.groups()];
611    let mut out = vec![T::Native::default(); raw.len()];
612    blocks.for_each(array.nulls(), |g, _, p| {
613        let v = raw[p];
614        let next = match best[g] {
615            Some(b) if (is_max && !v.is_gt(b)) || (!is_max && !v.is_lt(b)) => b,
616            _ => v,
617        };
618        best[g] = Some(next);
619        out[p] = next;
620    });
621    Arc::new(
622        PrimitiveArray::<T>::new(out.into(), array.nulls().cloned())
623            .with_data_type(array.data_type().clone()),
624    )
625}
626
627/// Range of values (`max - min`), like `numpy.ptp`.
628pub fn ptp(a: &NdArrowArray, axes: Option<&[&str]>) -> Result<NdArrowArray> {
629    sub(&max(a, axes)?, &min(a, axes)?)
630}
631
632/// The `n`-th discrete difference along `dim`, like `numpy.diff`.
633pub fn diff(a: &NdArrowArray, dim: &str, n: usize) -> Result<NdArrowArray> {
634    let mut out = a.clone();
635    for _ in 0..n {
636        let upper = isel(
637            &out,
638            &[(dim, Selector::Slice(Slice::new(Some(1), None, 1)))],
639        )?;
640        let lower = isel(
641            &out,
642            &[(dim, Selector::Slice(Slice::new(None, Some(-1), 1)))],
643        )?;
644        out = sub(&upper, &lower)?;
645    }
646    Ok(out)
647}
648
649#[cfg(test)]
650mod tests {
651    use arrow::array::{Float32Array, Float64Array, Int8Array, Int32Array};
652    use arrow::datatypes::{Float64Type, Int64Type};
653
654    use super::*;
655
656    fn nd(values: ArrayRef, dims: &[(&str, usize)]) -> NdArrowArray {
657        NdArrowArray::try_new(values, Dimensions::try_new(dims.iter().copied()).unwrap()).unwrap()
658    }
659
660    fn grid() -> NdArrowArray {
661        nd(
662            Arc::new(Int32Array::from_iter_values(0..24)),
663            &[("a", 2), ("b", 3), ("c", 4)],
664        )
665    }
666
667    fn i64s(a: &NdArrowArray) -> Vec<Option<i64>> {
668        a.values().as_primitive::<Int64Type>().iter().collect()
669    }
670
671    fn f64s(a: &NdArrowArray) -> Vec<Option<f64>> {
672        a.values().as_primitive::<Float64Type>().iter().collect()
673    }
674
675    #[test]
676    fn sum_over_each_axis_layout() {
677        let g = grid();
678        let all = sum(&g, None).unwrap();
679        assert!(all.is_scalar());
680        assert_eq!(i64s(&all), vec![Some(276)]);
681
682        let last = sum(&g, Some(&["c"])).unwrap();
683        assert_eq!(last.dims().names(), vec!["a", "b"]);
684        assert_eq!(i64s(&last), [6, 22, 38, 54, 70, 86].map(Some).to_vec());
685
686        let first = sum(&g, Some(&["a"])).unwrap();
687        assert_eq!(first.shape(), vec![3, 4]);
688        assert_eq!(i64s(&first)[0], Some(12));
689
690        let outer = sum(&g, Some(&["a", "c"])).unwrap();
691        assert_eq!(outer.dims().names(), vec!["b"]);
692        assert_eq!(i64s(&outer), vec![Some(60), Some(92), Some(124)]);
693    }
694
695    #[test]
696    fn nulls_are_skipped() {
697        let a = nd(
698            Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0), None])),
699            &[("x", 2), ("y", 2)],
700        );
701        assert_eq!(
702            f64s(&sum(&a, Some(&["y"])).unwrap()),
703            vec![Some(1.0), Some(3.0)]
704        );
705        assert_eq!(
706            f64s(&mean(&a, Some(&["x"])).unwrap()),
707            vec![Some(2.0), None]
708        );
709        assert_eq!(
710            i64s(&count(&a, Some(&["x"])).unwrap()),
711            vec![Some(2), Some(0)]
712        );
713        assert_eq!(
714            f64s(&sum(&a, Some(&["x"])).unwrap()),
715            vec![Some(4.0), Some(0.0)],
716            "an all-null group sums to zero"
717        );
718    }
719
720    #[test]
721    fn min_max_and_arg() {
722        let a = nd(
723            Arc::new(Int32Array::from(vec![3, 1, 4, 1, 5, 9])),
724            &[("x", 2), ("y", 3)],
725        );
726        let m = max(&a, Some(&["y"])).unwrap();
727        assert_eq!(
728            m.values()
729                .as_primitive::<arrow::datatypes::Int32Type>()
730                .values(),
731            &[4, 9]
732        );
733        assert_eq!(
734            i64s(&argmin(&a, Some(&["y"])).unwrap()),
735            vec![Some(1), Some(0)]
736        );
737        assert_eq!(i64s(&argmax(&a, None).unwrap()), vec![Some(5)]);
738        assert!(argmax(&grid(), Some(&["a", "b"])).is_err());
739    }
740
741    #[test]
742    fn var_std_quantile() {
743        let a = nd(
744            Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])),
745            &[("x", 4)],
746        );
747        assert_eq!(f64s(&var(&a, None, 0).unwrap()), vec![Some(1.25)]);
748        assert_eq!(f64s(&var(&a, None, 1).unwrap())[0].unwrap(), 5.0 / 3.0);
749        assert_eq!(f64s(&std(&a, None, 0).unwrap()), vec![Some(1.25f64.sqrt())]);
750        assert_eq!(f64s(&median(&a, None).unwrap()), vec![Some(2.5)]);
751        assert_eq!(f64s(&quantile(&a, None, 0.25).unwrap()), vec![Some(1.75)]);
752        assert!(quantile(&a, None, 1.5).is_err());
753
754        let f = nd(Arc::new(Float32Array::from(vec![1.0, 2.0])), &[("x", 2)]);
755        assert_eq!(mean(&f, None).unwrap().data_type(), &DataType::Float32);
756    }
757
758    #[test]
759    fn any_all_on_numbers_and_bools() {
760        let a = nd(
761            Arc::new(Int8Array::from(vec![0, 0, 2, 0])),
762            &[("x", 2), ("y", 2)],
763        );
764        let any_y: Vec<bool> = any(&a, Some(&["y"]))
765            .unwrap()
766            .values()
767            .as_boolean()
768            .iter()
769            .map(Option::unwrap)
770            .collect();
771        assert_eq!(any_y, vec![false, true]);
772        assert!(!all(&a, None).unwrap().values().as_boolean().value(0));
773    }
774
775    #[test]
776    fn sum_widens_like_numpy() {
777        let a = nd(Arc::new(Int8Array::from(vec![100, 100])), &[("x", 2)]);
778        assert_eq!(i64s(&sum(&a, None).unwrap()), vec![Some(200)]);
779    }
780
781    #[test]
782    fn cumsum_and_diff() {
783        let a = nd(
784            Arc::new(Int32Array::from(vec![Some(1), None, Some(3), Some(4)])),
785            &[("x", 2), ("y", 2)],
786        );
787        assert_eq!(
788            i64s(&cumsum(&a, "x").unwrap()),
789            vec![Some(1), None, Some(4), Some(4)]
790        );
791        assert_eq!(
792            i64s(&cumprod(&a, "y").unwrap()),
793            vec![Some(1), None, Some(3), Some(12)]
794        );
795
796        let b = nd(Arc::new(Int32Array::from(vec![1, 4, 9, 16])), &[("x", 4)]);
797        let d = diff(&b, "x", 2).unwrap();
798        assert_eq!(d.shape(), vec![2]);
799        assert_eq!(
800            d.values()
801                .as_primitive::<arrow::datatypes::Int32Type>()
802                .values(),
803            &[2, 2]
804        );
805    }
806}
nd_arrow_array/compute/reduce.rs

nd_arrow_array/compute/
reduce.rs
