Skip to main content

rich_ext/table/
group.rs

1//! Grouping rows by a column, with per-group aggregates.
2//!
3//! [`GroupBy`] names the grouped column and the [`Aggregate`]s to compute for
4//! each group. Groups appear in the order their first row appears, so sort by
5//! the grouped column first to get them in key order. [`TableData`] renders
6//! each group as a header row (styled `table.group`), its rows, and a summary
7//! row (styled `table.aggregate`).
8//!
9//! ```
10//! use rich_ext::table::{Aggregate, GroupBy, Value};
11//!
12//! let rows = vec![
13//!     vec![Value::from("api"), Value::Int(3)],
14//!     vec![Value::from("web"), Value::Int(5)],
15//!     vec![Value::from("api"), Value::Int(4)],
16//! ];
17//! let by = GroupBy::new(0).aggregate(Aggregate::sum(1));
18//! let groups = by.groups(&rows, &[0, 1, 2]);
19//! assert_eq!(groups[0].key.plain(), "api");
20//! assert_eq!(groups[0].rows, [0, 2]);
21//! assert_eq!(groups[0].aggregates[0], Value::Int(7));
22//! ```
23//!
24//! [`TableData`]: super::TableData
25
26use std::collections::HashMap;
27use std::fmt;
28use std::sync::Arc;
29
30use super::Value;
31use super::sort::{Compare, compare_values};
32
33/// A custom aggregate: the group's cells of one column in, a value out.
34pub type AggregateFn = Arc<dyn Fn(&[&Value]) -> Value + Send + Sync>;
35
36#[derive(Clone)]
37enum Kind {
38    Count,
39    Sum,
40    Min,
41    Max,
42    Mean,
43    Custom(AggregateFn),
44}
45
46/// A summary computed over one column of a group of rows.
47#[derive(Clone)]
48pub struct Aggregate {
49    column: usize,
50    kind: Kind,
51}
52
53impl fmt::Debug for Aggregate {
54    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
55        let kind = match self.kind {
56            Kind::Count => "count",
57            Kind::Sum => "sum",
58            Kind::Min => "min",
59            Kind::Max => "max",
60            Kind::Mean => "mean",
61            Kind::Custom(_) => "custom",
62        };
63        f.debug_struct("Aggregate")
64            .field("column", &self.column)
65            .field("kind", &kind)
66            .finish()
67    }
68}
69
70impl Aggregate {
71    fn new(column: usize, kind: Kind) -> Self {
72        Aggregate { column, kind }
73    }
74
75    /// The number of non-empty cells. Shown as a plain integer, not through
76    /// the column's formatter.
77    pub fn count(column: usize) -> Self {
78        Self::new(column, Kind::Count)
79    }
80
81    /// The sum of the numeric cells: an `Int` when every one is an integer
82    /// (and the sum fits), else a `Float`; `Null` when there are none.
83    pub fn sum(column: usize) -> Self {
84        Self::new(column, Kind::Sum)
85    }
86
87    /// The smallest non-empty cell, in natural order.
88    pub fn min(column: usize) -> Self {
89        Self::new(column, Kind::Min)
90    }
91
92    /// The largest non-empty cell, in natural order.
93    pub fn max(column: usize) -> Self {
94        Self::new(column, Kind::Max)
95    }
96
97    /// The mean of the numeric cells as a `Float`; `Null` when there are none.
98    pub fn mean(column: usize) -> Self {
99        Self::new(column, Kind::Mean)
100    }
101
102    /// Any summary: `f` receives the group's cells of `column`, in row order.
103    pub fn custom(column: usize, f: impl Fn(&[&Value]) -> Value + Send + Sync + 'static) -> Self {
104        Self::new(column, Kind::Custom(Arc::new(f)))
105    }
106
107    /// The column this aggregate summarises.
108    pub fn column(&self) -> usize {
109        self.column
110    }
111
112    /// Whether this is a [`count`](Aggregate::count), which is displayed
113    /// without the column's formatter.
114    pub fn is_count(&self) -> bool {
115        matches!(self.kind, Kind::Count)
116    }
117
118    /// Compute the aggregate over one column's cells.
119    pub fn compute(&self, values: &[&Value]) -> Value {
120        let non_empty = || values.iter().copied().filter(|v| !v.is_empty());
121        match &self.kind {
122            Kind::Count => Value::from(non_empty().count()),
123            Kind::Sum => {
124                let numbers: Vec<&Value> = values
125                    .iter()
126                    .copied()
127                    .filter(|v| v.as_f64().is_some())
128                    .collect();
129                if numbers.is_empty() {
130                    return Value::Null;
131                }
132                let ints: Option<i64> = numbers.iter().try_fold(0i64, |sum, v| match v {
133                    Value::Int(n) => sum.checked_add(*n),
134                    _ => None,
135                });
136                ints.map_or_else(
137                    || Value::Float(numbers.iter().filter_map(|v| v.as_f64()).sum()),
138                    Value::Int,
139                )
140            }
141            Kind::Min => non_empty()
142                .min_by(|a, b| compare_values(a, b, Compare::Natural))
143                .cloned()
144                .unwrap_or_default(),
145            Kind::Max => non_empty()
146                // `max_by` keeps the last of equal maxima; the first is wanted.
147                .reduce(|best, v| {
148                    if compare_values(v, best, Compare::Natural).is_gt() {
149                        v
150                    } else {
151                        best
152                    }
153                })
154                .cloned()
155                .unwrap_or_default(),
156            Kind::Mean => {
157                let numbers: Vec<f64> = values.iter().filter_map(|v| v.as_f64()).collect();
158                if numbers.is_empty() {
159                    Value::Null
160                } else {
161                    Value::Float(numbers.iter().sum::<f64>() / numbers.len() as f64)
162                }
163            }
164            Kind::Custom(f) => f(values),
165        }
166    }
167
168    /// Compute the aggregate over `rows` (indices into `data`).
169    pub fn over<R: AsRef<[Value]>>(&self, data: &[R], rows: &[usize]) -> Value {
170        const NULL: Value = Value::Null;
171        let cells: Vec<&Value> = rows
172            .iter()
173            .map(|&i| data[i].as_ref().get(self.column).unwrap_or(&NULL))
174            .collect();
175        self.compute(&cells)
176    }
177}
178
179/// Group rows by one column, with aggregates per group.
180#[derive(Clone, Debug)]
181pub struct GroupBy {
182    column: usize,
183    aggregates: Vec<Aggregate>,
184    label: String,
185}
186
187impl GroupBy {
188    /// Group by `column`, with no aggregates and the summary label `subtotal`.
189    pub fn new(column: usize) -> Self {
190        GroupBy {
191            column,
192            aggregates: Vec::new(),
193            label: "subtotal".to_string(),
194        }
195    }
196
197    /// Add an aggregate to each group's summary row.
198    #[must_use]
199    pub fn aggregate(mut self, aggregate: Aggregate) -> Self {
200        self.aggregates.push(aggregate);
201        self
202    }
203
204    /// The text in the summary row's first cell, followed by that column's
205    /// own aggregate if it has one (`subtotal: 2`). Empty for none.
206    #[must_use]
207    pub fn label(mut self, label: impl Into<String>) -> Self {
208        self.label = label.into();
209        self
210    }
211
212    /// The grouped column.
213    pub fn column(&self) -> usize {
214        self.column
215    }
216
217    /// The aggregates, in the order they were added.
218    pub fn aggregates(&self) -> &[Aggregate] {
219        &self.aggregates
220    }
221
222    /// The summary row's label.
223    pub fn summary_label(&self) -> &str {
224        &self.label
225    }
226
227    /// Group `rows`, visiting them in `order` (indices, e.g. from
228    /// [`sorted_indices`](super::sort::sorted_indices)). Groups are keyed by
229    /// the cell's plain display string (so `Int(1)` and `"1"` share a group)
230    /// and appear in order of their first row; empty cells form one group.
231    pub fn groups<R: AsRef<[Value]>>(&self, rows: &[R], order: &[usize]) -> Vec<Group> {
232        let mut groups: Vec<Group> = Vec::new();
233        let mut by_key: HashMap<String, usize> = HashMap::new();
234        for &row in order {
235            let key = rows[row]
236                .as_ref()
237                .get(self.column)
238                .cloned()
239                .unwrap_or_default();
240            let slot = *by_key.entry(key.plain()).or_insert_with(|| {
241                groups.push(Group {
242                    key,
243                    rows: Vec::new(),
244                    aggregates: Vec::new(),
245                });
246                groups.len() - 1
247            });
248            groups[slot].rows.push(row);
249        }
250        for group in &mut groups {
251            group.aggregates = self
252                .aggregates
253                .iter()
254                .map(|aggregate| aggregate.over(rows, &group.rows))
255                .collect();
256        }
257        groups
258    }
259}
260
261/// One group of rows.
262#[derive(Clone, Debug)]
263pub struct Group {
264    /// The grouped column's value (from the group's first row).
265    pub key: Value,
266    /// The group's rows, as indices in visiting order.
267    pub rows: Vec<usize>,
268    /// One value per [`GroupBy::aggregate`], in the same order.
269    pub aggregates: Vec<Value>,
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275
276    #[test]
277    fn aggregates_over_mixed_cells() {
278        let cells = [
279            Value::Int(3),
280            Value::Null,
281            Value::Float(1.5),
282            Value::from("x"),
283            Value::Int(-2),
284        ];
285        let refs: Vec<&Value> = cells.iter().collect();
286        assert_eq!(Aggregate::count(0).compute(&refs), Value::Int(4));
287        assert_eq!(Aggregate::sum(0).compute(&refs), Value::Float(2.5));
288        assert_eq!(Aggregate::min(0).compute(&refs), Value::Int(-2));
289        assert_eq!(Aggregate::max(0).compute(&refs), Value::from("x"));
290        assert_eq!(Aggregate::mean(0).compute(&refs), Value::Float(2.5 / 3.0));
291        assert_eq!(Aggregate::sum(0).compute(&[]), Value::Null);
292        let ints = [Value::Int(i64::MAX), Value::Int(1)];
293        let refs: Vec<&Value> = ints.iter().collect();
294        assert_eq!(
295            Aggregate::sum(0).compute(&refs),
296            Value::Float(i64::MAX as f64 + 1.0)
297        );
298    }
299}