mirror of
https://github.com/storytold/gridcraft.git
synced 2026-10-24 10:46:50 -06:00
Let GROUPBY and PIVOTBY take a LAMBDA as their function
Fixes #903 Co-Authored-By: Claude Haiku 5.5 <noreply@anthropic.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Haiku 5.5
parent
e6b6b34970
commit
cdc8177fae
+33
-5
@@ -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)),
|
||||
}
|
||||
@@ -1318,7 +1346,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;
|
||||
|
||||
@@ -611,6 +611,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();
|
||||
|
||||
@@ -573,6 +573,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);
|
||||
@@ -732,10 +749,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);
|
||||
@@ -748,11 +765,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()
|
||||
};
|
||||
@@ -798,11 +815,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);
|
||||
@@ -818,12 +835,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));
|
||||
@@ -873,7 +890,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 => {
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user