runlength_arrow/
slices.rs

1//! Aggregating many slices of sequences without copying them.
2
3use std::sync::Arc;
4
5use arrow_array::{
6    Array, ArrayRef, ArrowPrimitiveType, Float64Array, Int64Array, PrimitiveArray, RecordBatch,
7};
8use arrow_buffer::ScalarBuffer;
9use arrow_cast::cast;
10use arrow_schema::DataType;
11use rayon::prelude::*;
12use runlength::Num;
13
14use crate::batch::ArrowRleBatch;
15use crate::error::Result;
16use crate::numeric::with_numeric_type;
17use crate::rle::ArrowRle;
18
19/// How [`ArrowRle::slice_agg`] reduces each slice. Null values are skipped,
20/// as in polars and pandas.
21#[derive(Clone, Copy, Debug, PartialEq, Eq)]
22pub enum SliceAgg {
23    /// The sum over all positions, in the values' type; 0 for an empty slice.
24    Sum,
25    /// The number of positions with a non-null value (`Int64`).
26    Count,
27    /// The mean over non-null positions (`Float64`), null if there are none.
28    Mean,
29    /// The variance with `ddof` delta degrees of freedom (`Float64`), null if
30    /// there are `ddof` or fewer non-null positions.
31    Var { ddof: i64 },
32    /// The square root of [`SliceAgg::Var`].
33    Std { ddof: i64 },
34    /// The smallest value ignoring `NaN`, in the values' type; null if none.
35    Min,
36    /// The largest value ignoring `NaN`, in the values' type; null if none.
37    Max,
38}
39
40/// Minimum slices per rayon task; reducing one slice is cheap.
41const SLICES_PER_TASK: usize = 1024;
42
43fn check_slices(starts: &[i64], ends: &[i64]) -> Result<()> {
44    if starts.len() != ends.len() {
45        return Err(runlength::Error::LengthMismatch {
46            what: "starts and ends",
47            left: starts.len(),
48            right: ends.len(),
49        }
50        .into());
51    }
52    for (index, (&start, &end)) in starts.iter().zip(ends).enumerate() {
53        if start < 0 || end < start {
54            return Err(runlength::Error::InvalidInterval { index, start, end }.into());
55        }
56    }
57    Ok(())
58}
59
60/// The runs of many sequences stored back to back, with the slices to
61/// reduce: slice `k` covers `starts[k]..stops[k]` of the sequence at runs
62/// `range(k)`, or nothing if `range(k)` is `None`.
63struct Slices<'a, T: ArrowPrimitiveType, R> {
64    ends: &'a [i64],
65    values: &'a PrimitiveArray<T>,
66    starts: &'a [i64],
67    stops: &'a [i64],
68    range: R,
69}
70
71impl<T, R> Slices<'_, T, R>
72where
73    T: ArrowPrimitiveType,
74    T::Native: Num,
75    R: Fn(usize) -> Option<(usize, usize)> + Sync,
76{
77    /// `(length, value)` for each run slice `k` touches; `None` for nulls.
78    fn runs(&self, k: usize) -> impl Iterator<Item = (i64, Option<T::Native>)> + Clone + '_ {
79        let (lo, hi) = (self.range)(k).unwrap_or((0, 0));
80        let ends = &self.ends[lo..hi];
81        let len = ends.last().copied().unwrap_or(0);
82        let (start, stop) = (self.starts[k].min(len), self.stops[k].min(len));
83        let (first, last) = if start < stop {
84            (
85                ends.partition_point(|&e| e <= start),
86                ends.partition_point(|&e| e < stop) + 1,
87            )
88        } else {
89            (0, 0)
90        };
91        (first..last).map(move |i| {
92            let run_start = if i == 0 { 0 } else { ends[i - 1] }.max(start);
93            let run_end = ends[i].min(stop);
94            let value = self
95                .values
96                .is_valid(lo + i)
97                .then(|| self.values.value(lo + i));
98            (run_end - run_start, value)
99        })
100    }
101
102    fn map<X: Send>(&self, f: impl Fn(usize) -> X + Send + Sync) -> Vec<X>
103    where
104        Self: Sync,
105    {
106        (0..self.starts.len())
107            .into_par_iter()
108            .with_min_len(SLICES_PER_TASK)
109            .map(f)
110            .collect()
111    }
112
113    fn count(&self, k: usize) -> i64 {
114        self.runs(k).filter_map(|(w, v)| v.map(|_| w)).sum::<i64>()
115    }
116
117    fn mean(&self, k: usize) -> Option<f64> {
118        let (n, total) = self.runs(k).fold((0, 0.0), |(n, total), (w, v)| match v {
119            Some(v) => (n + w, total + w as f64 * v.to_f64()),
120            None => (n, total),
121        });
122        (n > 0).then(|| total / n as f64)
123    }
124
125    fn variance(&self, k: usize, ddof: i64) -> Option<f64> {
126        let n = self.count(k);
127        if n - ddof <= 0 {
128            return None;
129        }
130        let mean = self.mean(k)?;
131        let squares: f64 = self
132            .runs(k)
133            .filter_map(|(w, v)| v.map(|v| w as f64 * (v.to_f64() - mean).powi(2)))
134            .sum();
135        Some(squares / (n - ddof) as f64)
136    }
137
138    fn extreme(&self, k: usize, max: bool) -> Option<T::Native> {
139        self.runs(k)
140            .filter_map(|(_, v)| v.filter(|v| v.partial_cmp(v).is_some()))
141            .fold(None, |best, v| match best {
142                Some(b) if (max && v <= b) || (!max && v >= b) => Some(b),
143                _ => Some(v),
144            })
145    }
146
147    fn aggregate(&self, agg: SliceAgg) -> ArrayRef
148    where
149        Self: Sync,
150    {
151        match agg {
152            SliceAgg::Sum => {
153                let sums = self.map(|k| {
154                    self.runs(k)
155                        .filter_map(|(w, v)| v.map(|v| (w, v)))
156                        .fold(T::Native::default(), |acc, (w, v)| {
157                            acc + v * T::Native::from_i64(w)
158                        })
159                });
160                Arc::new(PrimitiveArray::<T>::new(ScalarBuffer::from(sums), None))
161            }
162            SliceAgg::Count => Arc::new(Int64Array::from(self.map(|k| self.count(k)))),
163            SliceAgg::Mean => Arc::new(Float64Array::from(self.map(|k| self.mean(k)))),
164            SliceAgg::Var { ddof } => {
165                Arc::new(Float64Array::from(self.map(|k| self.variance(k, ddof))))
166            }
167            SliceAgg::Std { ddof } => Arc::new(Float64Array::from(
168                self.map(|k| self.variance(k, ddof).map(f64::sqrt)),
169            )),
170            SliceAgg::Min => Arc::new(
171                self.map(|k| self.extreme(k, false))
172                    .into_iter()
173                    .collect::<PrimitiveArray<T>>(),
174            ),
175            SliceAgg::Max => Arc::new(
176                self.map(|k| self.extreme(k, true))
177                    .into_iter()
178                    .collect::<PrimitiveArray<T>>(),
179            ),
180        }
181    }
182}
183
184/// Reduces the slices over `values`, which must be numeric or boolean
185/// (booleans count as 0 and 1).
186fn aggregate(
187    ends: &[i64],
188    values: &ArrayRef,
189    starts: &[i64],
190    stops: &[i64],
191    range: impl Fn(usize) -> Option<(usize, usize)> + Sync,
192    agg: SliceAgg,
193) -> Result<ArrayRef> {
194    check_slices(starts, stops)?;
195    if values.data_type() == &DataType::Boolean {
196        let ints = cast(values, &DataType::Int64)?;
197        let out = aggregate(ends, &ints, starts, stops, range, agg)?;
198        return match agg {
199            SliceAgg::Min | SliceAgg::Max => Ok(cast(&out, &DataType::Boolean)?),
200            _ => Ok(out),
201        };
202    }
203    with_numeric_type!(values, |T| {
204        let slices = Slices {
205            ends,
206            values: values
207                .as_any()
208                .downcast_ref::<PrimitiveArray<T>>()
209                .expect("the type was matched"),
210            starts,
211            stops,
212            range,
213        };
214        Ok(slices.aggregate(agg))
215    })
216}
217
218impl ArrowRle {
219    /// Reduces many slices without copying them, in parallel: entry `k` of
220    /// the result reduces positions `starts[k]..ends[k]` (clamped to the
221    /// sequence) with `agg`. Values must be numeric or boolean.
222    ///
223    /// ```
224    /// use std::sync::Arc;
225    /// use arrow_array::{Array, Float64Array, Int64Array};
226    /// use runlength_arrow::{ArrowRle, SliceAgg};
227    ///
228    /// let rle = ArrowRle::from_runs(&[3, 2], Arc::new(Int64Array::from(vec![1, 4]))).unwrap();
229    /// let means = rle.slice_agg(&[0, 2], &[2, 5], SliceAgg::Mean).unwrap();
230    /// let means = means.as_any().downcast_ref::<Float64Array>().unwrap();
231    /// assert_eq!(means.values().to_vec(), vec![1.0, 3.0]);
232    /// ```
233    pub fn slice_agg(&self, starts: &[i64], ends: &[i64], agg: SliceAgg) -> Result<ArrayRef> {
234        let n = self.num_runs();
235        aggregate(
236            self.ends(),
237            self.values(),
238            starts,
239            ends,
240            |_| Some((0, n)),
241            agg,
242        )
243    }
244}
245
246impl ArrowRleBatch {
247    /// Like [`ArrowRle::slice_agg`], where slice `k` is in the sequence of row
248    /// `k` of `keys`. A key not in the batch gives an empty slice. `keys` must
249    /// have the same column names and types as the batch's keys.
250    pub fn slice_agg(
251        &self,
252        keys: &RecordBatch,
253        starts: &[i64],
254        ends: &[i64],
255        agg: SliceAgg,
256    ) -> Result<ArrayRef> {
257        if keys.num_rows() != starts.len() {
258            return Err(runlength::Error::LengthMismatch {
259                what: "keys and starts",
260                left: keys.num_rows(),
261                right: starts.len(),
262            }
263            .into());
264        }
265        let positions = self.positions(keys)?;
266        let offsets = self.offsets();
267        aggregate(
268            self.ends(),
269            self.values(),
270            starts,
271            ends,
272            |k| positions[k].map(|p| (offsets[p], offsets[p + 1])),
273            agg,
274        )
275    }
276}
runlength_arrow/slices.rs

runlength_arrow/
slices.rs
