diff --git a/crates/functions/src/dynamic.rs b/crates/functions/src/dynamic.rs index 3a66ca8..e016759 100644 --- a/crates/functions/src/dynamic.rs +++ b/crates/functions/src/dynamic.rs @@ -785,6 +785,12 @@ fn groupby(args: &[Arg], _c: &mut dyn Ctx) -> R { 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 { 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 { let ckeys: Vec> = 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> = 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 { 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::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 { } out.push(line); } - if vals.cols > 1 || p.show_headers { + if value_row { let mut line: Vec = 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 = (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 = (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); + } }