nd_arrow_array/compute/
partial.rs

1//! Reduction states that merge exactly across chunks.
2//!
3//! Reduce each chunk to a [`PartialReduction`], merge the partials in any
4//! order or tree shape, then [`finish`](PartialReduction::finish). The result
5//! equals one reduction over all chunks (up to float rounding). This is the
6//! base for chunked and lazy execution.
7
8use std::sync::Arc;
9
10use arrow::array::{Array, AsArray, Float64Array};
11use arrow::compute::cast;
12use arrow::compute::kernels::{boolean, cmp, zip::zip};
13use arrow::datatypes::{DataType, Float64Type, Int64Type};
14
15use crate::NdArrowArray;
16use crate::compute::elementwise::{add, logical_and, logical_or, mul};
17use crate::compute::reduce::{Reduction, count, reduce};
18use crate::dimensions::Dimensions;
19use crate::error::{Result, nd_err};
20
21#[derive(Debug, Clone)]
22enum State {
23    /// Sum, product, any or all: one array that merges with one kernel.
24    Single(NdArrowArray),
25    /// Minimum or maximum, with nulls for empty groups.
26    Extreme(NdArrowArray),
27    /// Sum and count, both `Float64`.
28    Mean { sum: Vec<f64>, count: Vec<u64> },
29    /// Count, mean and sum of squared deviations (Chan et al.).
30    Moments {
31        count: Vec<u64>,
32        mean: Vec<f64>,
33        m2: Vec<f64>,
34    },
35}
36
37/// The mergeable state of one reduction over one chunk.
38#[derive(Debug, Clone)]
39pub struct PartialReduction {
40    op: Reduction,
41    dims: Dimensions,
42    input_type: DataType,
43    state: State,
44}
45
46impl PartialReduction {
47    /// Reduce one chunk. Median, quantiles and arg reductions cannot merge.
48    pub fn new(a: &NdArrowArray, axes: Option<&[&str]>, op: Reduction) -> Result<Self> {
49        let f64_input = || crate::compute::cast(a, &DataType::Float64);
50        let (state, dims) = match op {
51            Reduction::Sum
52            | Reduction::Prod
53            | Reduction::Count
54            | Reduction::Any
55            | Reduction::All => {
56                let r = reduce(a, axes, op)?;
57                let dims = r.dims().clone();
58                (State::Single(r), dims)
59            }
60            Reduction::Min | Reduction::Max => {
61                let r = reduce(a, axes, op)?;
62                let dims = r.dims().clone();
63                (State::Extreme(r), dims)
64            }
65            Reduction::Mean => {
66                let x = f64_input()?;
67                let s = reduce(&x, axes, Reduction::Sum)?;
68                let n = count(&x, axes)?;
69                let dims = s.dims().clone();
70                (
71                    State::Mean {
72                        sum: f64_values(&s),
73                        count: u64_values(&n),
74                    },
75                    dims,
76                )
77            }
78            Reduction::Var { .. } | Reduction::Std { .. } => {
79                let x = f64_input()?;
80                let n = count(&x, axes)?;
81                let mean = reduce(&x, axes, Reduction::Mean)?;
82                let var = reduce(&x, axes, Reduction::Var { ddof: 0 })?;
83                let dims = n.dims().clone();
84                let count = u64_values(&n);
85                let m2 = f64_values(&var)
86                    .iter()
87                    .zip(&count)
88                    .map(|(v, &c)| if c == 0 { 0.0 } else { v * c as f64 })
89                    .collect();
90                (
91                    State::Moments {
92                        count,
93                        mean: f64_values(&mean),
94                        m2,
95                    },
96                    dims,
97                )
98            }
99            Reduction::Median | Reduction::Quantile(_) | Reduction::ArgMin | Reduction::ArgMax => {
100                return Err(nd_err!(NotSupported, "{op:?} has no mergeable state"));
101            }
102        };
103        Ok(Self {
104            op,
105            dims,
106            input_type: a.data_type().clone(),
107            state,
108        })
109    }
110
111    /// The dimensions of the result.
112    pub fn dims(&self) -> &Dimensions {
113        &self.dims
114    }
115
116    /// Combine with the partial of another chunk.
117    pub fn merge(&self, other: &Self) -> Result<Self> {
118        if self.op != other.op || self.dims != other.dims {
119            return Err(nd_err!(
120                InvalidArgument,
121                "cannot merge {:?} over {} with {:?} over {}",
122                self.op,
123                self.dims,
124                other.op,
125                other.dims
126            ));
127        }
128        let state = match (&self.state, &other.state) {
129            (State::Single(a), State::Single(b)) => State::Single(match self.op {
130                Reduction::Prod => mul(a, b)?,
131                Reduction::Any => logical_or(a, b)?,
132                Reduction::All => logical_and(a, b)?,
133                _ => add(a, b)?,
134            }),
135            (State::Extreme(a), State::Extreme(b)) => {
136                State::Extreme(merge_extreme(a, b, self.op == Reduction::Max)?)
137            }
138            (State::Mean { sum: sa, count: ca }, State::Mean { sum: sb, count: cb }) => {
139                State::Mean {
140                    sum: sa.iter().zip(sb).map(|(a, b)| a + b).collect(),
141                    count: ca.iter().zip(cb).map(|(a, b)| a + b).collect(),
142                }
143            }
144            (
145                State::Moments {
146                    count: na,
147                    mean: ma,
148                    m2: qa,
149                },
150                State::Moments {
151                    count: nb,
152                    mean: mb,
153                    m2: qb,
154                },
155            ) => {
156                let mut count = Vec::with_capacity(na.len());
157                let mut mean = Vec::with_capacity(na.len());
158                let mut m2 = Vec::with_capacity(na.len());
159                for i in 0..na.len() {
160                    let (n1, n2) = (na[i] as f64, nb[i] as f64);
161                    let n = n1 + n2;
162                    let (c, m, q) = if n2 == 0.0 {
163                        (na[i], ma[i], qa[i])
164                    } else if n1 == 0.0 {
165                        (nb[i], mb[i], qb[i])
166                    } else {
167                        let delta = mb[i] - ma[i];
168                        (
169                            na[i] + nb[i],
170                            ma[i] + delta * n2 / n,
171                            qa[i] + qb[i] + delta * delta * n1 * n2 / n,
172                        )
173                    };
174                    count.push(c);
175                    mean.push(m);
176                    m2.push(q);
177                }
178                State::Moments { count, mean, m2 }
179            }
180            _ => unreachable!("same op gives the same state"),
181        };
182        Ok(Self {
183            op: self.op,
184            dims: self.dims.clone(),
185            input_type: self.input_type.clone(),
186            state,
187        })
188    }
189
190    /// The reduction result.
191    pub fn finish(&self) -> Result<NdArrowArray> {
192        let float_out = |values: Vec<Option<f64>>| -> Result<NdArrowArray> {
193            let mut out: Arc<dyn Array> = Arc::new(Float64Array::from(values));
194            if matches!(self.input_type, DataType::Float32 | DataType::Float16) {
195                out = cast(out.as_ref(), &DataType::Float32)?;
196            }
197            NdArrowArray::try_new(out, self.dims.clone())
198        };
199        match &self.state {
200            State::Single(a) | State::Extreme(a) => Ok(a.clone()),
201            State::Mean { sum, count } => float_out(
202                sum.iter()
203                    .zip(count)
204                    .map(|(&s, &n)| (n > 0).then(|| s / n as f64))
205                    .collect(),
206            ),
207            State::Moments { count, m2, .. } => {
208                let (ddof, sqrt) = match self.op {
209                    Reduction::Var { ddof } => (ddof, false),
210                    Reduction::Std { ddof } => (ddof, true),
211                    _ => unreachable!(),
212                };
213                float_out(
214                    count
215                        .iter()
216                        .zip(m2)
217                        .map(|(&n, &q)| {
218                            (n as usize > ddof).then(|| {
219                                let v = q / (n as usize - ddof) as f64;
220                                if sqrt { v.sqrt() } else { v }
221                            })
222                        })
223                        .collect(),
224                )
225            }
226        }
227    }
228}
229
230/// Merge partials in a balanced tree, which keeps float error low.
231pub fn merge_all(partials: &[PartialReduction]) -> Result<PartialReduction> {
232    match partials {
233        [] => Err(nd_err!(InvalidArgument, "no partial reductions to merge")),
234        [single] => Ok(single.clone()),
235        many => {
236            let (left, right) = many.split_at(many.len() / 2);
237            merge_all(left)?.merge(&merge_all(right)?)
238        }
239    }
240}
241
242fn f64_values(a: &NdArrowArray) -> Vec<f64> {
243    let values = cast(a.values().as_ref(), &DataType::Float64).expect("numeric");
244    values
245        .as_primitive::<Float64Type>()
246        .iter()
247        .map(|v| v.unwrap_or(0.0))
248        .collect()
249}
250
251fn u64_values(a: &NdArrowArray) -> Vec<u64> {
252    a.values()
253        .as_primitive::<Int64Type>()
254        .values()
255        .iter()
256        .map(|&v| v as u64)
257        .collect()
258}
259
260/// Elementwise minimum or maximum that skips nulls.
261fn merge_extreme(a: &NdArrowArray, b: &NdArrowArray, is_max: bool) -> Result<NdArrowArray> {
262    let (x, y) = (a.values(), b.values());
263    let better = if is_max {
264        cmp::gt(y, x)?
265    } else {
266        cmp::lt(y, x)?
267    };
268    let a_missing = arrow::compute::is_null(x.as_ref())?;
269    let b_present = arrow::compute::is_not_null(y.as_ref())?;
270    let take_b = boolean::or_kleene(&a_missing, &boolean::and_kleene(&b_present, &better)?)?;
271    let values = zip(&take_b, y, x)?;
272    NdArrowArray::try_new(values, a.dims().clone())
273}
274
275#[cfg(test)]
276mod tests {
277    use arrow::array::Float64Array;
278
279    use super::*;
280    use crate::compute::shape::{Selector, Slice, isel};
281
282    fn data() -> NdArrowArray {
283        let values: Vec<Option<f64>> = (0..60)
284            .map(|i| (i % 7 != 3).then(|| ((i * 37) % 23) as f64 - 11.0))
285            .collect();
286        NdArrowArray::try_new(
287            Arc::new(Float64Array::from(values)),
288            Dimensions::try_new([("time", 12), ("x", 5)]).unwrap(),
289        )
290        .unwrap()
291    }
292
293    fn chunks(a: &NdArrowArray) -> Vec<NdArrowArray> {
294        [(0, 5), (5, 6), (6, 12)]
295            .iter()
296            .map(|&(s, e)| isel(a, &[("time", Selector::Slice(Slice::range(s, e)))]).unwrap())
297            .collect()
298    }
299
300    fn close(a: &NdArrowArray, b: &NdArrowArray) -> bool {
301        let x = cast(a.values().as_ref(), &DataType::Float64).unwrap();
302        let y = cast(b.values().as_ref(), &DataType::Float64).unwrap();
303        let (x, y) = (
304            x.as_primitive::<Float64Type>(),
305            y.as_primitive::<Float64Type>(),
306        );
307        x.len() == y.len()
308            && x.iter().zip(y.iter()).all(|(p, q)| match (p, q) {
309                (None, None) => true,
310                (Some(p), Some(q)) => (p - q).abs() < 1e-9,
311                _ => false,
312            })
313    }
314
315    #[test]
316    fn merged_chunks_equal_full_reduction() {
317        let a = data();
318        let ops = [
319            Reduction::Sum,
320            Reduction::Prod,
321            Reduction::Count,
322            Reduction::Mean,
323            Reduction::Min,
324            Reduction::Max,
325            Reduction::Var { ddof: 0 },
326            Reduction::Std { ddof: 1 },
327        ];
328        for op in ops {
329            for axes in [Some(&["time"][..]), None] {
330                let partials: Vec<PartialReduction> = chunks(&a)
331                    .iter()
332                    .map(|c| PartialReduction::new(c, axes, op).unwrap())
333                    .collect();
334                let merged = merge_all(&partials).unwrap().finish().unwrap();
335                let full = reduce(&a, axes, op).unwrap();
336                assert!(close(&merged, &full), "{op:?} over {axes:?}");
337            }
338        }
339    }
340
341    #[test]
342    fn unsupported_and_mismatched() {
343        let a = data();
344        assert!(PartialReduction::new(&a, None, Reduction::Median).is_err());
345        let p = PartialReduction::new(&a, None, Reduction::Sum).unwrap();
346        let q = PartialReduction::new(&a, Some(&["x"]), Reduction::Sum).unwrap();
347        assert!(p.merge(&q).is_err());
348    }
349}
nd_arrow_array/compute/partial.rs

nd_arrow_array/compute/
partial.rs
