Check the full PIVOTBY layout size before building column headers (#537)

crates/functions/src/dynamic.rs

@@ -785,6 +785,12 @@ fn groupby(args: &[Arg], _c: &mut dyn Ctx) -> R<Value> {
     from_rows(out, &Value::Empty)
 }
 
+/// Height and width of a PIVOTBY grid: `header_rows` (one per column field, plus the value-name row) and
+/// `body_rows` (row keys plus the total row), each `lead` row-field cells wide plus `value_cols` cells per column key.
+fn pivot_dims(header_rows: usize, body_rows: usize, col_keys: usize, value_cols: usize, lead: usize) -> (usize, usize) {
+    (header_rows.saturating_add(body_rows), col_keys.saturating_mul(value_cols).saturating_add(lead))
+}
+
 fn pivotby(args: &[Arg], _c: &mut dyn Ctx) -> R<Value> {
     let rf = arr_arg(args, 0)?;
     let cf = arr_arg(args, 1)?;
@@ -801,9 +807,6 @@ fn pivotby(args: &[Arg], _c: &mut dyn Ctx) -> R<Value> {
     let ckeys: Vec<Vec<Value>> = p.rows.iter().map(|&r| key_row(&cf, r)).collect();
     let rk = distinct_keys(&rkeys, if row_sort.unsigned_abs() as usize <= rf.cols { row_sort } else { 0 });
     let ck = distinct_keys(&ckeys, if col_sort.unsigned_abs() as usize <= cf.cols { col_sort } else { 0 });
-    if rk.len().saturating_add(2).saturating_mul((ck.len() + 1).saturating_mul(vals.cols).saturating_add(rf.cols)) > MAX_CELLS {
-        return Err(CellError::Num);
-    }
     let mut cell: HashMap<(String, String), Vec<usize>> = HashMap::new();
     for ((r, c), &i) in rkeys.iter().zip(ckeys.iter()).zip(p.rows.iter()) {
         cell.entry((row_key(r), row_key(c))).or_default().push(i);
@@ -832,6 +835,14 @@ fn pivotby(args: &[Arg], _c: &mut dyn Ctx) -> R<Value> {
         cols.insert(0, None);
     }
     let lead = rf.cols;
+    let value_row = vals.cols > 1 || p.show_headers;
+    let header_rows = cf.cols.saturating_add(usize::from(value_row));
+    let body_rows = rk.len().saturating_add(usize::from(row_total != 0));
+    // Reject before building any header or body row: every column-field header row is as wide as the whole grid.
+    let (height, width) = pivot_dims(header_rows, body_rows, cols.len(), vals.cols, lead);
+    if height.saturating_mul(width) > MAX_CELLS {
+        return Err(CellError::Num);
+    }
     let mut out: Vec<Vec<Value>> = Vec::new();
     // Column header rows: one per column field, plus value names when several value columns.
     for level in 0..cf.cols {
@@ -847,7 +858,7 @@ fn pivotby(args: &[Arg], _c: &mut dyn Ctx) -> R<Value> {
         }
         out.push(line);
     }
-    if vals.cols > 1 || p.show_headers {
+    if value_row {
         let mut line: Vec<Value> = if p.show_headers { p.field_headers.iter().take(lead).cloned().collect() } else { vec![Value::Empty; lead] };
         line.resize(lead, Value::Empty);
         for _ in &cols {
@@ -1225,4 +1236,29 @@ mod tests {
             ]
         );
     }
+
+    #[test]
+    fn pivotby_size_check_counts_column_header_rows() {
+        // One row, n column fields and n value columns: n header rows + value-name row + body row + total row,
+        // each 1 + 2n wide. The size checked up front matches the grid PIVOTBY builds.
+        let dims = |k: u32, totals: bool| {
+            let fields: Vec<f64> = (0..k).map(f64::from).collect();
+            let mut args = vec![av(row(&[1.0])), av(row(&fields)), av(row(&fields)), t("SUM")];
+            if !totals {
+                args.extend([empty(), n(0.0), empty(), n(0.0)]);
+            }
+            let r = rows_of(&ev("PIVOTBY", args));
+            (r.len(), r.first().map_or(0, Vec::len))
+        };
+        assert_eq!(dims(2, true), super::pivot_dims(3, 2, 2, 2, 1));
+        assert_eq!(dims(10, true), (13, 21));
+        assert_eq!(dims(10, true), super::pivot_dims(11, 2, 2, 10, 1));
+        assert_eq!(dims(10, false), (12, 11));
+        assert_eq!(dims(10, false), super::pivot_dims(11, 1, 1, 10, 1));
+        // n = 2300: the 2300 header rows alone are 10.58M cells, so the check must count them before any is built.
+        let (h, w) = super::pivot_dims(2301, 2, 2, 2300, 1);
+        assert!(2300 * w > crate::util::MAX_CELLS && h * w > crate::util::MAX_CELLS);
+        let wide: Vec<f64> = (0..2300).map(f64::from).collect();
+        is_err(ev("PIVOTBY", vec![av(row(&[1.0])), av(row(&wide)), av(row(&wide)), t("SUM")]), CellError::Num);
+    }
 }