improv_engine/
scalar.rs

1//! A pure scalar evaluator: one formula, one value per input, no dimensions.
2//!
3//! This is the engine's per-cell semantics without the dataflow around it. It
4//! calls the very operators the dataflow builder maps over its collections
5//! (`apply_binary`, `apply_unary`, `apply_scalar`), so a formula evaluated
6//! here gives the bit-identical value the engine gives at a single coordinate.
7//! `pg_improv`'s `LANGUAGE improv` evaluates SQL function bodies with it.
8//!
9//! Semantics:
10//! * a `Ref` to a measure absent from `env` is a missing input: the result is
11//!   `None` wherever it depends on it (the dataflow emits no cell there);
12//! * division by zero and domain errors are `NaN` (the engine's convention);
13//! * aggregations and dimension projections (`SUM(x OVER T)`, `x BY T`) need a
14//!   model with categories and are rejected with [`ScalarError::NeedsDimensions`].
15
16use crate::compiler::{desugar_lazy, is_aggregation, lazy};
17use crate::dataflow::{apply_binary, apply_unary};
18use crate::CellValue;
19use improv_core_model::functions::{self, Kind};
20use improv_core_model::{Expr, FuncId, MeasureId};
21use std::collections::HashMap;
22
23/// Why a formula cannot be evaluated as a scalar.
24#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
25pub enum ScalarError {
26    /// An aggregation or a dimension projection: these range over categories,
27    /// and a scalar evaluation has none.
28    #[error(
29        "{0} needs dimensions; aggregate over categories in an Improv model, not a scalar formula"
30    )]
31    NeedsDimensions(String),
32    /// A function id that is neither a scalar built-in nor an aggregation.
33    #[error("unknown function {0:?}")]
34    UnknownFunction(FuncId),
35    /// A built-in called with the wrong number of arguments.
36    #[error("{want} takes other arguments than the {got} given")]
37    Arity {
38        func: FuncId,
39        want: String,
40        got: usize,
41    },
42}
43
44/// Evaluate `expr` with each referenced measure's value taken from `env`.
45///
46/// `Ok(None)` means a referenced input is missing (SQL NULL in, NULL out).
47pub fn eval(
48    expr: &Expr,
49    env: &HashMap<MeasureId, CellValue>,
50) -> Result<Option<CellValue>, ScalarError> {
51    check(expr)?;
52    Ok(eval_checked(expr, env))
53}
54
55/// Reject what a scalar evaluation cannot do, before evaluating anything, so
56/// the error does not depend on which inputs happen to be NULL.
57pub fn check(expr: &Expr) -> Result<(), ScalarError> {
58    match expr {
59        Expr::Literal(_) => Ok(()),
60        Expr::Category(_) => Err(ScalarError::NeedsDimensions("a category argument".into())),
61        Expr::Ref(_, spec) => {
62            if spec.by.is_empty() && spec.over.is_empty() && spec.except.is_empty() {
63                Ok(())
64            } else {
65                Err(ScalarError::NeedsDimensions(
66                    "a dimension projection (BY/OVER/EXCEPT)".into(),
67                ))
68            }
69        }
70        Expr::UnaryOp(_, a) => check(a),
71        Expr::BinaryOp(_, a, b) => check(a).and_then(|()| check(b)),
72        Expr::Call(func, args) => {
73            if is_aggregation(*func) {
74                return Err(ScalarError::NeedsDimensions(
75                    "an aggregation (SUM(x OVER …))".into(),
76                ));
77            }
78            let def = functions::by_id(*func).ok_or(ScalarError::UnknownFunction(*func))?;
79            if def.category == functions::Category::Dimensional {
80                return Err(ScalarError::NeedsDimensions(format!(
81                    "{} (it reads a category)",
82                    def.name
83                )));
84            }
85            if !def.accepts(args.len()) {
86                return Err(ScalarError::Arity {
87                    func: *func,
88                    want: def.signature(),
89                    got: args.len(),
90                });
91            }
92            args.iter().try_for_each(check)
93        }
94    }
95}
96
97fn eval_checked(expr: &Expr, env: &HashMap<MeasureId, CellValue>) -> Option<CellValue> {
98    match expr {
99        Expr::Literal(v) => CellValue::from_model_value(v),
100        Expr::Category(_) => None,
101        Expr::Ref(m, _) => env.get(m).cloned(),
102        Expr::UnaryOp(op, a) => Some(apply_unary(*op, &eval_checked(a, env)?)),
103        Expr::BinaryOp(op, a, b) => {
104            let (a, b) = (eval_checked(a, env)?, eval_checked(b, env)?);
105            Some(apply_binary(*op, &a, &b))
106        }
107        Expr::Call(func, args) if func.0 == lazy::IF => {
108            // Lazy: only the selected branch is evaluated (an error in the
109            // other never surfaces); no else-branch means no value.
110            match eval_checked(&args[0], env)? {
111                CellValue::Bool(true) => eval_checked(&args[1], env),
112                CellValue::Bool(false) => eval_checked(args.get(2)?, env),
113                other => Some(other),
114            }
115        }
116        // COALESCE: the first argument with a value.
117        Expr::Call(func, args) if func.0 == lazy::COALESCE => {
118            args.iter().find_map(|a| eval_checked(a, env))
119        }
120        Expr::Call(func, args) if functions::by_id(*func).is_some_and(|d| d.kind == Kind::Lazy) => {
121            eval_checked(&desugar_lazy(*func, args).ok()?, env)
122        }
123        Expr::Call(func, args) => {
124            let vals = args
125                .iter()
126                .map(|a| eval_checked(a, env))
127                .collect::<Option<Vec<CellValue>>>()?;
128            Some(crate::builtins::eval(*func, &vals))
129        }
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136    use improv_core_model::{
137        BinaryOp, CategoryId, DimensionSpec, Formula, Measure, MeasureKind, Model, Name, Value,
138        ValueType,
139    };
140    use proptest::prelude::*;
141
142    const A: MeasureId = MeasureId(1);
143    const B: MeasureId = MeasureId(2);
144
145    fn r(m: MeasureId) -> Expr {
146        Expr::Ref(m, DimensionSpec::default())
147    }
148    fn n(x: f64) -> Expr {
149        Expr::Literal(Value::Number(x))
150    }
151    fn t(s: &str) -> Expr {
152        Expr::Literal(Value::Text(s.into()))
153    }
154    fn bin(op: BinaryOp, a: Expr, b: Expr) -> Expr {
155        Expr::BinaryOp(op, Box::new(a), Box::new(b))
156    }
157    fn call(f: u32, args: Vec<Expr>) -> Expr {
158        Expr::Call(FuncId(f), args)
159    }
160    fn env(a: f64, b: f64) -> HashMap<MeasureId, CellValue> {
161        HashMap::from([(A, CellValue::num(a)), (B, CellValue::num(b))])
162    }
163    fn num(e: &Expr, a: f64, b: f64) -> f64 {
164        eval(e, &env(a, b)).unwrap().unwrap().as_num().unwrap()
165    }
166    fn boolean(e: &Expr) -> bool {
167        match eval(e, &env(2.0, 3.0)).unwrap().unwrap() {
168            CellValue::Bool(b) => b,
169            other => panic!("not a bool: {other:?}"),
170        }
171    }
172
173    #[test]
174    fn arithmetic() {
175        assert_eq!(num(&bin(BinaryOp::Add, r(A), r(B)), 2.0, 3.0), 5.0);
176        assert_eq!(num(&bin(BinaryOp::Sub, r(A), r(B)), 2.0, 3.0), -1.0);
177        assert_eq!(num(&bin(BinaryOp::Mul, r(A), r(B)), 2.0, 3.0), 6.0);
178        assert_eq!(num(&bin(BinaryOp::Div, r(A), r(B)), 3.0, 2.0), 1.5);
179        let neg = Expr::UnaryOp(improv_core_model::UnaryOp::Neg, Box::new(r(A)));
180        assert_eq!(num(&neg, 2.0, 0.0), -2.0);
181    }
182
183    fn value(e: &Expr, a: f64, b: f64) -> CellValue {
184        eval(e, &env(a, b)).unwrap().unwrap()
185    }
186
187    #[test]
188    fn division_by_zero_is_an_error() {
189        // Mirrors the dataflow: x / 0 is #DIV/0! (including 0 / 0).
190        let div0 = CellValue::Err(crate::builtins::err::DIV0);
191        assert_eq!(value(&bin(BinaryOp::Div, r(A), r(B)), 1.0, 0.0), div0);
192        assert_eq!(value(&bin(BinaryOp::Div, r(A), r(B)), 0.0, 0.0), div0);
193        assert_eq!(value(&bin(BinaryOp::Div, n(-1.0), n(0.0)), 0.0, 0.0), div0);
194    }
195
196    /// IF evaluates only the branch it selects; the IF family desugars to it.
197    #[test]
198    fn if_is_lazy() {
199        let div = bin(BinaryOp::Div, r(A), r(B)); // #DIV/0! when B is 0
200        let pos = bin(BinaryOp::Gt, r(B), n(0.0));
201        let iff = |args| call(30, args);
202        assert_eq!(
203            value(&iff(vec![pos.clone(), div.clone(), n(-1.0)]), 6.0, 2.0),
204            CellValue::num(3.0)
205        );
206        assert_eq!(
207            value(&iff(vec![pos.clone(), div.clone(), n(-1.0)]), 6.0, 0.0),
208            CellValue::num(-1.0)
209        );
210        assert_eq!(
211            eval(&iff(vec![pos.clone(), div.clone()]), &env(6.0, 0.0)),
212            Ok(None)
213        );
214        // An error test is the result; a missing selected branch is no cell.
215        let bad = bin(BinaryOp::Gt, div.clone(), n(0.0));
216        assert_eq!(
217            value(&iff(vec![bad, n(1.0), n(2.0)]), 1.0, 0.0),
218            CellValue::Err(crate::builtins::err::DIV0)
219        );
220        let only_a = HashMap::from([(A, CellValue::num(1.0))]);
221        assert_eq!(
222            eval(
223                &iff(vec![bin(BinaryOp::Gt, r(A), n(0.0)), r(B), n(0.0)]),
224                &only_a
225            ),
226            Ok(None)
227        );
228        // IFERROR and friends.
229        assert_eq!(
230            value(&call(34, vec![div.clone(), n(9.0)]), 1.0, 0.0),
231            CellValue::num(9.0)
232        );
233        assert_eq!(
234            value(&call(34, vec![div.clone(), n(9.0)]), 4.0, 2.0),
235            CellValue::num(2.0)
236        );
237        assert_eq!(eval(&call(36, vec![div.clone()]), &env(1.0, 0.0)), Ok(None));
238        assert_eq!(
239            value(&call(37, vec![r(A), n(5.0)]), 0.0, 0.0),
240            CellValue::num(5.0)
241        );
242        assert_eq!(
243            value(&call(33, vec![r(A), t("x"), t("y")]), 2.9, 0.0),
244            CellValue::Text("y".into())
245        );
246        assert_eq!(
247            value(&call(33, vec![r(A), t("x")]), 2.0, 0.0),
248            CellValue::Err(crate::builtins::err::VALUE)
249        );
250        assert_eq!(
251            value(
252                &call(
253                    31,
254                    vec![
255                        bin(BinaryOp::Gt, r(A), n(5.0)),
256                        t("big"),
257                        bin(BinaryOp::Gt, r(A), n(1.0)),
258                        t("mid")
259                    ]
260                ),
261                2.0,
262                0.0
263            ),
264            CellValue::Text("mid".into())
265        );
266        assert_eq!(
267            value(
268                &call(31, vec![bin(BinaryOp::Gt, r(A), n(5.0)), t("big")]),
269                2.0,
270                0.0
271            ),
272            CellValue::Err(crate::builtins::err::NA)
273        );
274        assert_eq!(
275            value(
276                &call(
277                    32,
278                    vec![r(A), n(1.0), t("one"), n(2.0), t("two"), t("many")]
279                ),
280                2.0,
281                0.0
282            ),
283            CellValue::Text("two".into())
284        );
285        assert_eq!(
286            value(&call(32, vec![r(A), n(1.0), t("one"), t("many")]), 7.0, 0.0),
287            CellValue::Text("many".into())
288        );
289        assert_eq!(
290            value(&call(35, vec![call(53, vec![]), n(0.0)]), 0.0, 0.0),
291            CellValue::num(0.0)
292        );
293    }
294
295    #[test]
296    fn comparisons_and_logic() {
297        assert!(boolean(&bin(BinaryOp::Lt, r(A), r(B))));
298        assert!(boolean(&bin(BinaryOp::Le, r(A), r(A))));
299        assert!(!boolean(&bin(BinaryOp::Gt, r(A), r(B))));
300        assert!(boolean(&bin(BinaryOp::Ge, r(B), r(A))));
301        assert!(boolean(&bin(BinaryOp::Eq, r(A), n(2.0))));
302        assert!(boolean(&bin(BinaryOp::Ne, r(A), r(B))));
303        assert!(boolean(&bin(BinaryOp::And, r(A), r(B))));
304        assert!(!boolean(&bin(BinaryOp::And, r(A), n(0.0))));
305        assert!(boolean(&bin(BinaryOp::Or, n(0.0), r(B))));
306        let not = Expr::UnaryOp(improv_core_model::UnaryOp::Not, Box::new(n(0.0)));
307        assert!(boolean(&not));
308    }
309
310    #[test]
311    fn text_equality() {
312        let e = HashMap::from([(A, CellValue::Text("x".into()))]);
313        let eq = |rhs: &str| eval(&bin(BinaryOp::Eq, r(A), t(rhs)), &e).unwrap().unwrap();
314        assert_eq!(eq("x"), CellValue::Bool(true));
315        assert_eq!(eq("y"), CellValue::Bool(false));
316    }
317
318    #[test]
319    fn scalar_builtins() {
320        for (id, x, want) in [
321            (10, -2.5, 2.5),
322            (11, 2.5, 3.0),
323            (12, 2.5, 2.0),
324            (13, 2.1, 3.0),
325            (14, 9.0, 3.0),
326            (15, 4.0, -4.0),
327        ] {
328            assert_eq!(num(&call(id, vec![r(A)]), x, 0.0), want, "func {id}");
329        }
330        assert_eq!(
331            value(&call(14, vec![r(A)]), -1.0, 0.0),
332            CellValue::Err(crate::builtins::err::NUM)
333        );
334        assert_eq!(num(&call(20, vec![r(A), r(B)]), 1.0, 2.0), 1.0);
335        assert_eq!(num(&call(21, vec![r(A), r(B)]), 1.0, 2.0), 2.0);
336    }
337
338    #[test]
339    fn missing_input_propagates() {
340        let only_a = HashMap::from([(A, CellValue::num(1.0))]);
341        for e in [
342            bin(BinaryOp::Add, r(A), r(B)),
343            call(21, vec![r(A), r(B)]),
344            Expr::UnaryOp(improv_core_model::UnaryOp::Neg, Box::new(r(B))),
345            bin(BinaryOp::Or, n(1.0), r(B)), // no short-circuit, like the dataflow
346        ] {
347            assert_eq!(eval(&e, &only_a), Ok(None), "{e:?}");
348        }
349        assert_eq!(
350            eval(&bin(BinaryOp::Add, r(A), n(1.0)), &only_a),
351            Ok(Some(CellValue::num(2.0)))
352        );
353    }
354
355    #[test]
356    fn rejects_what_needs_dimensions() {
357        let none = HashMap::new();
358        for f in 1..=4 {
359            assert!(matches!(
360                eval(&call(f, vec![r(A)]), &none),
361                Err(ScalarError::NeedsDimensions(_))
362            ));
363        }
364        let cats = vec![CategoryId(1)];
365        for spec in [
366            DimensionSpec {
367                by: cats.clone(),
368                ..DimensionSpec::default()
369            },
370            DimensionSpec {
371                over: cats.clone(),
372                ..DimensionSpec::default()
373            },
374            DimensionSpec {
375                except: cats.clone(),
376                ..DimensionSpec::default()
377            },
378        ] {
379            assert!(matches!(
380                eval(&Expr::Ref(A, spec), &none),
381                Err(ScalarError::NeedsDimensions(_))
382            ));
383        }
384        // Rejected even when the inputs are missing (no NULL-dependent errors).
385        let nested = bin(BinaryOp::Add, r(B), call(1, vec![r(A)]));
386        assert!(eval(&nested, &none).is_err());
387        assert_eq!(
388            eval(&call(99, vec![]), &none),
389            Err(ScalarError::UnknownFunction(FuncId(99)))
390        );
391        assert!(matches!(
392            eval(&call(312, vec![r(A)]), &none),
393            Err(ScalarError::Arity { got: 1, .. })
394        ));
395    }
396
397    /// Run `expr` through the dataflow engine on a 0-dimension model with
398    /// inputs A and B (absent when `None`).
399    fn via_dataflow(expr: &Expr, a: Option<f64>, b: Option<f64>) -> Option<CellValue> {
400        let mut m = Model::new();
401        let measure = |id, name: &str, kind| Measure {
402            id,
403            name: Name(name.into()),
404            value_type: ValueType::Number,
405            categories: vec![],
406            kind,
407            description: None,
408        };
409        m.add_measure(measure(A, "A", MeasureKind::Input));
410        m.add_measure(measure(B, "B", MeasureKind::Input));
411        let target = MeasureId(3);
412        let derived = MeasureKind::Derived(Formula::new(expr.clone()));
413        m.add_measure(measure(target, "R", derived));
414        for (id, v) in [(A, a), (B, b)] {
415            if let Some(v) = v {
416                m.set_input(id, Default::default(), Value::Number(v));
417            }
418        }
419        let out = crate::dataflow::evaluate(&m, &[target]).expect("dataflow evaluate");
420        out.get(&target)
421            .and_then(|cells| cells.get(&Vec::new()).cloned())
422    }
423
424    fn num_expr() -> impl Strategy<Value = Expr> {
425        let leaf = prop_oneof![
426            Just(r(A)),
427            Just(r(B)),
428            (-3i8..=3).prop_map(|x| n(f64::from(x) / 2.0)),
429        ];
430        leaf.prop_recursive(3, 12, 2, |inner| {
431            let op = prop_oneof![
432                Just(BinaryOp::Add),
433                Just(BinaryOp::Sub),
434                Just(BinaryOp::Mul),
435                Just(BinaryOp::Div),
436                Just(BinaryOp::Pow),
437            ];
438            // Every one-number and two-number scalar in the registry.
439            let one: Vec<u32> = improv_core_model::functions::all()
440                .filter(|d| {
441                    d.kind == Kind::Scalar
442                        && d.required == 1
443                        && d.param(0) == functions::Arg::Number
444                        && matches!(
445                            d.ret,
446                            functions::Ret::Fixed(improv_core_model::ValueType::Number)
447                        )
448                })
449                .map(|d| d.id.0)
450                .collect();
451            let two: Vec<u32> = improv_core_model::functions::all()
452                .filter(|d| {
453                    d.kind == Kind::Scalar
454                        && d.accepts(2)
455                        && d.param(0) == functions::Arg::Number
456                        && d.param(1) == functions::Arg::Number
457                        && matches!(
458                            d.ret,
459                            functions::Ret::Fixed(improv_core_model::ValueType::Number)
460                        )
461                })
462                .map(|d| d.id.0)
463                .collect();
464            let cmp = (inner.clone(), inner.clone()).prop_map(|(a, b)| bin(BinaryOp::Gt, a, b));
465            prop_oneof![
466                (op, inner.clone(), inner.clone()).prop_map(|(op, a, b)| bin(op, a, b)),
467                (prop::sample::select(one), inner.clone()).prop_map(|(f, a)| call(f, vec![a])),
468                (prop::sample::select(two), inner.clone(), inner.clone())
469                    .prop_map(|(f, a, b)| call(f, vec![a, b])),
470                (cmp.clone(), inner.clone(), inner.clone())
471                    .prop_map(|(c, a, b)| call(30, vec![c, a, b])),
472                (cmp, inner.clone()).prop_map(|(c, a)| call(30, vec![c, a])),
473                (inner.clone(), inner).prop_map(|(a, b)| call(34, vec![a, b])),
474            ]
475        })
476    }
477
478    fn any_expr() -> impl Strategy<Value = Expr> {
479        let cmp = prop_oneof![
480            Just(BinaryOp::Lt),
481            Just(BinaryOp::Le),
482            Just(BinaryOp::Gt),
483            Just(BinaryOp::Ge),
484            Just(BinaryOp::Eq),
485            Just(BinaryOp::Ne),
486        ];
487        prop_oneof![
488            num_expr(),
489            (cmp, num_expr(), num_expr()).prop_map(|(op, a, b)| bin(op, a, b)),
490        ]
491    }
492
493    proptest! {
494        #![proptest_config(ProptestConfig::with_cases(256))]
495        /// The scalar evaluator and the dataflow engine agree, bit for bit,
496        /// on every formula over two scalar inputs (including missing ones).
497        #[test]
498        fn matches_dataflow(
499            expr in any_expr(),
500            a in proptest::option::weighted(0.9, -4i8..=4),
501            b in proptest::option::weighted(0.9, -4i8..=4),
502        ) {
503            let (a, b) = (a.map(f64::from), b.map(f64::from));
504            let mut env = HashMap::new();
505            for (id, v) in [(A, a), (B, b)] {
506                if let Some(v) = v {
507                    env.insert(id, CellValue::num(v));
508                }
509            }
510            let scalar = eval(&expr, &env).expect("scalar eval");
511            prop_assert_eq!(scalar, via_dataflow(&expr, a, b), "{:?}", expr);
512        }
513    }
514}
improv_engine/scalar.rs

improv_engine/
scalar.rs
