diff --git a/crates/calc/src/eval.rs b/crates/calc/src/eval.rs index dd50d832c..9220e9702 100644 --- a/crates/calc/src/eval.rs +++ b/crates/calc/src/eval.rs @@ -90,6 +90,8 @@ struct FnCtx<'a, 'h> { ev: &'a mut Evaluator<'h>, /// The criteria argument of a database function, when it is a reference. criteria: Option, + /// The LAMBDA given as the function argument of GROUPBY/PIVOTBY. + lambda: Option>, } impl Ctx for FnCtx<'_, '_> { @@ -124,6 +126,19 @@ impl Ctx for FnCtx<'_, '_> { let mut ev = Evaluator { host: &mut *self.ev.host, sheet: area.sheet, at, env, depth: self.ev.depth }; Some(matches!(ev.value(&shifted).scalar().to_bool(), Ok(true))) } + fn has_lambda(&self) -> bool { + self.lambda.is_some() + } + /// A LAMBDA must give one value per group; an array result is #CALC!. + fn call_lambda(&mut self, arg: Value) -> Option { + let l = self.lambda.clone()?; + let [p] = l.params.as_slice() else { return Some(Value::Error(CellError::Value)) }; + let r = self.ev.call_lambda_values(&l, vec![(p.clone(), Ev::V(arg))]); + Some(match self.ev.deref(r) { + Value::Array(a) if a.data.len() != 1 => Value::Error(CellError::Calc), + v => v.scalar(), + }) + } } impl<'h> Evaluator<'h> { @@ -618,7 +633,7 @@ impl<'h> Evaluator<'h> { let value = self.value(a); evaluated.push(Arg { value, from_ref: false }); } - let mut ctx = FnCtx { ev: self, criteria: None }; + let mut ctx = FnCtx { ev: self, criteria: None, lambda: None }; return gridcraft_functions::call(spec, &evaluated, &mut ctx).pipe(Ev::V); } let i = match index.to_number() { @@ -817,8 +832,21 @@ impl<'h> Evaluator<'h> { let mut criteria_shape = None; // Database functions: the criteria range (argument 2), whose formulas are computed criteria. let mut criteria = None; + // GROUPBY/PIVOTBY: the LAMBDA given as their function. + let mut lambda = None; + let lambda_at = match name { + "GROUPBY" => Some(2), + "PIVOTBY" => Some(3), + _ => None, + }; for (i, a) in args.iter().enumerate() { let mut ev = self.eval(a); + if lambda_at == Some(i) + && let Ev::L(l) = &ev + { + lambda = Some(l.clone()); + ev = Ev::V(Value::Empty); + } if spec.category == gridcraft_functions::Category::Database && i == 2 && let Ev::R(areas) = &ev @@ -853,7 +881,7 @@ impl<'h> Evaluator<'h> { }; evaluated.push(Arg { value, from_ref }); } - let mut ctx = FnCtx { ev: self, criteria }; + let mut ctx = FnCtx { ev: self, criteria, lambda }; gridcraft_functions::call(spec, &evaluated, &mut ctx).pipe(Ev::V) } @@ -907,13 +935,13 @@ impl<'h> Evaluator<'h> { }; let args = [Arg { value, from_ref }]; if kind != Agg::Count { - match gridcraft_functions::call(sum, &args, &mut FnCtx { ev: self, criteria: None }) { + match gridcraft_functions::call(sum, &args, &mut FnCtx { ev: self, criteria: None, lambda: None }) { Value::Number(n) => acc.sum += n, other => return Err(Ev::V(other)), } } if kind != Agg::Sum { - match gridcraft_functions::call(count, &args, &mut FnCtx { ev: self, criteria: None }) { + match gridcraft_functions::call(count, &args, &mut FnCtx { ev: self, criteria: None, lambda: None }) { Value::Number(n) => acc.count = acc.count.saturating_add(n as u64), other => return Err(Ev::V(other)), } @@ -1329,7 +1357,7 @@ impl<'h> Evaluator<'h> { } } } - let mut ctx = FnCtx { ev: self, criteria: None }; + let mut ctx = FnCtx { ev: self, criteria: None, lambda: None }; return match gridcraft_functions::call(spec, &xargs, &mut ctx) { Value::Number(p) if p >= 1.0 => { let k = p as u32 - 1; diff --git a/crates/calc/src/tests.rs b/crates/calc/src/tests.rs index 376e9a7bc..f02197413 100644 --- a/crates/calc/src/tests.rs +++ b/crates/calc/src/tests.rs @@ -634,6 +634,36 @@ fn functions_given_by_name() { assert_eq!(t.num("J1"), 20.0); } +#[test] +fn groupby_and_pivotby_take_a_lambda() { + let mut t = T::new(); + for (i, (k, c, v)) in [("a", "x", 1), ("b", "x", 2), ("a", "y", 3)].iter().enumerate() { + t.set(&format!("A{}", i + 1), k); + t.set(&format!("B{}", i + 1), c); + t.set(&format!("C{}", i + 1), &v.to_string()); + } + t.set("E1", "=GROUPBY(A1:A3,C1:C3,LAMBDA(x,SUM(x)))"); + assert_eq!(t.get("E1"), Value::from("a")); + assert_eq!(t.num("F1"), 4.0); + assert_eq!(t.get("E2"), Value::from("b")); + assert_eq!(t.num("F2"), 2.0); + assert_eq!(t.get("E3"), Value::from("Total")); + assert_eq!(t.num("F3"), 6.0); + // A LET-bound lambda, called once per group with that group's values. + t.set("H1", "=LET(f,LAMBDA(v,ROWS(v)*10),GROUPBY(A1:A3,C1:C3,f))"); + assert_eq!(t.num("I1"), 20.0); + assert_eq!(t.num("I3"), 30.0); + // PIVOTBY: rows a, b, Total; columns x, y, Total. + t.set("K1", "=PIVOTBY(A1:A3,B1:B3,C1:C3,LAMBDA(x,SUM(x)))"); + assert_eq!(t.get("L1"), Value::from("x")); + assert_eq!(t.num("L2"), 1.0); + assert_eq!(t.num("M2"), 3.0); + assert_eq!(t.num("N4"), 6.0); + // Neither a function nor a lambda. + t.set("P1", "=GROUPBY(A1:A3,C1:C3,TRUE)"); + assert_eq!(t.get("P1"), Value::Error(CellError::Value)); +} + #[test] fn lambda_argument_count() { let mut t = T::new(); diff --git a/crates/functions/src/dynamic.rs b/crates/functions/src/dynamic.rs index 0e5c5aa17..ee3e10a3e 100644 --- a/crates/functions/src/dynamic.rs +++ b/crates/functions/src/dynamic.rs @@ -576,6 +576,23 @@ fn parse_agg(v: &Value) -> R { } } +/// The function argument of GROUPBY/PIVOTBY: a built-in aggregation, or `None` for a LAMBDA (called +/// through the context). +fn agg_arg(args: &[Arg], i: usize, c: &dyn Ctx) -> R> { + if c.has_lambda() { Ok(None) } else { parse_agg(&arg(args, i)?.value).map(Some) } +} + +/// Aggregates one group with `f`, or with the LAMBDA (given the group's values as a column) when `None`. +fn aggregate_with(f: Option, vals: &[&Value], c: &mut dyn Ctx) -> Value { + match f { + Some(f) => aggregate(f, vals), + None => { + let col = Value::from(Array::column(vals.iter().map(|v| (*v).clone()).collect())); + c.call_lambda(col).unwrap_or(Value::Error(CellError::Calc)) + } + } +} + fn aggregate(f: Agg, vals: &[&Value]) -> Value { if f == Agg::CountA { return Value::Number(vals.iter().filter(|v| !matches!(v, Value::Empty)).count() as f64); @@ -735,10 +752,10 @@ fn opt_arr(args: &[Arg], i: usize) -> R> { if has(args, i) { Ok(Some(arr_arg(args, i)?)) } else { Ok(None) } } -fn groupby(args: &[Arg], _c: &mut dyn Ctx) -> R { +fn groupby(args: &[Arg], c: &mut dyn Ctx) -> R { let rf = arr_arg(args, 0)?; let vals = arr_arg(args, 1)?; - let f = parse_agg(&arg(args, 2)?.value)?; + let f = agg_arg(args, 2, c)?; let fh = opt_count(args, 3)?; let total = opt_count(args, 4)?.unwrap_or(1); let sort = opt_count(args, 5)?.unwrap_or(1); @@ -751,11 +768,11 @@ fn groupby(args: &[Arg], _c: &mut dyn Ctx) -> R { } let sort_in_keys = if sort.unsigned_abs() as usize <= rf.cols { sort } else { 0 }; let mut ordered = distinct_keys(&keys, sort_in_keys); - let agg_row = |rows: &[usize]| -> Vec { + let mut agg_row = |rows: &[usize]| -> Vec { (0..vals.cols) - .map(|c| { - let vs: Vec<&Value> = rows.iter().filter_map(|&r| vals.get(r, c)).collect(); - aggregate(f, &vs) + .map(|col| { + let vs: Vec<&Value> = rows.iter().filter_map(|&r| vals.get(r, col)).collect(); + aggregate_with(f, &vs, c) }) .collect() }; @@ -801,11 +818,11 @@ fn pivot_dims(header_rows: usize, body_rows: usize, col_keys: usize, value_cols: (header_rows.saturating_add(body_rows), col_keys.saturating_mul(value_cols).saturating_add(lead)) } -fn pivotby(args: &[Arg], _c: &mut dyn Ctx) -> R { +fn pivotby(args: &[Arg], ctx: &mut dyn Ctx) -> R { let rf = arr_arg(args, 0)?; let cf = arr_arg(args, 1)?; let vals = arr_arg(args, 2)?; - let f = parse_agg(&arg(args, 3)?.value)?; + let f = agg_arg(args, 3, ctx)?; let fh = opt_count(args, 4)?; let row_total = opt_count(args, 5)?.unwrap_or(1); let row_sort = opt_count(args, 6)?.unwrap_or(1); @@ -821,12 +838,12 @@ fn pivotby(args: &[Arg], _c: &mut dyn Ctx) -> R { 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); } - let agg = |rows: &[usize], c: usize| -> Value { + let mut agg = |rows: &[usize], c: usize| -> Value { if rows.is_empty() { return Value::Empty; } let vs: Vec<&Value> = rows.iter().filter_map(|&r| vals.get(r, c)).collect(); - aggregate(f, &vs) + aggregate_with(f, &vs, ctx) }; let rows_for = |r: Option<&Vec>, c: Option<&Vec>| -> Vec { let rkey = r.map(|k| row_key(k)); @@ -876,7 +893,7 @@ fn pivotby(args: &[Arg], _c: &mut dyn Ctx) -> R { } out.push(line); } - let body_line = |r: Option<&Vec>| -> Vec { + let mut body_line = |r: Option<&Vec>| -> Vec { let mut line: Vec = match r { Some(k) => k.clone(), None => { diff --git a/crates/functions/src/lib.rs b/crates/functions/src/lib.rs index 148d68221..d7ccc8726 100644 --- a/crates/functions/src/lib.rs +++ b/crates/functions/src/lib.rs @@ -111,6 +111,15 @@ pub trait Ctx { fn computed_criterion(&mut self, _row: usize, _col: usize, _db: &Array, _record: usize) -> Option { None } + /// For GROUPBY/PIVOTBY: whether their function argument is a LAMBDA (its value is then not used). + fn has_lambda(&self) -> bool { + false + } + /// Calls that LAMBDA with `arg` (one group's values) and returns its single result; `None` when + /// there is no lambda. + fn call_lambda(&mut self, _arg: Value) -> Option { + None + } } pub type FnImpl = fn(&[Arg], &mut dyn Ctx) -> Value;