runlength_arrow/
transform.rs

1use std::collections::HashMap;
2
3use arrow_array::cast::AsArray;
4use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, UInt64Array};
5use arrow_schema::{ArrowError, DataType};
6use arrow_select::concat::concat as concat_arrays;
7use arrow_select::interleave::interleave;
8use arrow_select::take::take;
9use arrow_select::zip::zip;
10use rayon::prelude::*;
11use runlength::{Rle, RleLike, plan_many};
12
13use crate::batch::{ArrowRleBatch, key_converter, key_rows};
14use crate::error::{Error, Result};
15use crate::kernels::Kernel;
16use crate::rle::{ArrowRle, default_fill};
17
18/// Joins sequences end to end and merges identical neighbours, including at
19/// the joins. All sequences must have the same value type.
20pub fn concat(rles: &[ArrowRle]) -> Result<ArrowRle> {
21    let Some(first) = rles.first() else {
22        return Ok(ArrowRle::empty(&DataType::Null));
23    };
24    let mut ends = Vec::with_capacity(rles.iter().map(|r| r.num_runs()).sum());
25    let mut offset = 0;
26    for rle in rles {
27        ends.extend(rle.ends().iter().map(|&e| e + offset));
28        offset += rle.len();
29    }
30    let values: Vec<&dyn Array> = rles.iter().map(|r| r.values().as_ref()).collect();
31    let values = if values.len() == 1 {
32        first.values().clone()
33    } else {
34        concat_arrays(&values)?
35    };
36    ArrowRle::try_new(ends, values)?.merge()
37}
38
39/// Combines several sequences position by position, and merges the result.
40///
41/// `f` receives each input's values gathered onto the segments of the
42/// [`plan_many`] (equal-length arrays, one per input) and must return one
43/// value per segment. Shorter inputs are padded with their fill value (one
44/// optional fill per input); with none, unequal lengths are an error.
45pub fn zip_many<F>(rles: &[ArrowRle], fills: &[Option<&dyn Array>], f: F) -> Result<ArrowRle>
46where
47    F: FnOnce(&[ArrayRef]) -> std::result::Result<ArrayRef, ArrowError>,
48{
49    if fills.len() != rles.len() {
50        return Err(runlength::Error::LengthMismatch {
51            what: "fills and sequences",
52            left: fills.len(),
53            right: rles.len(),
54        }
55        .into());
56    }
57    let max = rles.iter().map(|r| r.len()).max().unwrap_or(0);
58    for (rle, fill) in rles.iter().zip(fills) {
59        if rle.len() < max && fill.is_none() {
60            return Err(runlength::Error::UnequalLength {
61                left: rle.len(),
62                right: max,
63            }
64            .into());
65        }
66    }
67    let ends: Vec<&[i64]> = rles.iter().map(|r| r.ends()).collect();
68    let plan = plan_many(&ends);
69    let gathered = rles
70        .iter()
71        .zip(fills)
72        .zip(&plan.idx)
73        .map(|((rle, fill), idx)| crate::rle::gather(rle.values(), idx, *fill))
74        .collect::<Result<Vec<_>>>()?;
75    let values = f(&gathered)?;
76    if values.len() != plan.ends.len() {
77        return Err(Error::WrongLength {
78            expected: plan.ends.len(),
79            actual: values.len(),
80        });
81    }
82    ArrowRle::try_new(plan.ends, values)?.merge()
83}
84
85impl ArrowRle {
86    /// Shifts values `distance` positions right (left if negative), keeping
87    /// the length; positions shifted in hold `fill`, a single value of the
88    /// same type. The result is merged.
89    pub fn shift(&self, distance: i64, fill: &dyn Array) -> Result<ArrowRle> {
90        crate::rle::check_fill(Some(fill), self.data_type())?;
91        let len = self.len();
92        if distance == 0 {
93            return self.clone().merge();
94        }
95        let fill_run = |n: i64| ArrowRle::from_runs(&[n], crate::rle::array_ref(fill));
96        if distance.unsigned_abs() >= len as u64 {
97            return fill_run(len);
98        }
99        if distance > 0 {
100            concat(&[fill_run(distance)?, self.slice(0, len - distance)])
101        } else {
102            concat(&[self.slice(-distance, len), fill_run(-distance)?])
103        }
104    }
105
106    /// The sequence back to front.
107    pub fn reversed(&self) -> Result<ArrowRle> {
108        let rle = Rle::from_ends(self.ends().to_vec(), vec![(); self.num_runs()])?.reversed();
109        let n = self.num_runs() as u64;
110        let indices = UInt64Array::from_iter_values((0..n).rev());
111        ArrowRle::try_new(
112            rle.into_parts().0,
113            take(self.values().as_ref(), &indices, None)?,
114        )
115    }
116
117    /// For a boolean sequence, the half-open ranges `(starts, ends)` where it
118    /// is true (nulls count as false), with touching ranges joined.
119    pub fn true_ranges(&self) -> Result<(Vec<i64>, Vec<i64>)> {
120        let values = self
121            .values()
122            .as_boolean_opt()
123            .ok_or_else(|| Error::UnsupportedType(self.data_type().clone()))?;
124        let flags: Vec<bool> = values.iter().map(|v| v == Some(true)).collect();
125        let rle = Rle::from_ends(self.ends().to_vec(), flags)?;
126        Ok(rle.ranges_where(|&v| v))
127    }
128
129    /// Replaces nulls with `fill`, a single value of the same type.
130    pub fn fill_null(&self, fill: &dyn Array) -> Result<ArrowRle> {
131        crate::rle::check_fill(Some(fill), self.data_type())?;
132        if self.values().null_count() == 0 {
133            return Ok(self.clone());
134        }
135        let is_null = arrow_arith::boolean::is_null(self.values().as_ref())?;
136        let fill = arrow_array::Scalar::new(crate::rle::array_ref(fill));
137        let values = zip(&is_null, &fill, self.values())?;
138        ArrowRle::try_new(self.ends().to_vec(), values)?.merge()
139    }
140}
141
142impl ArrowRleBatch {
143    /// Combines the sequences whose keys agree on the `keep` columns, folding
144    /// each group with `kernel` in key order, in parallel across groups.
145    ///
146    /// For example, keeping only `Chromosome` with [`kernels::add`] adds the
147    /// strands of each chromosome together. Sequences of different lengths
148    /// are padded with [`default_fill`].
149    ///
150    /// [`kernels::add`]: crate::kernels::add
151    pub fn collapse(&self, keep: &[&str], kernel: Kernel) -> Result<ArrowRleBatch> {
152        let keys = self.keys();
153        let indices = keep
154            .iter()
155            .map(|name| keys.schema().index_of(name))
156            .collect::<std::result::Result<Vec<_>, _>>()?;
157        let kept = keys.project(&indices)?;
158        let converter = key_converter(&kept)?;
159        let rows = key_rows(&converter, &kept)?;
160
161        let mut groups: Vec<Vec<usize>> = Vec::new();
162        let mut first_rows = Vec::new();
163        let mut positions = HashMap::new();
164        for (i, row) in rows.iter().enumerate() {
165            let g = *positions.entry(row).or_insert_with(|| {
166                first_rows.push(i as u64);
167                groups.push(Vec::new());
168                groups.len() - 1
169            });
170            groups[g].push(i);
171        }
172
173        let data_type = self.values().data_type();
174        let fill = default_fill(data_type);
175        let rles = groups
176            .par_iter()
177            .enumerate()
178            .map(|(position, members)| {
179                let fold = || {
180                    let mut acc = self.get_index(members[0]).expect("in range").merge()?;
181                    for &i in &members[1..] {
182                        let next = self.get_index(i).expect("in range");
183                        acc = acc.zip_with(&next, fill.as_deref(), fill.as_deref(), kernel)?;
184                    }
185                    Ok(acc)
186                };
187                fold().map_err(|e: Error| e.in_key(position))
188            })
189            .collect::<Result<Vec<_>>>()?;
190
191        let first_rows = UInt64Array::from(first_rows);
192        let columns = kept
193            .columns()
194            .iter()
195            .map(|c| take(c.as_ref(), &first_rows, None))
196            .collect::<std::result::Result<Vec<_>, _>>()?;
197        let new_keys = RecordBatch::try_new(kept.schema(), columns)?;
198        let out_type = rles
199            .first()
200            .map_or(data_type.clone(), |r| r.data_type().clone());
201        ArrowRleBatch::from_rles(new_keys, rles, &out_type)
202    }
203}
204
205impl ArrowRleBatch {
206    /// Shifts every sequence, in parallel; see [`ArrowRle::shift`].
207    pub fn shift(&self, distance: i64, fill: &dyn Array) -> Result<ArrowRleBatch> {
208        self.map_rles(|r| r.shift(distance, fill))
209    }
210
211    /// Reverses every sequence, in parallel.
212    pub fn reversed(&self) -> Result<ArrowRleBatch> {
213        self.map_rles(ArrowRle::reversed)
214    }
215
216    /// Transforms every sequence in parallel, keeping the keys. All results
217    /// must have the same value type.
218    pub fn map_rles<F>(&self, f: F) -> Result<ArrowRleBatch>
219    where
220        F: Fn(&ArrowRle) -> Result<ArrowRle> + Sync,
221    {
222        let rles = (0..self.len())
223            .into_par_iter()
224            .map(|i| f(&self.get_index(i).expect("in range")).map_err(|e| e.in_key(i)))
225            .collect::<Result<Vec<_>>>()?;
226        let data_type = rles
227            .first()
228            .map_or(self.values().data_type().clone(), |r| r.data_type().clone());
229        ArrowRleBatch::from_rles(self.keys().clone(), rles, &data_type)
230    }
231
232    /// Joins batches end to end key by key, in parallel: each key's sequences
233    /// are concatenated in batch order, and a batch without the key adds
234    /// nothing. Keys appear in order of first appearance. All batches must
235    /// have the same key columns and value type.
236    pub fn concat(batches: &[ArrowRleBatch]) -> Result<ArrowRleBatch> {
237        let Some(first) = batches.first() else {
238            return Err(Error::NoInputs);
239        };
240        let converter = key_converter(first.keys())?;
241        let mut groups: Vec<Vec<(usize, usize)>> = Vec::new();
242        let mut positions = HashMap::new();
243        for (b, batch) in batches.iter().enumerate() {
244            first.check_key_schema(batch.keys())?;
245            let rows = key_rows(&converter, batch.keys())?;
246            for (i, row) in rows.iter().enumerate() {
247                let g = *positions.entry(row.owned()).or_insert_with(|| {
248                    groups.push(Vec::new());
249                    groups.len() - 1
250                });
251                groups[g].push((b, i));
252            }
253        }
254        let rles = groups
255            .par_iter()
256            .enumerate()
257            .map(|(position, members)| {
258                let parts: Vec<ArrowRle> = members
259                    .iter()
260                    .map(|&(b, i)| batches[b].get_index(i).expect("in range"))
261                    .collect();
262                concat(&parts).map_err(|e| e.in_key(position))
263            })
264            .collect::<Result<Vec<_>>>()?;
265        let sources: Vec<(usize, usize)> = groups.iter().map(|g| g[0]).collect();
266        let columns = (0..first.keys().num_columns())
267            .map(|c| {
268                let arrays: Vec<&dyn Array> = batches
269                    .iter()
270                    .map(|b| b.keys().column(c).as_ref())
271                    .collect();
272                interleave(&arrays, &sources)
273            })
274            .collect::<std::result::Result<Vec<_>, _>>()?;
275        let keys = RecordBatch::try_new_with_options(
276            first.keys().schema(),
277            columns,
278            &RecordBatchOptions::new().with_row_count(Some(sources.len())),
279        )?;
280        ArrowRleBatch::from_rles(keys, rles, first.values().data_type())
281    }
282}
runlength_arrow/transform.rs

runlength_arrow/
transform.rs
