miniextendr_api/dataframe/
group.rs

1//! Group-level iteration over a [`DataFrame`] key column.
2//!
3//! Two rungs, cheapest first:
4//!
5//! 1. **Typed rows, grouped Rust-side** — after `Vec::<Row>::from_dataframe(&df)?`,
6//!    grouping is plain Rust. [`group_rows`] makes the idiom discoverable:
7//!
8//!    ```ignore
9//!    let rows: Vec<Obs> = Vec::from_dataframe(&df)?;
10//!    let by_site = group_rows(rows, |r| r.site.clone());
11//!    // by_site: BTreeMap<String, Vec<Obs>> — plain Rust data, rayon-safe.
12//!    ```
13//!
14//! 2. **Untyped, index-based** — [`DataFrame::group_by`] computes group indices
15//!    once (single pass, main thread) without extracting rows:
16//!
17//!    ```ignore
18//!    let grouped = df.group_by("site")?;
19//!    for (key, rows) in grouped.iter() { /* key: &GroupKey, rows: &[usize] */ }
20//!    let mut out = NamedDataFrameListBuilder::with_capacity(grouped.len());
21//!    for (key, sub) in grouped.frames() {
22//!        out = out.push(key.label(), sub);
23//!    }
24//!    ```
25//!
26//! # Key semantics (vs R `split()`)
27//!
28//! - **Group order**: factor keys follow level order (empty levels kept, like
29//!   `split()`); character keys sort in byte order (R sorts in locale collation
30//!   order — identical for ASCII); integer keys sort numerically; logical keys
31//!   order `FALSE`, `TRUE`.
32//! - **`NA` keys form one group, ordered last** — a deliberate deviation from
33//!   `split()`, which silently drops NA-keyed rows. A literal NA *level*
34//!   (`addNA(f)`) also surfaces as [`GroupKey::Na`].
35//! - **Double key columns are an error**: grouping on floating point is a
36//!   footgun — `cut()` or `factor()` the column first.
37
38use std::collections::BTreeMap;
39
40use super::{DataFrame, DataFrameError, FromDataFrame};
41use crate::{SEXP, SEXPTYPE, SexpExt};
42
43// region: group_rows — typed-rows grouping helper (rung 1)
44
45/// Group already-extracted rows by a key function.
46///
47/// Plain Rust — no SEXP contact, so the result is `Send` (given `T: Send`) and
48/// safe to iterate with rayon. Keys order by `Ord`; give NA-able keys a home by
49/// keying on `Option<T>` (`None` sorts first) or a custom enum.
50pub fn group_rows<T, K, F>(rows: Vec<T>, key: F) -> BTreeMap<K, Vec<T>>
51where
52    K: Ord,
53    F: Fn(&T) -> K,
54{
55    let mut groups: BTreeMap<K, Vec<T>> = BTreeMap::new();
56    for row in rows {
57        groups.entry(key(&row)).or_default().push(row);
58    }
59    groups
60}
61// endregion
62
63// region: GroupKey
64
65/// The key of one group produced by [`DataFrame::group_by`].
66#[derive(Debug, Clone, PartialEq, Eq, Hash)]
67pub enum GroupKey {
68    /// A character or factor-level key.
69    Str(String),
70    /// An integer key.
71    Int(i32),
72    /// A logical key.
73    Bool(bool),
74    /// The NA-keyed group (always ordered last).
75    Na,
76}
77
78impl GroupKey {
79    /// R-facing label for this key — suitable as a name in a result list
80    /// (matches how R prints the value: `TRUE`/`FALSE`, `NA`, digits).
81    pub fn label(&self) -> String {
82        match self {
83            GroupKey::Str(s) => s.clone(),
84            GroupKey::Int(i) => i.to_string(),
85            GroupKey::Bool(true) => "TRUE".to_string(),
86            GroupKey::Bool(false) => "FALSE".to_string(),
87            GroupKey::Na => "NA".to_string(),
88        }
89    }
90}
91
92impl std::fmt::Display for GroupKey {
93    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94        match self {
95            GroupKey::Str(s) => f.write_str(s),
96            GroupKey::Int(i) => write!(f, "{}", i),
97            GroupKey::Bool(true) => f.write_str("TRUE"),
98            GroupKey::Bool(false) => f.write_str("FALSE"),
99            GroupKey::Na => f.write_str("NA"),
100        }
101    }
102}
103// endregion
104
105// region: GroupedDataFrame
106
107/// A [`DataFrame`] partitioned by the values of one key column.
108///
109/// Produced by [`DataFrame::group_by`]. Holds the source frame plus one
110/// `(key, row-indices)` pair per group; nothing is copied until you ask for
111/// [`frames`](Self::frames) or [`extract`](Self::extract).
112///
113/// # GC rooting
114///
115/// The source frame is preserved on R's precious list
116/// (`R_PreserveObject`) for this struct's lifetime and released on drop —
117/// order-independent, unlike the PROTECT stack, so the struct can be held
118/// across arbitrary allocations (e.g. a locally built frame from
119/// [`DataFrame::builder`], which is unprotected once `build()` returns).
120/// Without this, the per-group allocations in [`frames`](Self::frames) /
121/// [`extract`](Self::extract) could collect the source mid-iteration.
122/// Main-thread-only (holds a SEXP; `!Send`).
123pub struct GroupedDataFrame {
124    source: DataFrame,
125    groups: Vec<(GroupKey, Vec<usize>)>,
126}
127
128impl Drop for GroupedDataFrame {
129    fn drop(&mut self) {
130        // SAFETY: main thread (construction invariant — SEXP is !Send);
131        // releases the preserve taken in `group_by`.
132        unsafe { crate::sys::R_ReleaseObject(self.source.as_sexp()) };
133    }
134}
135
136impl GroupedDataFrame {
137    /// Number of groups (empty factor levels included).
138    pub fn len(&self) -> usize {
139        self.groups.len()
140    }
141
142    /// Whether there are no groups.
143    pub fn is_empty(&self) -> bool {
144        self.groups.is_empty()
145    }
146
147    /// The frame this grouping was computed from.
148    pub fn source(&self) -> &DataFrame {
149        &self.source
150    }
151
152    /// Iterate `(key, row-indices)` pairs in group order. Indices are 0-based
153    /// rows of [`source`](Self::source).
154    pub fn iter(&self) -> impl Iterator<Item = (&GroupKey, &[usize])> {
155        self.groups.iter().map(|(k, idx)| (k, idx.as_slice()))
156    }
157
158    /// Iterate `(key, sub-frame)` pairs, materialising each group as its own
159    /// [`DataFrame`] via [`DataFrame::select_rows`].
160    ///
161    /// Main thread only. Each yielded frame is **unprotected** — root it
162    /// before the next iteration allocates (e.g. push it straight into a
163    /// [`NamedDataFrameListBuilder`](crate::dataframe::NamedDataFrameListBuilder),
164    /// which protects on push) or convert it to Rust data immediately.
165    pub fn frames(&self) -> impl Iterator<Item = (&GroupKey, DataFrame)> {
166        self.groups
167            .iter()
168            .map(|(k, idx)| (k, self.source.select_rows(idx)))
169    }
170
171    /// Extract typed rows once, then partition them by group.
172    ///
173    /// One `Vec::<T>::from_dataframe` pass over the whole frame, then a
174    /// move-partition by the stored indices — no per-group R subsetting and no
175    /// `Clone` bound. The result is plain Rust data (rayon-safe afterwards).
176    pub fn extract<T>(&self) -> Result<Vec<(GroupKey, Vec<T>)>, DataFrameError>
177    where
178        Vec<T>: FromDataFrame,
179    {
180        let rows = Vec::<T>::from_dataframe(&self.source)?;
181        let mut slots: Vec<Option<T>> = rows.into_iter().map(Some).collect();
182        Ok(self
183            .groups
184            .iter()
185            .map(|(key, idx)| {
186                let group_rows: Vec<T> = idx
187                    .iter()
188                    .map(|&i| slots[i].take().expect("group indices are disjoint"))
189                    .collect();
190                (key.clone(), group_rows)
191            })
192            .collect())
193    }
194}
195// endregion
196
197// region: DataFrame::group_by
198
199impl DataFrame {
200    /// Partition this frame's rows by the values of the named column.
201    ///
202    /// Computes group indices in a single pass on the main thread. Supported
203    /// key columns: factor (fast path — levels are the keys, level order kept,
204    /// empty levels included), character, integer, and logical. Double columns
205    /// error — `cut()` or `factor()` the column in R first.
206    ///
207    /// NA keys form one group, ordered last (unlike R `split()`, which drops
208    /// NA-keyed rows). See the [module docs](self) for the full key semantics.
209    pub fn group_by(&self, col: &str) -> Result<GroupedDataFrame, DataFrameError> {
210        let column = self
211            .column_raw(col)
212            .ok_or_else(|| DataFrameError::NoSuchColumn(col.to_string()))?;
213        let groups = if column.is_factor() {
214            factor_groups(column)
215        } else {
216            match column.type_of() {
217                SEXPTYPE::STRSXP => character_groups(column),
218                SEXPTYPE::INTSXP => integer_groups(column),
219                SEXPTYPE::LGLSXP => logical_groups(column),
220                other => {
221                    return Err(DataFrameError::UnsupportedGroupColumn {
222                        column: col.to_string(),
223                        type_of: format!("{:?}", other),
224                    });
225                }
226            }
227        };
228        // Root the source for the GroupedDataFrame's lifetime (see its GC
229        // rooting docs). R_PreserveObject conses onto the precious list —
230        // an allocation that can itself GC — so PROTECT across the call.
231        unsafe {
232            let _guard = crate::OwnedProtect::new(self.sexp);
233            crate::sys::R_PreserveObject(self.sexp);
234        }
235        Ok(GroupedDataFrame {
236            source: *self,
237            groups,
238        })
239    }
240}
241
242/// Factor fast path: levels are the keys (level order, empty levels kept).
243/// NA codes — and a literal NA level from `addNA()` — land in [`GroupKey::Na`].
244fn factor_groups(column: SEXP) -> Vec<(GroupKey, Vec<usize>)> {
245    // SAFETY: factor columns are INTSXP; as_slice handles empty vectors.
246    let codes: &[i32] = unsafe { column.as_slice() };
247    let levels = column.get_levels();
248    let n_levels: usize = if levels.is_nil() { 0 } else { levels.len() };
249
250    let mut buckets: Vec<Vec<usize>> = vec![Vec::new(); n_levels];
251    let mut na_bucket: Vec<usize> = Vec::new();
252    for (row, &code) in codes.iter().enumerate() {
253        if code == i32::MIN {
254            na_bucket.push(row);
255        } else {
256            buckets[(code - 1) as usize].push(row);
257        }
258    }
259
260    let mut groups: Vec<(GroupKey, Vec<usize>)> = Vec::with_capacity(n_levels + 1);
261    for (lvl, bucket) in buckets.into_iter().enumerate() {
262        let key = match levels.string_elt_str(lvl as isize) {
263            Some(label) => GroupKey::Str(label.to_string()),
264            None => GroupKey::Na, // addNA() level
265        };
266        groups.push((key, bucket));
267    }
268    if !na_bucket.is_empty() {
269        groups.push((GroupKey::Na, na_bucket));
270    }
271    groups
272}
273
274/// Character keys: byte-order sort (BTreeMap), NA last.
275fn character_groups(column: SEXP) -> Vec<(GroupKey, Vec<usize>)> {
276    let n = column.len() as isize;
277    let mut map: BTreeMap<String, Vec<usize>> = BTreeMap::new();
278    let mut na_bucket: Vec<usize> = Vec::new();
279    for i in 0..n {
280        match column.string_elt_str(i) {
281            Some(s) => map.entry(s.to_string()).or_default().push(i as usize),
282            None => na_bucket.push(i as usize),
283        }
284    }
285    let mut groups: Vec<(GroupKey, Vec<usize>)> = map
286        .into_iter()
287        .map(|(k, idx)| (GroupKey::Str(k), idx))
288        .collect();
289    if !na_bucket.is_empty() {
290        groups.push((GroupKey::Na, na_bucket));
291    }
292    groups
293}
294
295/// Integer keys: numeric sort (BTreeMap), NA (`i32::MIN`) last.
296fn integer_groups(column: SEXP) -> Vec<(GroupKey, Vec<usize>)> {
297    // SAFETY: INTSXP column; as_slice handles empty vectors.
298    let values: &[i32] = unsafe { column.as_slice() };
299    let mut map: BTreeMap<i32, Vec<usize>> = BTreeMap::new();
300    let mut na_bucket: Vec<usize> = Vec::new();
301    for (row, &v) in values.iter().enumerate() {
302        if v == i32::MIN {
303            na_bucket.push(row);
304        } else {
305            map.entry(v).or_default().push(row);
306        }
307    }
308    let mut groups: Vec<(GroupKey, Vec<usize>)> = map
309        .into_iter()
310        .map(|(k, idx)| (GroupKey::Int(k), idx))
311        .collect();
312    if !na_bucket.is_empty() {
313        groups.push((GroupKey::Na, na_bucket));
314    }
315    groups
316}
317
318/// Logical keys: `FALSE` then `TRUE` (R's sort order), NA last.
319/// Only keys present in the data appear.
320fn logical_groups(column: SEXP) -> Vec<(GroupKey, Vec<usize>)> {
321    let n = column.len() as isize;
322    let mut false_bucket: Vec<usize> = Vec::new();
323    let mut true_bucket: Vec<usize> = Vec::new();
324    let mut na_bucket: Vec<usize> = Vec::new();
325    for i in 0..n {
326        match column.logical_elt(i) {
327            0 => false_bucket.push(i as usize),
328            v if v == i32::MIN => na_bucket.push(i as usize),
329            _ => true_bucket.push(i as usize),
330        }
331    }
332    let mut groups: Vec<(GroupKey, Vec<usize>)> = Vec::with_capacity(3);
333    if !false_bucket.is_empty() {
334        groups.push((GroupKey::Bool(false), false_bucket));
335    }
336    if !true_bucket.is_empty() {
337        groups.push((GroupKey::Bool(true), true_bucket));
338    }
339    if !na_bucket.is_empty() {
340        groups.push((GroupKey::Na, na_bucket));
341    }
342    groups
343}
344// endregion
345
346#[cfg(test)]
347mod tests {
348    use super::*;
349
350    #[test]
351    fn group_rows_partitions_and_orders_by_key() {
352        let rows = vec![("b", 1), ("a", 2), ("b", 3), ("c", 4)];
353        let grouped = group_rows(rows, |r| r.0);
354        let keys: Vec<&str> = grouped.keys().copied().collect();
355        assert_eq!(keys, vec!["a", "b", "c"]);
356        assert_eq!(grouped["b"], vec![("b", 1), ("b", 3)]);
357    }
358
359    #[test]
360    fn group_rows_option_key_gives_na_a_home() {
361        let rows = vec![(Some(2), "x"), (None, "y"), (Some(1), "z")];
362        let grouped = group_rows(rows, |r| r.0);
363        let keys: Vec<Option<i32>> = grouped.keys().copied().collect();
364        assert_eq!(keys, vec![None, Some(1), Some(2)]);
365    }
366
367    #[test]
368    fn group_key_labels_match_r_printing() {
369        assert_eq!(GroupKey::Str("a".into()).label(), "a");
370        assert_eq!(GroupKey::Int(-3).label(), "-3");
371        assert_eq!(GroupKey::Bool(true).label(), "TRUE");
372        assert_eq!(GroupKey::Bool(false).label(), "FALSE");
373        assert_eq!(GroupKey::Na.label(), "NA");
374        assert_eq!(GroupKey::Na.to_string(), "NA");
375    }
376}
miniextendr_api/dataframe/group.rs

miniextendr_api/dataframe/
group.rs
