diff --git a/crates/engine/src/cmd/whatif/consolidate.rs b/crates/engine/src/cmd/whatif/consolidate.rs index 149b72bc7..702baa312 100644 --- a/crates/engine/src/cmd/whatif/consolidate.rs +++ b/crates/engine/src/cmd/whatif/consolidate.rs @@ -61,6 +61,24 @@ enum Key { const MAX_CONSOLIDATE_CELLS: u64 = 5_000_000; +/// The cells a consolidated table writes: every row key meets every column key, plus a label for +/// each key on the sides that have labels. +fn output_cells(rows: usize, cols: usize, top: bool, left: bool) -> u64 { + let (r, c) = (rows as u64, cols as u64); + r.saturating_mul(c).saturating_add(if top { c } else { 0 }).saturating_add(if left { r } else { 0 }) +} + +#[cfg(test)] +#[test] +fn the_output_limit_counts_label_cells() { + // A million row keys and four column keys are 4 million cells, 5 million with the row labels. + assert_eq!(output_cells(1_000_000, 4, false, false), 4_000_000); + assert_eq!(output_cells(1_000_000, 4, false, true), 5_000_000); + assert!(output_cells(1_000_000, 5, false, true) > MAX_CONSOLIDATE_CELLS); + assert_eq!(output_cells(3, 2, true, true), 6 + 2 + 3); + assert_eq!(output_cells(usize::MAX, usize::MAX, true, true), u64::MAX); +} + pub(super) fn consolidate(s: &mut Session, p: &Json) -> Result { const CMD: &str = "data.consolidate"; let func = str_param(p, "function").unwrap_or("sum").to_string(); @@ -89,20 +107,20 @@ pub(super) fn consolidate(s: &mut Session, p: &Json) -> Result { } None => (d.wb.active_sheet(), d.selection.active), }; + // Keys in first-seen order, each with its index: a lookup is one probe, not a scan. let mut row_keys: Vec = Vec::new(); let mut col_keys: Vec = Vec::new(); + let mut row_index: HashMap = HashMap::new(); + let mut col_index: HashMap = HashMap::new(); let mut row_label: HashMap = HashMap::new(); let mut col_label: HashMap = HashMap::new(); let mut aggs: HashMap<(usize, usize), Agg> = HashMap::new(); let mut total: u64 = 0; - let index_of = |keys: &mut Vec, k: Key| -> usize { - match keys.iter().position(|x| *x == k) { - Some(i) => i, - None => { - keys.push(k); - keys.len() - 1 - } - } + let index_of = |keys: &mut Vec, index: &mut HashMap, k: Key| -> usize { + *index.entry(k).or_insert_with_key(|k| { + keys.push(k.clone()); + keys.len() - 1 + }) }; for t in &srcs { let (si, r) = parse_range(&d.wb, t, CMD)?; @@ -131,7 +149,7 @@ pub(super) fn consolidate(s: &mut Session, p: &Json) -> Result { } else { Key::Pos(col - c0) }; - cols.push((col, index_of(&mut col_keys, k))); + cols.push((col, index_of(&mut col_keys, &mut col_index, k))); } for row in r0..=r.end.row { let k = if left { @@ -145,7 +163,7 @@ pub(super) fn consolidate(s: &mut Session, p: &Json) -> Result { } else { Key::Pos(row - r0) }; - let ri = index_of(&mut row_keys, k); + let ri = index_of(&mut row_keys, &mut row_index, k); for (col, ci) in &cols { aggs.entry((ri, *ci)).or_default().add(&sh.value(CellRef::new(row, *col))); } @@ -157,6 +175,9 @@ pub(super) fn consolidate(s: &mut Session, p: &Json) -> Result { { return Err(EngineError::Other("The consolidated data doesn't fit below and right of the destination.".into())); } + if output_cells(row_keys.len(), col_keys.len(), top, left) > MAX_CONSOLIDATE_CELLS { + return Err(bad(CMD, "the consolidated table would be too large")); + } let mut cells: Vec<(CellRef, Value)> = Vec::new(); if top { for (ci, k) in col_keys.iter().enumerate() { diff --git a/crates/engine/src/cmd/whatif/tests.rs b/crates/engine/src/cmd/whatif/tests.rs index c1476e401..5e5da0e4c 100644 --- a/crates/engine/src/cmd/whatif/tests.rs +++ b/crates/engine/src/cmd/whatif/tests.rs @@ -297,3 +297,16 @@ fn consolidate_by_position_and_label() { assert!(s.execute("data.consolidate", json!({"sources": []})).is_err()); assert!(s.execute("data.consolidate", json!({"sources": ["A1"], "function": "median"})).is_err()); } + +/// A hundred thousand positional rows take one key lookup each; a tall source and a wide one +/// whose keys would meet in a table of 160 million cells are refused before it is built. +#[test] +fn consolidate_keys_are_indexed_and_the_output_is_bounded() { + let mut s = s(); + s.execute("range.setValues", json!({"range": "A1", "values": [[1], [2]]})).unwrap(); + let r = s.execute("data.consolidate", json!({"sources": ["A1:A100000"], "destination": "C1"})).unwrap(); + assert_eq!(r["rows"], 100_000); + assert_eq!(n(&s, "Sheet1!C2"), 2.0); + let r = s.execute("data.consolidate", json!({"sources": ["A1:A10000", "A1:XFC1"], "destination": "A1"})); + assert!(r.unwrap_err().to_string().contains("too large")); +}