sqlly_datatable/pivot/
engine.rs

1//! Pure pivot computation: source rows in, [`PivotResult`] out.
2//!
3//! No GPUI types and **no mutation of the source data** — the engine borrows
4//! the rows, buckets them into hierarchical row/column group trees, and
5//! streams every value through [`Accumulator`]s at every rollup level:
6//! leaf × leaf intersections, group subtotals, axis totals, and the grand
7//! total. Aggregating once per level (instead of aggregating aggregates)
8//! keeps `Avg`/`Count` correct at every level.
9
10use crate::config::ResolvedColumnFormat;
11use crate::data::{compare_cells, CellValue, Column};
12use crate::format::format_cell;
13use crate::pivot::aggregation::Accumulator;
14use crate::pivot::config::PivotConfig;
15use std::collections::HashMap;
16
17/// Sentinel node key meaning "the total across this whole axis". Used as a
18/// key into [`PivotResult::values`] alongside real node ids.
19pub const TOTAL_KEY: usize = usize::MAX;
20
21/// One group node on the row or column axis.
22#[derive(Clone, Debug, PartialEq)]
23pub struct PivotNode {
24    /// Formatted grouping value (or the configured blank label).
25    pub label: String,
26    /// The raw grouping value, kept for ordering with
27    /// [`crate::data::compare_cells`].
28    pub sort_key: CellValue,
29    /// 0 = outermost field on this axis.
30    pub depth: usize,
31    /// Parent node id; `None` for roots.
32    pub parent: Option<usize>,
33    /// Child node ids (empty at the innermost depth).
34    pub children: Vec<usize>,
35    /// This node's subtotal across the entire opposite axis.
36    pub total: CellValue,
37}
38
39impl PivotNode {
40    /// `true` when this node is at the innermost depth of its axis.
41    #[must_use]
42    pub fn is_leaf(&self) -> bool {
43        self.children.is_empty()
44    }
45}
46
47/// The complete computed pivot. Cheap to share behind an `Arc`; the paint
48/// path never touches the source rows again.
49#[derive(Clone, Debug, Default)]
50pub struct PivotResult {
51    /// Arena of row-axis nodes; tree edges via `parent`/`children`.
52    pub row_nodes: Vec<PivotNode>,
53    /// Arena of column-axis nodes.
54    pub col_nodes: Vec<PivotNode>,
55    /// Row-axis roots (depth 0), in canonical ascending order.
56    pub row_roots: Vec<usize>,
57    /// Column-axis roots (depth 0), in canonical ascending order.
58    pub col_roots: Vec<usize>,
59    /// Aggregated value for every `(row key, col key)` pair where a key is a
60    /// node id or [`TOTAL_KEY`]. Contains every rollup level, so collapsed
61    /// groups and subtotal rows/columns read straight from this map.
62    pub values: HashMap<(usize, usize), CellValue>,
63    /// The value across all rows and columns.
64    pub grand_total: CellValue,
65    /// Number of row fields (row tree depth).
66    pub row_depth: usize,
67    /// Number of column fields (column tree depth).
68    pub col_depth: usize,
69    /// Display names of the row fields, outermost first.
70    pub row_field_names: Vec<String>,
71    /// Display names of the column fields, outermost first.
72    pub col_field_names: Vec<String>,
73    /// Header caption for the value area, e.g. `"Sum of Amount"`.
74    pub value_caption: String,
75    /// How many source rows were included (after source filters).
76    pub source_row_count: usize,
77}
78
79impl PivotResult {
80    /// Aggregated value for `(row_key, col_key)` where either key may be
81    /// [`TOTAL_KEY`]. Missing intersections (no source rows) are `None`.
82    #[must_use]
83    pub fn value(&self, row_key: usize, col_key: usize) -> &CellValue {
84        self.values
85            .get(&(row_key, col_key))
86            .unwrap_or(&CellValue::None)
87    }
88
89    /// Ids of row-axis leaves in canonical depth-first order.
90    #[must_use]
91    pub fn row_leaves(&self) -> Vec<usize> {
92        collect_leaves(&self.row_nodes, &self.row_roots)
93    }
94
95    /// Ids of column-axis leaves in canonical depth-first order.
96    #[must_use]
97    pub fn col_leaves(&self) -> Vec<usize> {
98        collect_leaves(&self.col_nodes, &self.col_roots)
99    }
100}
101
102fn collect_leaves(nodes: &[PivotNode], roots: &[usize]) -> Vec<usize> {
103    let mut out = Vec::new();
104    let mut stack: Vec<usize> = roots.iter().rev().copied().collect();
105    while let Some(id) = stack.pop() {
106        let node = &nodes[id];
107        if node.is_leaf() {
108            out.push(id);
109        } else {
110            for &child in node.children.iter().rev() {
111                stack.push(child);
112            }
113        }
114    }
115    out
116}
117
118/// Incrementally builds one axis' group tree while source rows stream by.
119struct AxisBuilder {
120    nodes: Vec<PivotNode>,
121    roots: Vec<usize>,
122    /// `(parent id or TOTAL_KEY for roots, label) -> node id`.
123    index: HashMap<(usize, String), usize>,
124}
125
126impl AxisBuilder {
127    fn new() -> Self {
128        Self {
129            nodes: Vec::new(),
130            roots: Vec::new(),
131            index: HashMap::new(),
132        }
133    }
134
135    /// Resolve the node path for one source row down this axis. Returns the
136    /// node id at every depth (outermost first). Empty when the axis has no
137    /// fields.
138    fn path_for_row(
139        &mut self,
140        row: &[CellValue],
141        fields: &[usize],
142        formats: &[ResolvedColumnFormat],
143        blank_label: &str,
144    ) -> Vec<usize> {
145        let mut path = Vec::with_capacity(fields.len());
146        let mut parent_key = TOTAL_KEY;
147        for (depth, &field) in fields.iter().enumerate() {
148            let cell = row.get(field).unwrap_or(&CellValue::None);
149            let label = if matches!(cell, CellValue::None) {
150                blank_label.to_owned()
151            } else {
152                format_cell(cell, &formats[field]).0
153            };
154            let id = match self.index.get(&(parent_key, label.clone())) {
155                Some(&id) => id,
156                None => {
157                    let id = self.nodes.len();
158                    self.nodes.push(PivotNode {
159                        label: label.clone(),
160                        sort_key: cell.clone(),
161                        depth,
162                        parent: (parent_key != TOTAL_KEY).then_some(parent_key),
163                        children: Vec::new(),
164                        total: CellValue::None,
165                    });
166                    self.index.insert((parent_key, label), id);
167                    if parent_key == TOTAL_KEY {
168                        self.roots.push(id);
169                    } else {
170                        self.nodes[parent_key].children.push(id);
171                    }
172                    id
173                }
174            };
175            path.push(id);
176            parent_key = id;
177        }
178        path
179    }
180
181    /// Sort every sibling list by the raw grouping value (ascending) so the
182    /// initial presentation is deterministic and natural.
183    fn sort_canonical(&mut self) {
184        let keys: Vec<CellValue> = self.nodes.iter().map(|n| n.sort_key.clone()).collect();
185        let by_key = |a: &usize, b: &usize| compare_cells(&keys[*a], &keys[*b]);
186        self.roots.sort_by(by_key);
187        for node in &mut self.nodes {
188            node.children.sort_by(by_key);
189        }
190    }
191}
192
193/// Compute a pivot over `rows[source_rows...]`.
194///
195/// * `columns` / `formats` describe the source schema (formats drive group
196///   labels, so date grouping follows the column's display format).
197/// * `source_rows` selects which rows participate — the caller applies any
198///   source filters first. Pass `0..rows.len()` for everything.
199///
200/// The source slices are only read; nothing is cloned except the small
201/// grouping values that become node labels/sort keys.
202#[must_use]
203pub fn compute_pivot(
204    columns: &[Column],
205    rows: &[Vec<CellValue>],
206    source_rows: &[usize],
207    config: &PivotConfig,
208    formats: &[ResolvedColumnFormat],
209) -> PivotResult {
210    let mut row_axis = AxisBuilder::new();
211    let mut col_axis = AxisBuilder::new();
212    let mut accs: HashMap<(usize, usize), Accumulator> = HashMap::new();
213
214    let value_field = config.value_field;
215    let agg = config.aggregation;
216
217    for &row_idx in source_rows {
218        let Some(row) = rows.get(row_idx) else {
219            continue;
220        };
221        let row_path = row_axis.path_for_row(row, &config.row_fields, formats, &config.blank_label);
222        let col_path =
223            col_axis.path_for_row(row, &config.column_fields, formats, &config.blank_label);
224        let Some(vf) = value_field else {
225            continue;
226        };
227        let value = row.get(vf).unwrap_or(&CellValue::None);
228        // Ingest at every rollup level: each ancestor (incl. leaf) on the row
229        // path plus the axis total, crossed with the same on the column path.
230        for &rk in row_path.iter().chain(std::iter::once(&TOTAL_KEY)) {
231            for &ck in col_path.iter().chain(std::iter::once(&TOTAL_KEY)) {
232                accs.entry((rk, ck))
233                    .or_insert_with(|| Accumulator::new(agg))
234                    .ingest(value);
235            }
236        }
237    }
238
239    row_axis.sort_canonical();
240    col_axis.sort_canonical();
241
242    let values: HashMap<(usize, usize), CellValue> =
243        accs.iter().map(|(k, acc)| (*k, acc.finish())).collect();
244
245    let mut row_nodes = row_axis.nodes;
246    let mut col_nodes = col_axis.nodes;
247    for (id, node) in row_nodes.iter_mut().enumerate() {
248        node.total = values
249            .get(&(id, TOTAL_KEY))
250            .cloned()
251            .unwrap_or(CellValue::None);
252    }
253    for (id, node) in col_nodes.iter_mut().enumerate() {
254        node.total = values
255            .get(&(TOTAL_KEY, id))
256            .cloned()
257            .unwrap_or(CellValue::None);
258    }
259    let grand_total = values
260        .get(&(TOTAL_KEY, TOTAL_KEY))
261        .cloned()
262        .unwrap_or(CellValue::None);
263
264    let field_name = |idx: usize| {
265        columns
266            .get(idx)
267            .map(|c| c.name.clone())
268            .unwrap_or_else(|| format!("column {idx}"))
269    };
270    let value_caption = match value_field {
271        Some(vf) => agg.caption(&field_name(vf)),
272        None => "Values".to_owned(),
273    };
274
275    PivotResult {
276        row_roots: row_axis.roots,
277        col_roots: col_axis.roots,
278        row_nodes,
279        col_nodes,
280        values,
281        grand_total,
282        row_depth: config.row_fields.len(),
283        col_depth: config.column_fields.len(),
284        row_field_names: config.row_fields.iter().map(|&f| field_name(f)).collect(),
285        col_field_names: config
286            .column_fields
287            .iter()
288            .map(|&f| field_name(f))
289            .collect(),
290        value_caption,
291        source_row_count: source_rows.len(),
292    }
293}
294
295#[cfg(test)]
296#[allow(clippy::unwrap_used)]
297mod tests {
298    use super::*;
299    use crate::config::GridConfig;
300    use crate::data::ColumnKind;
301    use crate::pivot::aggregation::AggregationFn;
302    use CellValue::{Decimal, Integer, None as Null, Text};
303
304    /// region / product / year / amount ledger used across tests.
305    fn fixture() -> (Vec<Column>, Vec<Vec<CellValue>>, Vec<ResolvedColumnFormat>) {
306        let columns = vec![
307            Column::new("region", ColumnKind::Text, 100.0),
308            Column::new("product", ColumnKind::Text, 100.0),
309            Column::new("year", ColumnKind::Integer, 80.0),
310            Column::new("amount", ColumnKind::Decimal, 100.0),
311        ];
312        let r = |region: &str, product: &str, year: i64, amount: f64| {
313            vec![
314                Text(region.into()),
315                Text(product.into()),
316                Integer(year),
317                Decimal(amount),
318            ]
319        };
320        let rows = vec![
321            r("Europe", "Widget", 2023, 10.0),
322            r("Europe", "Widget", 2024, 20.0),
323            r("Europe", "Gadget", 2023, 5.0),
324            r("Asia", "Widget", 2023, 7.0),
325            r("Asia", "Gadget", 2024, 3.0),
326            r("Asia", "Gadget", 2024, 4.0),
327        ];
328        let formats = GridConfig::default().resolve_all(&columns);
329        (columns, rows, formats)
330    }
331
332    fn all_rows(rows: &[Vec<CellValue>]) -> Vec<usize> {
333        (0..rows.len()).collect()
334    }
335
336    fn config(rows: &[usize], cols: &[usize], value: usize, agg: AggregationFn) -> PivotConfig {
337        PivotConfig {
338            row_fields: rows.to_vec(),
339            column_fields: cols.to_vec(),
340            value_field: Some(value),
341            aggregation: agg,
342            ..PivotConfig::default()
343        }
344    }
345
346    fn node_by_label<'a>(result: &'a PivotResult, label: &str) -> (usize, &'a PivotNode) {
347        result
348            .row_nodes
349            .iter()
350            .enumerate()
351            .find(|(_, n)| n.label == label)
352            .unwrap()
353    }
354
355    #[test]
356    fn single_row_and_column_field_sum() {
357        let (columns, rows, formats) = fixture();
358        let cfg = config(&[0], &[2], 3, AggregationFn::Sum);
359        let result = compute_pivot(&columns, &rows, &all_rows(&rows), &cfg, &formats);
360
361        assert_eq!(result.row_roots.len(), 2); // Asia, Europe
362        assert_eq!(result.col_roots.len(), 2); // 2023, 2024
363                                               // Canonical order is ascending by raw value.
364        assert_eq!(result.row_nodes[result.row_roots[0]].label, "Asia");
365        assert_eq!(result.row_nodes[result.row_roots[1]].label, "Europe");
366
367        let (europe, europe_node) = node_by_label(&result, "Europe");
368        let y2023 = result.col_roots[0];
369        let y2024 = result.col_roots[1];
370        assert_eq!(result.value(europe, y2023), &Decimal(15.0));
371        assert_eq!(result.value(europe, y2024), &Decimal(20.0));
372        assert_eq!(europe_node.total, Decimal(35.0));
373
374        let (asia, asia_node) = node_by_label(&result, "Asia");
375        assert_eq!(result.value(asia, y2023), &Decimal(7.0));
376        assert_eq!(result.value(asia, y2024), &Decimal(7.0));
377        assert_eq!(asia_node.total, Decimal(14.0));
378
379        // Column totals and grand total.
380        assert_eq!(result.col_nodes[y2023].total, Decimal(22.0));
381        assert_eq!(result.col_nodes[y2024].total, Decimal(27.0));
382        assert_eq!(result.grand_total, Decimal(49.0));
383    }
384
385    #[test]
386    fn two_row_fields_build_two_level_tree_with_subtotals() {
387        let (columns, rows, formats) = fixture();
388        let cfg = config(&[0, 1], &[], 3, AggregationFn::Sum);
389        let result = compute_pivot(&columns, &rows, &all_rows(&rows), &cfg, &formats);
390
391        assert_eq!(result.row_depth, 2);
392        let (europe, europe_node) = node_by_label(&result, "Europe");
393        assert_eq!(europe_node.depth, 0);
394        assert_eq!(europe_node.children.len(), 2); // Gadget, Widget
395        let child_labels: Vec<&str> = europe_node
396            .children
397            .iter()
398            .map(|&c| result.row_nodes[c].label.as_str())
399            .collect();
400        assert_eq!(child_labels, vec!["Gadget", "Widget"]);
401        // Group subtotal at the parent level.
402        assert_eq!(europe_node.total, Decimal(35.0));
403        // Leaf value: Europe->Widget across (no column fields => TOTAL col).
404        let widget = europe_node
405            .children
406            .iter()
407            .copied()
408            .find(|&c| result.row_nodes[c].label == "Widget")
409            .unwrap();
410        assert_eq!(result.value(widget, TOTAL_KEY), &Decimal(30.0));
411        assert_eq!(result.row_nodes[widget].parent, Some(europe));
412        // Leaves enumerate depth-first: Asia(Gadget,Widget), Europe(Gadget,Widget).
413        let leaves = result.row_leaves();
414        let leaf_labels: Vec<&str> = leaves
415            .iter()
416            .map(|&l| result.row_nodes[l].label.as_str())
417            .collect();
418        assert_eq!(leaf_labels, vec!["Gadget", "Widget", "Gadget", "Widget"]);
419    }
420
421    #[test]
422    fn count_and_avg_levels_are_computed_from_source_not_from_child_aggregates() {
423        let (columns, rows, formats) = fixture();
424        let cfg = config(&[0], &[], 3, AggregationFn::Avg);
425        let result = compute_pivot(&columns, &rows, &all_rows(&rows), &cfg, &formats);
426        let (_, europe) = node_by_label(&result, "Europe");
427        // Europe rows: 10, 20, 5 -> avg 35/3, NOT avg of per-product avgs.
428        match &europe.total {
429            Decimal(v) => assert!((v - 35.0 / 3.0).abs() < 1e-9),
430            other => panic!("expected Decimal, got {other:?}"),
431        }
432        match &result.grand_total {
433            Decimal(v) => assert!((v - 49.0 / 6.0).abs() < 1e-9),
434            other => panic!("expected Decimal, got {other:?}"),
435        }
436    }
437
438    #[test]
439    fn count_aggregation_reports_integers() {
440        let (columns, rows, formats) = fixture();
441        let cfg = config(&[0], &[2], 3, AggregationFn::Count);
442        let result = compute_pivot(&columns, &rows, &all_rows(&rows), &cfg, &formats);
443        let (asia, _) = node_by_label(&result, "Asia");
444        let y2024 = result.col_roots[1];
445        assert_eq!(result.value(asia, y2024), &Integer(2));
446        assert_eq!(result.grand_total, Integer(6));
447    }
448
449    #[test]
450    fn empty_source_produces_empty_result_without_panic() {
451        let (columns, _, formats) = fixture();
452        let rows: Vec<Vec<CellValue>> = vec![];
453        let cfg = config(&[0], &[2], 3, AggregationFn::Sum);
454        let result = compute_pivot(&columns, &rows, &[], &cfg, &formats);
455        assert!(result.row_roots.is_empty());
456        assert!(result.col_roots.is_empty());
457        assert_eq!(result.grand_total, Null);
458        assert_eq!(result.source_row_count, 0);
459    }
460
461    #[test]
462    fn null_grouping_values_bucket_under_blank_label() {
463        let (columns, mut rows, formats) = fixture();
464        rows.push(vec![
465            Null,
466            Text("Widget".into()),
467            Integer(2023),
468            Decimal(2.0),
469        ]);
470        rows.push(vec![
471            Null,
472            Text("Gadget".into()),
473            Integer(2023),
474            Decimal(3.0),
475        ]);
476        let cfg = config(&[0], &[], 3, AggregationFn::Sum);
477        let result = compute_pivot(&columns, &rows, &all_rows(&rows), &cfg, &formats);
478        let (blank, blank_node) = node_by_label(&result, "(blank)");
479        assert_eq!(blank_node.sort_key, Null);
480        assert_eq!(result.value(blank, TOTAL_KEY), &Decimal(5.0));
481        // Null sorts first, so (blank) is the first root.
482        assert_eq!(result.row_roots[0], blank);
483    }
484
485    #[test]
486    fn source_row_subset_excludes_filtered_rows() {
487        let (columns, rows, formats) = fixture();
488        let cfg = config(&[0], &[], 3, AggregationFn::Sum);
489        // Only the Asia rows (indices 3..6).
490        let result = compute_pivot(&columns, &rows, &[3, 4, 5], &cfg, &formats);
491        assert_eq!(result.row_roots.len(), 1);
492        assert_eq!(result.row_nodes[result.row_roots[0]].label, "Asia");
493        assert_eq!(result.grand_total, Decimal(14.0));
494        assert_eq!(result.source_row_count, 3);
495    }
496
497    #[test]
498    fn missing_intersections_read_as_none() {
499        let (columns, rows, formats) = fixture();
500        let cfg = config(&[0, 1], &[2], 3, AggregationFn::Sum);
501        let result = compute_pivot(&columns, &rows, &all_rows(&rows), &cfg, &formats);
502        // Asia -> Widget has no 2024 rows.
503        let (asia, asia_node) = node_by_label(&result, "Asia");
504        let widget = asia_node
505            .children
506            .iter()
507            .copied()
508            .find(|&c| result.row_nodes[c].label == "Widget")
509            .unwrap();
510        let y2024 = result.col_roots[1];
511        assert_eq!(result.value(widget, y2024), &Null);
512        assert_eq!(result.row_nodes[widget].parent, Some(asia));
513    }
514
515    #[test]
516    fn value_caption_reflects_aggregation_and_field() {
517        let (columns, rows, formats) = fixture();
518        let cfg = config(&[0], &[], 3, AggregationFn::Sum);
519        let result = compute_pivot(&columns, &rows, &all_rows(&rows), &cfg, &formats);
520        assert_eq!(result.value_caption, "Sum of amount");
521        assert_eq!(result.row_field_names, vec!["region".to_owned()]);
522    }
523
524    #[test]
525    fn source_rows_are_not_mutated() {
526        let (columns, rows, formats) = fixture();
527        let snapshot = rows.clone();
528        let cfg = config(&[0, 1], &[2], 3, AggregationFn::Avg);
529        let _ = compute_pivot(&columns, &rows, &all_rows(&rows), &cfg, &formats);
530        assert_eq!(rows, snapshot);
531    }
532
533    #[test]
534    fn boolean_and_date_grouping_fields_work() {
535        let columns = vec![
536            Column::new("flag", ColumnKind::Boolean, 80.0),
537            Column::new("when", ColumnKind::Date, 120.0),
538            Column::new("n", ColumnKind::Integer, 80.0),
539        ];
540        let rows = vec![
541            vec![CellValue::Boolean(true), CellValue::Date(0), Integer(1)],
542            vec![CellValue::Boolean(false), CellValue::Date(0), Integer(2)],
543            vec![
544                CellValue::Boolean(true),
545                CellValue::Date(86_400),
546                Integer(3),
547            ],
548        ];
549        let formats = GridConfig::default().resolve_all(&columns);
550        let cfg = config(&[0], &[1], 2, AggregationFn::Sum);
551        let result = compute_pivot(&columns, &rows, &all_rows(&rows), &cfg, &formats);
552        assert_eq!(result.row_roots.len(), 2);
553        assert_eq!(result.col_roots.len(), 2);
554        // Labels come from the resolved formats (default date is %Y-%m-%d).
555        assert_eq!(result.col_nodes[result.col_roots[0]].label, "1970-01-01");
556        let t = result
557            .row_roots
558            .iter()
559            .copied()
560            .find(|&r| result.row_nodes[r].label == "true")
561            .unwrap();
562        assert_eq!(result.row_nodes[t].total, Integer(4));
563    }
564}
sqlly_datatable/pivot/engine.rs

sqlly_datatable/pivot/
engine.rs
