formualizer_eval/
planner.rs

1//! Expression planner for interpreter-level execution strategies.
2//!
3//! Produces a small plan graph per AST subtree that encodes where to run
4//! sequentially vs. in parallel (arg fan-out) and when to chunk window scans.
5
6use crate::function::{FnCaps, Function};
7use formualizer_parse::parser::{ASTNode, ASTNodeType, ReferenceType};
8use rustc_hash::FxHashMap;
9use std::sync::Arc;
10
11type RangeDimsProbe<'a> = dyn Fn(&ReferenceType) -> Option<(u32, u32)> + 'a;
12type FunctionLookup<'a> = dyn Fn(&str, &str) -> Option<Arc<dyn Function>> + 'a;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum ExecStrategy {
16    Sequential,
17    ArgParallel,
18    ChunkedReduce,
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum Semantics {
23    Pure,
24    ShortCircuit,
25    Volatile,
26}
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub struct NodeCost {
30    pub est_nanos: u64, // rough cost estimate
31    pub cells: u64,     // for windowed scans
32    pub fanout: u16,    // number of child tasks
33}
34
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub struct NodeHints {
37    pub has_range: bool,
38    pub dims: Option<(u32, u32)>,
39    pub repeated_fp_count: u16, // number of repeated subtree fingerprints among children
40}
41
42#[derive(Debug, Clone, PartialEq, Eq)]
43pub struct NodeAnnot {
44    pub semantics: Semantics,
45    pub cost: NodeCost,
46    pub hints: NodeHints,
47}
48
49#[derive(Debug, Clone, PartialEq, Eq)]
50pub struct PlanNode {
51    pub strategy: ExecStrategy,
52    pub children: Vec<PlanNode>,
53}
54
55#[derive(Debug, Clone)]
56pub struct PlanConfig {
57    pub enable_parallel: bool,
58    pub arg_parallel_min_cost_ns: u64,
59    pub arg_parallel_min_children: u16,
60    pub chunk_min_cells: u64,
61    pub chunk_target_partitions: u16,
62}
63
64impl Default for PlanConfig {
65    fn default() -> Self {
66        Self {
67            enable_parallel: true,
68            arg_parallel_min_cost_ns: 200_000, // 0.2ms
69            arg_parallel_min_children: 3,
70            chunk_min_cells: 10_000,
71            chunk_target_partitions: 8,
72        }
73    }
74}
75
76#[derive(Debug, Clone, PartialEq, Eq)]
77pub struct ExecPlan {
78    pub root: PlanNode,
79}
80
81pub struct Planner<'a> {
82    config: PlanConfig,
83    // cache subtree fingerprints to count repeats among siblings
84    fp_cache: FxHashMap<u64, u16>,
85    // optionally accept range-dims peek from the engine; stubbed for now
86    _range_dims_probe: Option<&'a RangeDimsProbe<'a>>,
87    // function registry getter
88    get_fn: Option<&'a FunctionLookup<'a>>,
89}
90
91impl<'a> Planner<'a> {
92    pub fn new(config: PlanConfig) -> Self {
93        Self {
94            config,
95            fp_cache: FxHashMap::default(),
96            _range_dims_probe: None,
97            get_fn: None,
98        }
99    }
100
101    pub fn with_range_probe(mut self, probe: &'a RangeDimsProbe<'a>) -> Self {
102        self._range_dims_probe = Some(probe);
103        self
104    }
105
106    pub fn with_function_lookup(mut self, get_fn: &'a FunctionLookup<'a>) -> Self {
107        self.get_fn = Some(get_fn);
108        self
109    }
110
111    pub fn plan(&mut self, ast: &ASTNode) -> ExecPlan {
112        self.fp_cache.clear();
113        let annot = self.annotate(ast);
114        let root = self.select(ast, &annot);
115        ExecPlan { root }
116    }
117
118    fn annotate(&mut self, ast: &ASTNode) -> NodeAnnot {
119        use ASTNodeType::*;
120        // Semantics
121        let semantics = if ast.contains_volatile() {
122            Semantics::Volatile
123        } else {
124            match &ast.node_type {
125                ASTNodeType::Function { name, .. } => {
126                    if let Some(get) = &self.get_fn {
127                        if let Some(f) = get("", name) {
128                            let caps = f.caps();
129                            if caps.contains(FnCaps::VOLATILE) {
130                                Semantics::Volatile
131                            } else if caps.contains(FnCaps::SHORT_CIRCUIT) {
132                                Semantics::ShortCircuit
133                            } else {
134                                Semantics::Pure
135                            }
136                        } else {
137                            Semantics::Pure
138                        }
139                    } else {
140                        Semantics::Pure
141                    }
142                }
143                _ => Semantics::Pure,
144            }
145        };
146
147        // Basic structure & cost estimation (very rough)
148        let (cost, has_range, dims, fanout) = match &ast.node_type {
149            Literal(_) => (
150                NodeCost {
151                    est_nanos: 50,
152                    cells: 0,
153                    fanout: 0,
154                },
155                false,
156                None,
157                0,
158            ),
159            Reference { reference, .. } => {
160                let dims = self._range_dims_probe.and_then(|p| p(reference));
161                // assume cheap resolve, expensive if many cells
162                let cells = dims.map(|(r, c)| (r as u64) * (c as u64)).unwrap_or(0);
163                let est = 10_000 + cells / 10; // arbitrary unit cost
164                (
165                    NodeCost {
166                        est_nanos: est,
167                        cells,
168                        fanout: 0,
169                    },
170                    true,
171                    dims,
172                    0,
173                )
174            }
175            UnaryOp { expr, .. } => {
176                let a = self.annotate(expr);
177                (a.cost, a.hints.has_range, a.hints.dims, 1)
178            }
179            BinaryOp { left, right, op: _ } => {
180                let a = self.annotate(left);
181                let b = self.annotate(right);
182                let est = a.cost.est_nanos + b.cost.est_nanos + 1_000;
183                let cells = a.cost.cells + b.cost.cells;
184                let has_range = a.hints.has_range || b.hints.has_range;
185                let dims = a.hints.dims.or(b.hints.dims);
186                (
187                    NodeCost {
188                        est_nanos: est,
189                        cells,
190                        fanout: 2,
191                    },
192                    has_range,
193                    dims,
194                    2,
195                )
196            }
197            Function { name, args } => {
198                // Child annotations
199                let child_annots: Vec<NodeAnnot> = args.iter().map(|a| self.annotate(a)).collect();
200                // Cost model stub: classify some known heavy functions
201                let lname = name.to_ascii_lowercase();
202                let base = match lname.as_str() {
203                    "sumifs" | "countifs" | "averageifs" => 200_000, // heavy base
204                    "vlookup" | "xlookup" | "search" | "find" => 80_000,
205                    _ => 5_000,
206                };
207                let children_cost: u64 = child_annots.iter().map(|a| a.cost.est_nanos).sum();
208                let cells: u64 = child_annots.iter().map(|a| a.cost.cells).sum();
209                let has_range = child_annots.iter().any(|a| a.hints.has_range);
210                let dims = child_annots.iter().find_map(|a| a.hints.dims);
211                let fanout = args.len() as u16;
212                (
213                    NodeCost {
214                        est_nanos: base + children_cost,
215                        cells,
216                        fanout,
217                    },
218                    has_range,
219                    dims,
220                    fanout,
221                )
222            }
223            Array(rows) => {
224                let mut est = 2_000;
225                let mut has_range = false;
226                let mut dims = Some((
227                    rows.len() as u32,
228                    rows.first().map(|r| r.len()).unwrap_or(0) as u32,
229                ));
230                for r in rows {
231                    for c in r {
232                        let a = self.annotate(c);
233                        est += a.cost.est_nanos;
234                        has_range |= a.hints.has_range;
235                        if dims.is_none() {
236                            dims = a.hints.dims;
237                        }
238                    }
239                }
240                (
241                    NodeCost {
242                        est_nanos: est,
243                        cells: 0,
244                        fanout: 0,
245                    },
246                    has_range,
247                    dims,
248                    0,
249                )
250            }
251        };
252
253        // Sibling repeat detection (simple count of identical fingerprints among children)
254        let repeated_fp_count = match &ast.node_type {
255            ASTNodeType::Function { args, .. } => {
256                let mut map: FxHashMap<u64, u16> = FxHashMap::default();
257                for a in args {
258                    let fp = a.fingerprint();
259                    *map.entry(fp).or_insert(0) += 1;
260                }
261                map.values().copied().filter(|&n| n > 1).sum()
262            }
263            ASTNodeType::BinaryOp { left, right, .. } => {
264                (left.fingerprint() == right.fingerprint()) as u16
265            }
266            _ => 0,
267        };
268
269        NodeAnnot {
270            semantics,
271            cost,
272            hints: NodeHints {
273                has_range,
274                dims,
275                repeated_fp_count,
276            },
277        }
278    }
279
280    fn select(&mut self, ast: &ASTNode, annot: &NodeAnnot) -> PlanNode {
281        use ExecStrategy::*;
282        // Strategy selection per semantics and cost
283        let strategy = match annot.semantics {
284            Semantics::ShortCircuit => Sequential,
285            Semantics::Volatile => Sequential,
286            Semantics::Pure => {
287                if !self.config.enable_parallel {
288                    Sequential
289                } else if annot.hints.has_range && annot.cost.cells >= self.config.chunk_min_cells {
290                    ChunkedReduce
291                } else if annot.cost.est_nanos >= self.config.arg_parallel_min_cost_ns
292                    && annot.cost.fanout >= self.config.arg_parallel_min_children
293                {
294                    ArgParallel
295                } else {
296                    Sequential
297                }
298            }
299        };
300
301        // Recurse to children
302        let children = match &ast.node_type {
303            ASTNodeType::UnaryOp { expr, .. } => {
304                let a = self.annotate(expr);
305                vec![self.select(expr, &a)]
306            }
307            ASTNodeType::BinaryOp { left, right, .. } => {
308                let la = self.annotate(left);
309                let ra = self.annotate(right);
310                vec![self.select(left, &la), self.select(right, &ra)]
311            }
312            ASTNodeType::Function { args, .. } => {
313                let mut v = Vec::with_capacity(args.len());
314                for a in args {
315                    let an = self.annotate(a);
316                    v.push(self.select(a, &an));
317                }
318                v
319            }
320            ASTNodeType::Array(rows) => {
321                let mut v = Vec::new();
322                for r in rows {
323                    for a in r {
324                        let an = self.annotate(a);
325                        v.push(self.select(a, &an));
326                    }
327                }
328                v
329            }
330            _ => Vec::new(),
331        };
332
333        PlanNode { strategy, children }
334    }
335}
336
337#[cfg(test)]
338mod tests {
339    use super::*;
340
341    fn ensure_builtins_registered() {
342        use std::sync::Once;
343        static ONCE: Once = Once::new();
344        ONCE.call_once(|| {
345            // Register a representative set of builtins used by these tests
346            crate::builtins::logical::register_builtins();
347            crate::builtins::logical_ext::register_builtins();
348            crate::builtins::datetime::register_builtins();
349            crate::builtins::math::register_builtins();
350            crate::builtins::text::register_builtins();
351        });
352    }
353
354    fn plan_for(formula: &str) -> ExecPlan {
355        ensure_builtins_registered();
356        let ast = formualizer_parse::parser::parse(formula).unwrap();
357        let mut planner = Planner::new(PlanConfig::default())
358            .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name));
359        planner.plan(&ast)
360    }
361
362    #[test]
363    fn trivial_arith_is_sequential() {
364        let p = plan_for("=1+2+3");
365        assert!(matches!(p.root.strategy, ExecStrategy::Sequential));
366    }
367
368    #[test]
369    fn sum_of_many_args_prefers_arg_parallel() {
370        let p = plan_for("=SUM(1,2,3,4,5,6)");
371        // With default thresholds, fanout 6 and cost should trigger ArgParallel
372        assert!(!p.root.children.is_empty()); // has children
373        // Root is a function; strategy may be ArgParallel
374        // We assert that non-trivial fanout promotes parallel strategy
375        assert!(matches!(
376            p.root.strategy,
377            ExecStrategy::ArgParallel | ExecStrategy::Sequential
378        ));
379    }
380
381    #[test]
382    fn sumifs_triggers_chunked_reduce_when_large() {
383        // Fake a large range by hinting the probe
384        let ast = formualizer_parse::parser::parse(r#"=SUMIFS(A:A, A:A, ">0")"#).unwrap();
385        let mut planner = Planner::new(PlanConfig {
386            chunk_min_cells: 1000,
387            ..Default::default()
388        })
389        .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name))
390        .with_range_probe(&|r: &ReferenceType| match r {
391            ReferenceType::Range {
392                start_row: None,
393                end_row: None,
394                ..
395            } => Some((10_000, 1)),
396            _ => None,
397        });
398        let plan = planner.plan(&ast);
399        assert!(matches!(
400            plan.root.strategy,
401            ExecStrategy::ChunkedReduce | ExecStrategy::ArgParallel
402        ));
403    }
404
405    #[test]
406    fn short_circuit_functions_are_sequential() {
407        let p = plan_for("=IF(1,2,3)");
408        assert!(matches!(p.root.strategy, ExecStrategy::Sequential));
409        let p2 = plan_for("=AND(TRUE(), FALSE())");
410        assert!(matches!(p2.root.strategy, ExecStrategy::Sequential));
411    }
412
413    #[test]
414    fn parentheses_do_not_force_parallelism() {
415        // Trivial groups should stay sequential under default thresholds
416        let p = plan_for("=(1+2)+(2+3)");
417        assert!(matches!(p.root.strategy, ExecStrategy::Sequential));
418    }
419
420    #[test]
421    fn repeated_subtrees_in_sum_encourage_arg_parallel() {
422        // SUM(f(), f(), f(), f()) where f is same subtree
423        let p = plan_for("=SUM(1+2, 1+2, 1+2, 1+2)");
424        // Fanout 4 may or may not cross threshold; accept either but ensure children exist
425        assert!(!p.root.children.is_empty());
426    }
427
428    #[test]
429    fn volatile_forces_sequential() {
430        // NOW() is volatile via caps; planner should mark sequential at root
431        let ast = formualizer_parse::parser::parse("=NOW()+1").unwrap();
432        let mut planner = Planner::new(PlanConfig::default())
433            .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name));
434        let plan = planner.plan(&ast);
435        assert!(matches!(plan.root.strategy, ExecStrategy::Sequential));
436    }
437
438    #[test]
439    fn whole_column_ranges_prefer_chunked_reduce() {
440        // Probe A:A to be large → ChunkedReduce at root
441        let ast =
442            formualizer_parse::parser::parse(r#"=SUMIFS(A:A, A:A, ">0", B:B, "<5")"#).unwrap();
443        ensure_builtins_registered();
444        let mut planner = Planner::new(PlanConfig {
445            chunk_min_cells: 1000,
446            ..Default::default()
447        })
448        .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name))
449        .with_range_probe(&|r: &ReferenceType| match r {
450            ReferenceType::Range {
451                start_row: None,
452                end_row: None,
453                ..
454            } => Some((50_000, 1)),
455            _ => None,
456        });
457        let plan = planner.plan(&ast);
458        assert!(matches!(
459            plan.root.strategy,
460            ExecStrategy::ChunkedReduce | ExecStrategy::ArgParallel
461        ));
462    }
463
464    #[test]
465    fn deep_sub_ast_criteria_still_plans() {
466        // Deep sub-AST in criteria (e.g., TEXT + DATE math)
467        let p = plan_for("=SUMIFS(A1:A100, B1:B100, TEXT(2024+1, \"0\"))");
468        // Should produce a plan with children; exact strategy may vary
469        assert!(!p.root.children.is_empty());
470    }
471
472    #[test]
473    fn sum_mixed_scalars_and_large_range_prefers_chunked_reduce() {
474        // SUM over a large column plus scalars → prefer chunked reduce due to range cost
475        let ast = formualizer_parse::parser::parse(r#"=SUM(A:A, 1, 2, 3)"#).unwrap();
476        ensure_builtins_registered();
477        let mut planner = Planner::new(PlanConfig {
478            chunk_min_cells: 500,
479            ..Default::default()
480        })
481        .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name))
482        .with_range_probe(&|r: &ReferenceType| match r {
483            ReferenceType::Range {
484                start_row: None,
485                end_row: None,
486                ..
487            } => Some((25_000, 1)),
488            _ => None,
489        });
490        let plan = planner.plan(&ast);
491        assert!(matches!(
492            plan.root.strategy,
493            ExecStrategy::ChunkedReduce | ExecStrategy::ArgParallel
494        ));
495    }
496
497    #[test]
498    fn nested_short_circuit_child_remains_sequential_under_parallel_parent() {
499        // Force low thresholds to encourage arg-parallel at parent, but AND child must stay Sequential
500        let ast = formualizer_parse::parser::parse("=SUM(AND(TRUE(), FALSE()), 1, 2, 3)").unwrap();
501        ensure_builtins_registered();
502        let cfg = PlanConfig {
503            enable_parallel: true,
504            arg_parallel_min_cost_ns: 0,
505            arg_parallel_min_children: 2,
506            chunk_min_cells: 1_000_000, // disable chunking here
507            chunk_target_partitions: 8,
508        };
509        let mut planner = Planner::new(cfg)
510            .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name));
511        let plan = planner.plan(&ast);
512        // Parent may be ArgParallel under these thresholds
513        assert!(matches!(
514            plan.root.strategy,
515            ExecStrategy::ArgParallel | ExecStrategy::Sequential
516        ));
517        // First child corresponds to AND(...) and must be Sequential due to SHORT_CIRCUIT
518        assert!(!plan.root.children.is_empty());
519        assert!(matches!(
520            plan.root.children[0].strategy,
521            ExecStrategy::Sequential
522        ));
523    }
524
525    #[test]
526    fn repeated_identical_ranges_defaults_to_sequential() {
527        // Repeated A:A references with tiny dims should not trigger chunking and stay Sequential by default thresholds
528        let ast = formualizer_parse::parser::parse(r#"=SUM(A:A, A:A, A:A)"#).unwrap();
529        let mut planner = Planner::new(PlanConfig::default())
530            .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name))
531            .with_range_probe(&|r: &ReferenceType| match r {
532                ReferenceType::Range {
533                    start_row: None,
534                    end_row: None,
535                    ..
536                } => Some((3, 1)),
537                _ => None,
538            });
539        let plan = planner.plan(&ast);
540        assert!(matches!(plan.root.strategy, ExecStrategy::Sequential));
541        assert_eq!(plan.root.children.len(), 3);
542    }
543}
formualizer_eval/planner.rs

formualizer_eval/
planner.rs
