Merge PR #1469: Let GROUPBY and PIVOTBY take a LAMBDA as their function

This commit is contained in:
aiongg
2026-10-22 02:34:45 +00:00
4 changed files with 100 additions and 16 deletions
+33 -5
View File
@@ -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<Area>,
/// The LAMBDA given as the function argument of GROUPBY/PIVOTBY.
lambda: Option<Arc<Lambda>>,
}
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<Value> {
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;
+30
View File
@@ -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();
+28 -11
View File
@@ -576,6 +576,23 @@ fn parse_agg(v: &Value) -> R<Agg> {
}
}
/// 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<Option<Agg>> {
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<Agg>, 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<Option<Array>> {
if has(args, i) { Ok(Some(arr_arg(args, i)?)) } else { Ok(None) }
}
fn groupby(args: &[Arg], _c: &mut dyn Ctx) -> R<Value> {
fn groupby(args: &[Arg], c: &mut dyn Ctx) -> R<Value> {
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<Value> {
}
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<Value> {
let mut agg_row = |rows: &[usize]| -> Vec<Value> {
(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<Value> {
fn pivotby(args: &[Arg], ctx: &mut dyn Ctx) -> R<Value> {
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<Value> {
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<Value>>, c: Option<&Vec<Value>>| -> Vec<usize> {
let rkey = r.map(|k| row_key(k));
@@ -876,7 +893,7 @@ fn pivotby(args: &[Arg], _c: &mut dyn Ctx) -> R<Value> {
}
out.push(line);
}
let body_line = |r: Option<&Vec<Value>>| -> Vec<Value> {
let mut body_line = |r: Option<&Vec<Value>>| -> Vec<Value> {
let mut line: Vec<Value> = match r {
Some(k) => k.clone(),
None => {
+9
View File
@@ -111,6 +111,15 @@ pub trait Ctx {
fn computed_criterion(&mut self, _row: usize, _col: usize, _db: &Array, _record: usize) -> Option<bool> {
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<Value> {
None
}
}
pub type FnImpl = fn(&[Arg], &mut dyn Ctx) -> Value;