runlength_arrow/
coverage.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use arrow_array::cast::AsArray;
5use arrow_array::types::Int64Type;
6use arrow_array::{
7    Array, ArrayRef, ArrowPrimitiveType, Int64Array, LargeListArray, PrimitiveArray, RecordBatch,
8    RecordBatchOptions, StringArray, UInt64Array,
9};
10use arrow_buffer::{OffsetBuffer, ScalarBuffer};
11use arrow_schema::{DataType, Field, Schema};
12use arrow_select::take::take;
13use rayon::prelude::*;
14use runlength::{Rle, RleBatch};
15
16use crate::batch::{ArrowRleBatch, key_converter, key_rows};
17use crate::error::{Error, Result};
18use crate::numeric::with_numeric_type;
19use crate::rle::ArrowRle;
20
21fn to_array<T: ArrowPrimitiveType>(values: Vec<T::Native>) -> ArrayRef {
22    Arc::new(PrimitiveArray::<T>::new(ScalarBuffer::from(values), None))
23}
24
25fn native<T: ArrowPrimitiveType>(values: &dyn Array) -> Result<&[T::Native]> {
26    if values.null_count() > 0 {
27        return Err(Error::UnexpectedNulls);
28    }
29    Ok(values.as_primitive::<T>().values())
30}
31
32/// Coverage of half-open intervals `[start, end)`.
33///
34/// With `values`, each position holds the sum of the values of the intervals
35/// covering it, in the values' own numeric type (which must have no nulls).
36/// Without, it holds the number of covering intervals as `Int64`. Uncovered
37/// positions are 0. The result covers `0..max(end)`, or `0..length`.
38pub fn coverage(
39    starts: &[i64],
40    ends: &[i64],
41    values: Option<&dyn Array>,
42    length: Option<i64>,
43) -> Result<ArrowRle> {
44    let Some(values) = values else {
45        let rle = runlength::coverage_count(starts, ends, length)?;
46        return Ok(ArrowRle::from_rle::<Int64Type>(&rle));
47    };
48    with_numeric_type!(values, |T| {
49        let rle = runlength::coverage_sum(starts, ends, native::<T>(values)?, length)?;
50        Ok(ArrowRle::from_rle::<T>(&rle))
51    })
52}
53
54/// How [`coverage_agg`] combines the intervals covering a position.
55#[derive(Clone, Copy, Debug, PartialEq, Eq)]
56pub enum Agg {
57    /// The sum of the values (the values' numeric type).
58    Sum,
59    /// The number of intervals (`Int64`); values are ignored.
60    Count,
61    /// The largest value (numbers or strings), ignoring `NaN` where possible.
62    Max,
63    /// The smallest value (numbers or strings), ignoring `NaN` where possible.
64    Min,
65}
66
67fn from_options<T: ArrowPrimitiveType>(rle: Rle<Option<T::Native>>) -> ArrowRle {
68    let (ends, values) = rle.into_parts();
69    ArrowRle::try_new(
70        ends,
71        Arc::new(values.into_iter().collect::<PrimitiveArray<T>>()),
72    )
73    .expect("a core sequence has valid ends")
74}
75
76fn strings_from_options(rle: Rle<Option<&str>>) -> ArrowRle {
77    let (ends, values) = rle.into_parts();
78    ArrowRle::try_new(ends, Arc::new(values.into_iter().collect::<StringArray>()))
79        .expect("a core sequence has valid ends")
80}
81
82/// The core max/min coverage, for any ordered value type.
83fn core_extreme<V>(
84    starts: &[i64],
85    ends: &[i64],
86    values: &[V],
87    length: Option<i64>,
88    max: bool,
89) -> runlength::Result<Rle<Option<V>>>
90where
91    V: PartialOrd + Clone + runlength::Identical,
92{
93    if max {
94        runlength::coverage_max(starts, ends, values, length)
95    } else {
96        runlength::coverage_min(starts, ends, values, length)
97    }
98}
99
100/// Max/min coverage over numbers or strings.
101fn extreme(
102    starts: &[i64],
103    ends: &[i64],
104    values: &dyn Array,
105    length: Option<i64>,
106    max: bool,
107) -> Result<ArrowRle> {
108    if let Some(strings) = values.as_string_opt::<i32>() {
109        if strings.null_count() > 0 {
110            return Err(Error::UnexpectedNulls);
111        }
112        let strings: Vec<&str> = strings.iter().map(|s| s.expect("no nulls")).collect();
113        return Ok(strings_from_options(core_extreme(
114            starts, ends, &strings, length, max,
115        )?));
116    }
117    with_numeric_type!(values, |T| {
118        Ok(from_options::<T>(core_extreme(
119            starts,
120            ends,
121            native::<T>(values)?,
122            length,
123            max,
124        )?))
125    })
126}
127
128/// Coverage of half-open intervals `[start, end)` with any aggregation.
129///
130/// Unlike [`coverage`], positions no interval covers are null (for
131/// [`Agg::Count`] they are 0); use [`ArrowRle::fill_null`] to fill them.
132/// `values` are required except for [`Agg::Count`].
133pub fn coverage_agg(
134    starts: &[i64],
135    ends: &[i64],
136    values: Option<&dyn Array>,
137    length: Option<i64>,
138    agg: Agg,
139) -> Result<ArrowRle> {
140    let values = match (agg, values) {
141        (Agg::Count, _) => return coverage(starts, ends, None, length),
142        (_, Some(values)) => values,
143        (_, None) => return Err(Error::MissingValues),
144    };
145    match agg {
146        Agg::Sum => with_numeric_type!(values, |T| {
147            let rle = runlength::coverage_sum_opt(starts, ends, native::<T>(values)?, length)?;
148            Ok(from_options::<T>(rle))
149        }),
150        Agg::Max => extreme(starts, ends, values, length, true),
151        Agg::Min => extreme(starts, ends, values, length, false),
152        Agg::Count => unreachable!("handled above"),
153    }
154}
155
156/// [`coverage_agg`] per key, computed in parallel; see [`coverage_by_key`]
157/// for the keys.
158pub fn coverage_agg_by_key(
159    keys: &RecordBatch,
160    starts: &[i64],
161    ends: &[i64],
162    values: Option<&dyn Array>,
163    agg: Agg,
164) -> Result<ArrowRleBatch> {
165    if agg == Agg::Count {
166        return coverage_by_key(keys, starts, ends, None);
167    }
168    let values = values.ok_or(Error::MissingValues)?;
169    let (ids, first_rows) = group_rows(keys)?;
170    let groups = group_indices(&ids, first_rows.len());
171    let rles = groups
172        .par_iter()
173        .map(|idx| {
174            let starts: Vec<i64> = idx.iter().map(|&i| starts[i]).collect();
175            let ends: Vec<i64> = idx.iter().map(|&i| ends[i]).collect();
176            let indices = UInt64Array::from_iter_values(idx.iter().map(|&i| i as u64));
177            let values = take(values, &indices, None)?;
178            coverage_agg(&starts, &ends, Some(values.as_ref()), None, agg)
179        })
180        .collect::<Result<Vec<_>>>()?;
181    let first_rows = UInt64Array::from(first_rows);
182    let columns = keys
183        .columns()
184        .iter()
185        .map(|c| take(c.as_ref(), &first_rows, None))
186        .collect::<std::result::Result<Vec<_>, _>>()?;
187    let keys = RecordBatch::try_new(keys.schema(), columns)?;
188    let data_type = rles
189        .first()
190        .map_or(DataType::Null, |r| r.data_type().clone());
191    ArrowRleBatch::from_rles(keys, rles, &data_type)
192}
193
194fn group_indices(ids: &[usize], n: usize) -> Vec<Vec<usize>> {
195    let mut groups = vec![Vec::new(); n];
196    for (i, &g) in ids.iter().enumerate() {
197        groups[g].push(i);
198    }
199    groups
200}
201
202/// Numbers each key row by first appearance; returns the group of every row
203/// and the first row of every group.
204fn group_rows(keys: &RecordBatch) -> Result<(Vec<usize>, Vec<u64>)> {
205    let converter = key_converter(keys)?;
206    let rows = key_rows(&converter, keys)?;
207    let mut groups = HashMap::new();
208    let mut first_rows = Vec::new();
209    let ids = rows
210        .iter()
211        .enumerate()
212        .map(|(i, row)| {
213            *groups.entry(row).or_insert_with(|| {
214                first_rows.push(i as u64);
215                first_rows.len() - 1
216            })
217        })
218        .collect();
219    Ok((ids, first_rows))
220}
221
222fn into_arrow_batch<T: ArrowPrimitiveType>(
223    keys: &RecordBatch,
224    first_rows: Vec<u64>,
225    batch: RleBatch<usize, T::Native>,
226) -> Result<ArrowRleBatch> {
227    let first_rows = UInt64Array::from(first_rows);
228    let columns = keys
229        .columns()
230        .iter()
231        .map(|c| take(c.as_ref(), &first_rows, None))
232        .collect::<std::result::Result<Vec<_>, _>>()?;
233    let keys = RecordBatch::try_new(keys.schema(), columns)?;
234    let (_, offsets, ends, values) = batch.into_parts();
235    ArrowRleBatch::from_parts_unchecked(keys, offsets, ends, to_array::<T>(values))
236}
237
238/// [`coverage`] per key, computed in parallel.
239///
240/// Row `i` of `keys` (any number of columns, of any types) is the key of
241/// interval `i`. Keys appear in order of first appearance, and each sequence
242/// covers `0..max(end)` of its own intervals.
243pub fn coverage_by_key(
244    keys: &RecordBatch,
245    starts: &[i64],
246    ends: &[i64],
247    values: Option<&dyn Array>,
248) -> Result<ArrowRleBatch> {
249    let (ids, first_rows) = group_rows(keys)?;
250    let Some(values) = values else {
251        let batch = runlength::coverage_count_by_key(&ids, starts, ends)?;
252        return into_arrow_batch::<Int64Type>(keys, first_rows, batch);
253    };
254    with_numeric_type!(values, |T| {
255        let batch = runlength::coverage_sum_by_key(&ids, starts, ends, native::<T>(values)?)?;
256        into_arrow_batch::<T>(keys, first_rows, batch)
257    })
258}
259
260/// [`runlength::overlap_plan`] per key, as one long table: the key columns,
261/// then `start` and `end` of each segment and `ids`, a list of the intervals
262/// covering it (positions in the input, sorted). Keys appear in order of
263/// first appearance; each key's segments cover `0..max(end)` of its own
264/// intervals.
265///
266/// Exploding `ids` and joining interval values onto them gives any
267/// aggregation per segment, for example with a polars `group_by`.
268pub fn overlap_plan_by_key(
269    keys: &RecordBatch,
270    starts: &[i64],
271    ends: &[i64],
272) -> Result<RecordBatch> {
273    let (ids, first_rows) = group_rows(keys)?;
274    let (groups, plans) = runlength::overlap_plan_by_key(&ids, starts, ends)?;
275    debug_assert!(groups.iter().enumerate().all(|(i, &g)| i == g));
276    let segments: usize = plans.iter().map(|p| p.num_segments()).sum();
277    let mut key_rows = Vec::with_capacity(segments);
278    let mut seg_starts = Vec::with_capacity(segments);
279    let mut seg_ends = Vec::with_capacity(segments);
280    let mut list_offsets = Vec::with_capacity(segments + 1);
281    list_offsets.push(0i64);
282    let mut items = Vec::new();
283    for (plan, &row) in plans.iter().zip(&first_rows) {
284        let mut prev = 0;
285        for s in 0..plan.num_segments() {
286            key_rows.push(row);
287            seg_starts.push(prev);
288            seg_ends.push(plan.ends[s]);
289            prev = plan.ends[s];
290            items.extend(plan.ids(s).iter().map(|&i| i as i64));
291            list_offsets.push(items.len() as i64);
292        }
293    }
294    let key_rows = UInt64Array::from(key_rows);
295    let mut columns = keys
296        .columns()
297        .iter()
298        .map(|c| take(c.as_ref(), &key_rows, None))
299        .collect::<std::result::Result<Vec<_>, _>>()?;
300    let mut fields: Vec<Field> = keys
301        .schema()
302        .fields()
303        .iter()
304        .map(|f| f.as_ref().clone())
305        .collect();
306    let item = Arc::new(Field::new("item", DataType::Int64, false));
307    fields.push(Field::new("start", DataType::Int64, false));
308    fields.push(Field::new("end", DataType::Int64, false));
309    fields.push(Field::new("ids", DataType::LargeList(item.clone()), false));
310    columns.push(Arc::new(Int64Array::from(seg_starts)));
311    columns.push(Arc::new(Int64Array::from(seg_ends)));
312    columns.push(Arc::new(LargeListArray::new(
313        item,
314        OffsetBuffer::new(ScalarBuffer::from(list_offsets)),
315        Arc::new(Int64Array::from(items)),
316        None,
317    )));
318    Ok(RecordBatch::try_new_with_options(
319        Arc::new(Schema::new(fields)),
320        columns,
321        &RecordBatchOptions::new().with_row_count(Some(segments)),
322    )?)
323}
runlength_arrow/coverage.rs

runlength_arrow/
coverage.rs
