nd_arrow_array/compute/
histogram.rs

1//! Counting kernels: `histogram` and `bincount`.
2
3use std::sync::Arc;
4
5use arrow::array::{Array, AsArray, Float64Array, Int64Array};
6use arrow::compute::cast;
7use arrow::datatypes::{DataType, Float64Type, Int64Type};
8
9use crate::NdArrowArray;
10use crate::dimensions::Dimensions;
11use crate::error::{Result, nd_err};
12
13/// Count all elements of `a` per bin, like `numpy.histogram` with explicit edges.
14///
15/// `edges` is 1-D and increasing. Bins are half open, `[edges[i], edges[i+1])`,
16/// except the last bin, which includes its right edge. Nulls, NaN and values
17/// outside the edges are not counted. The result is `Int64` along `dim`.
18pub fn histogram(a: &NdArrowArray, edges: &NdArrowArray, dim: &str) -> Result<NdArrowArray> {
19    let edges = cast(edges.values().as_ref(), &DataType::Float64)?;
20    let edges = edges.as_primitive::<Float64Type>();
21    if edges.len() < 2 || edges.null_count() > 0 || edges.values().windows(2).any(|w| w[1] <= w[0])
22    {
23        return Err(nd_err!(
24            InvalidArgument,
25            "histogram edges must be at least two increasing, non-null values"
26        ));
27    }
28    let edges = edges.values();
29    let bins = edges.len() - 1;
30    let mut counts = vec![0i64; bins];
31    let values = cast(a.values().as_ref(), &DataType::Float64)?;
32    for v in values.as_primitive::<Float64Type>().iter().flatten() {
33        if v.is_nan() || v < edges[0] || v > edges[bins] {
34            continue;
35        }
36        let bin = edges
37            .partition_point(|&e| e <= v)
38            .saturating_sub(1)
39            .min(bins - 1);
40        counts[bin] += 1;
41    }
42    NdArrowArray::try_new(
43        Arc::new(Int64Array::from(counts)),
44        Dimensions::try_new([(dim, bins)])?,
45    )
46}
47
48/// Count each non-negative integer value, like `numpy.bincount`.
49///
50/// With `weights` (same dimensions as `a`), sums the weights instead and
51/// returns `Float64`. The result has at least `min_length` positions along
52/// `dim`. Null values are not counted. Negative values are an error.
53pub fn bincount(
54    a: &NdArrowArray,
55    weights: Option<&NdArrowArray>,
56    min_length: usize,
57    dim: &str,
58) -> Result<NdArrowArray> {
59    if !a.data_type().is_integer() {
60        return Err(nd_err!(
61            InvalidArgument,
62            "bincount needs integers, got {}",
63            a.data_type()
64        ));
65    }
66    let values = cast(a.values().as_ref(), &DataType::Int64)?;
67    let values = values.as_primitive::<Int64Type>();
68    if values.iter().flatten().any(|v| v < 0) {
69        return Err(nd_err!(
70            InvalidArgument,
71            "bincount values must be non-negative"
72        ));
73    }
74    let len = values
75        .iter()
76        .flatten()
77        .max()
78        .map_or(0, |m| m as usize + 1)
79        .max(min_length);
80    let dims = Dimensions::try_new([(dim, len)])?;
81    match weights {
82        None => {
83            let mut counts = vec![0i64; len];
84            for v in values.iter().flatten() {
85                counts[v as usize] += 1;
86            }
87            NdArrowArray::try_new(Arc::new(Int64Array::from(counts)), dims)
88        }
89        Some(w) => {
90            if w.dims() != a.dims() {
91                return Err(nd_err!(
92                    ShapeMismatch,
93                    "bincount weights {} do not match values {}",
94                    w.dims(),
95                    a.dims()
96                ));
97            }
98            let w = cast(w.values().as_ref(), &DataType::Float64)?;
99            let w = w.as_primitive::<Float64Type>();
100            let mut sums = vec![0.0f64; len];
101            for (v, weight) in values.iter().zip(w.iter()) {
102                if let (Some(v), Some(weight)) = (v, weight) {
103                    sums[v as usize] += weight;
104                }
105            }
106            NdArrowArray::try_new(Arc::new(Float64Array::from(sums)), dims)
107        }
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use arrow::array::Int32Array;
114
115    use super::*;
116
117    fn nd(values: Vec<f64>, dims: &[(&str, usize)]) -> NdArrowArray {
118        NdArrowArray::try_new(
119            Arc::new(Float64Array::from(values)),
120            Dimensions::try_new(dims.iter().copied()).unwrap(),
121        )
122        .unwrap()
123    }
124
125    #[test]
126    fn histogram_matches_numpy() {
127        let a = nd(
128            vec![0.5, 1.0, 1.5, 2.0, 2.0, 3.0, f64::NAN, -1.0],
129            &[("x", 8)],
130        );
131        let edges = nd(vec![0.0, 1.0, 2.0], &[("e", 3)]);
132        let h = histogram(&a, &edges, "bin").unwrap();
133        assert_eq!(h.values().as_primitive::<Int64Type>().values(), &[1, 4]);
134        assert!(histogram(&a, &nd(vec![1.0], &[("e", 1)]), "bin").is_err());
135    }
136
137    #[test]
138    fn bincount_with_and_without_weights() {
139        let a = NdArrowArray::try_new(
140            Arc::new(Int32Array::from(vec![Some(0), Some(2), Some(2), None])),
141            Dimensions::try_new([("x", 4)]).unwrap(),
142        )
143        .unwrap();
144        let c = bincount(&a, None, 5, "value").unwrap();
145        assert_eq!(
146            c.values().as_primitive::<Int64Type>().values(),
147            &[1, 0, 2, 0, 0]
148        );
149        let w = nd(vec![0.5, 1.0, 2.0, 9.0], &[("x", 4)]);
150        let s = bincount(&a, Some(&w), 0, "value").unwrap();
151        assert_eq!(
152            s.values().as_primitive::<Float64Type>().values(),
153            &[0.5, 0.0, 3.0]
154        );
155    }
156}
nd_arrow_array/compute/histogram.rs

nd_arrow_array/compute/
histogram.rs
