nd_arrow_array/compute/
methods.rs

1//! Method forms of common kernels, like NumPy array methods: `a.sum(..)`.
2
3use arrow::datatypes::DataType;
4
5use crate::NdArrowArray;
6use crate::compute::{self, Selector};
7use crate::dimensions::Dimensions;
8use crate::error::Result;
9
10impl NdArrowArray {
11    /// See [`compute::broadcast_to`].
12    pub fn broadcast_to(&self, target: &Dimensions) -> Result<Self> {
13        compute::broadcast_to(self, target)
14    }
15
16    /// See [`compute::transpose`].
17    pub fn transpose(&self, order: &[&str]) -> Result<Self> {
18        compute::transpose(self, order)
19    }
20
21    /// See [`compute::squeeze`].
22    pub fn squeeze(&self, dims: Option<&[&str]>) -> Result<Self> {
23        compute::squeeze(self, dims)
24    }
25
26    /// See [`compute::expand_dims`].
27    pub fn expand_dims(&self, name: &str, axis: usize) -> Result<Self> {
28        compute::expand_dims(self, name, axis)
29    }
30
31    /// See [`compute::isel`].
32    pub fn isel(&self, selections: &[(&str, Selector)]) -> Result<Self> {
33        compute::isel(self, selections)
34    }
35
36    /// See [`compute::cast`].
37    pub fn cast(&self, to: &DataType) -> Result<Self> {
38        compute::cast(self, to)
39    }
40
41    /// See [`compute::fill_null`].
42    pub fn fill_null(&self, value: &arrow::array::ArrayRef) -> Result<Self> {
43        compute::fill_null(self, value)
44    }
45
46    /// See [`compute::sum`].
47    pub fn sum(&self, axes: Option<&[&str]>) -> Result<Self> {
48        compute::sum(self, axes)
49    }
50
51    /// See [`compute::mean`].
52    pub fn mean(&self, axes: Option<&[&str]>) -> Result<Self> {
53        compute::mean(self, axes)
54    }
55
56    /// See [`compute::min`].
57    pub fn min(&self, axes: Option<&[&str]>) -> Result<Self> {
58        compute::min(self, axes)
59    }
60
61    /// See [`compute::max`].
62    pub fn max(&self, axes: Option<&[&str]>) -> Result<Self> {
63        compute::max(self, axes)
64    }
65
66    /// See [`compute::count`].
67    pub fn count(&self, axes: Option<&[&str]>) -> Result<Self> {
68        compute::count(self, axes)
69    }
70
71    /// See [`compute::std`].
72    pub fn std(&self, axes: Option<&[&str]>, ddof: usize) -> Result<Self> {
73        compute::std(self, axes, ddof)
74    }
75
76    /// See [`compute::var`].
77    pub fn var(&self, axes: Option<&[&str]>, ddof: usize) -> Result<Self> {
78        compute::var(self, axes, ddof)
79    }
80
81    /// See [`compute::median`].
82    pub fn median(&self, axes: Option<&[&str]>) -> Result<Self> {
83        compute::median(self, axes)
84    }
85
86    /// See [`compute::resample()`].
87    pub fn resample(
88        &self,
89        dim: &str,
90        times: &NdArrowArray,
91        freq: compute::Frequency,
92        op: compute::Reduction,
93    ) -> Result<(Self, Self)> {
94        compute::resample(self, dim, times, freq, op)
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use std::sync::Arc;
101
102    use arrow::array::{AsArray, Int32Array};
103    use arrow::datatypes::Int64Type;
104
105    use super::*;
106
107    #[test]
108    fn methods_chain() {
109        let a = NdArrowArray::try_new(
110            Arc::new(Int32Array::from_iter_values(0..6)),
111            Dimensions::try_new([("x", 2), ("y", 3)]).unwrap(),
112        )
113        .unwrap();
114        let s = a.transpose(&["y", "x"]).unwrap().sum(Some(&["x"])).unwrap();
115        assert_eq!(s.values().as_primitive::<Int64Type>().values(), &[3, 5, 7]);
116    }
117}
nd_arrow_array/compute/methods.rs

nd_arrow_array/compute/
methods.rs
