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    use formualizer_parse::Tokenizer;
341
342    fn ensure_builtins_registered() {
343        use std::sync::Once;
344        static ONCE: Once = Once::new();
345        ONCE.call_once(|| {
346            // Register a representative set of builtins used by these tests
347            crate::builtins::logical::register_builtins();
348            crate::builtins::logical_ext::register_builtins();
349            crate::builtins::datetime::register_builtins();
350            crate::builtins::math::register_builtins();
351            crate::builtins::text::register_builtins();
352        });
353    }
354
355    fn plan_for(formula: &str) -> ExecPlan {
356        ensure_builtins_registered();
357        let t = Tokenizer::new(formula).unwrap();
358        let mut parser = formualizer_parse::parser::Parser::new(t.items, false);
359        let ast = parser.parse().unwrap();
360        let mut planner = Planner::new(PlanConfig::default())
361            .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name));
362        planner.plan(&ast)
363    }
364
365    #[test]
366    fn trivial_arith_is_sequential() {
367        let p = plan_for("=1+2+3");
368        assert!(matches!(p.root.strategy, ExecStrategy::Sequential));
369    }
370
371    #[test]
372    fn sum_of_many_args_prefers_arg_parallel() {
373        let p = plan_for("=SUM(1,2,3,4,5,6)");
374        // With default thresholds, fanout 6 and cost should trigger ArgParallel
375        assert!(p.root.children.first().is_some()); // has children
376        // Root is a function; strategy may be ArgParallel
377        // We assert that non-trivial fanout promotes parallel strategy
378        assert!(matches!(
379            p.root.strategy,
380            ExecStrategy::ArgParallel | ExecStrategy::Sequential
381        ));
382    }
383
384    #[test]
385    fn sumifs_triggers_chunked_reduce_when_large() {
386        // Fake a large range by hinting the probe
387        let t = Tokenizer::new(r#"=SUMIFS(A:A, A:A, ">0")"#).unwrap();
388        let mut parser = formualizer_parse::parser::Parser::new(t.items, false);
389        let ast = parser.parse().unwrap();
390        let mut planner = Planner::new(PlanConfig {
391            chunk_min_cells: 1000,
392            ..Default::default()
393        })
394        .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name))
395        .with_range_probe(&|r: &ReferenceType| match r {
396            ReferenceType::Range {
397                start_row: None,
398                end_row: None,
399                ..
400            } => Some((10_000, 1)),
401            _ => None,
402        });
403        let plan = planner.plan(&ast);
404        assert!(matches!(
405            plan.root.strategy,
406            ExecStrategy::ChunkedReduce | ExecStrategy::ArgParallel
407        ));
408    }
409
410    #[test]
411    fn short_circuit_functions_are_sequential() {
412        let p = plan_for("=IF(1,2,3)");
413        assert!(matches!(p.root.strategy, ExecStrategy::Sequential));
414        let p2 = plan_for("=AND(TRUE(), FALSE())");
415        assert!(matches!(p2.root.strategy, ExecStrategy::Sequential));
416    }
417
418    #[test]
419    fn parentheses_do_not_force_parallelism() {
420        // Trivial groups should stay sequential under default thresholds
421        let p = plan_for("=(1+2)+(2+3)");
422        assert!(matches!(p.root.strategy, ExecStrategy::Sequential));
423    }
424
425    #[test]
426    fn repeated_subtrees_in_sum_encourage_arg_parallel() {
427        // SUM(f(), f(), f(), f()) where f is same subtree
428        let p = plan_for("=SUM(1+2, 1+2, 1+2, 1+2)");
429        // Fanout 4 may or may not cross threshold; accept either but ensure children exist
430        assert!(!p.root.children.is_empty());
431    }
432
433    #[test]
434    fn volatile_forces_sequential() {
435        // NOW() is volatile via caps; planner should mark sequential at root
436        let t = Tokenizer::new("=NOW()+1").unwrap();
437        let mut parser = formualizer_parse::parser::Parser::new(t.items, false);
438        let ast = parser.parse().unwrap();
439        let mut planner = Planner::new(PlanConfig::default())
440            .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name));
441        let plan = planner.plan(&ast);
442        assert!(matches!(plan.root.strategy, ExecStrategy::Sequential));
443    }
444
445    #[test]
446    fn whole_column_ranges_prefer_chunked_reduce() {
447        // Probe A:A to be large → ChunkedReduce at root
448        let t = Tokenizer::new(r#"=SUMIFS(A:A, A:A, ">0", B:B, "<5")"#).unwrap();
449        let mut parser = formualizer_parse::parser::Parser::new(t.items, false);
450        let ast = parser.parse().unwrap();
451        ensure_builtins_registered();
452        let mut planner = Planner::new(PlanConfig {
453            chunk_min_cells: 1000,
454            ..Default::default()
455        })
456        .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name))
457        .with_range_probe(&|r: &ReferenceType| match r {
458            ReferenceType::Range {
459                start_row: None,
460                end_row: None,
461                ..
462            } => Some((50_000, 1)),
463            _ => None,
464        });
465        let plan = planner.plan(&ast);
466        assert!(matches!(
467            plan.root.strategy,
468            ExecStrategy::ChunkedReduce | ExecStrategy::ArgParallel
469        ));
470    }
471
472    #[test]
473    fn deep_sub_ast_criteria_still_plans() {
474        // Deep sub-AST in criteria (e.g., TEXT + DATE math)
475        let p = plan_for("=SUMIFS(A1:A100, B1:B100, TEXT(2024+1, \"0\"))");
476        // Should produce a plan with children; exact strategy may vary
477        assert!(!p.root.children.is_empty());
478    }
479
480    #[test]
481    fn sum_mixed_scalars_and_large_range_prefers_chunked_reduce() {
482        // SUM over a large column plus scalars → prefer chunked reduce due to range cost
483        let t = Tokenizer::new(r#"=SUM(A:A, 1, 2, 3)"#).unwrap();
484        let mut parser = formualizer_parse::parser::Parser::new(t.items, false);
485        let ast = parser.parse().unwrap();
486        ensure_builtins_registered();
487        let mut planner = Planner::new(PlanConfig {
488            chunk_min_cells: 500,
489            ..Default::default()
490        })
491        .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name))
492        .with_range_probe(&|r: &ReferenceType| match r {
493            ReferenceType::Range {
494                start_row: None,
495                end_row: None,
496                ..
497            } => Some((25_000, 1)),
498            _ => None,
499        });
500        let plan = planner.plan(&ast);
501        assert!(matches!(
502            plan.root.strategy,
503            ExecStrategy::ChunkedReduce | ExecStrategy::ArgParallel
504        ));
505    }
506
507    #[test]
508    fn nested_short_circuit_child_remains_sequential_under_parallel_parent() {
509        // Force low thresholds to encourage arg-parallel at parent, but AND child must stay Sequential
510        let t = Tokenizer::new("=SUM(AND(TRUE(), FALSE()), 1, 2, 3)").unwrap();
511        let mut parser = formualizer_parse::parser::Parser::new(t.items, false);
512        let ast = parser.parse().unwrap();
513        ensure_builtins_registered();
514        let cfg = PlanConfig {
515            enable_parallel: true,
516            arg_parallel_min_cost_ns: 0,
517            arg_parallel_min_children: 2,
518            chunk_min_cells: 1_000_000, // disable chunking here
519            chunk_target_partitions: 8,
520        };
521        let mut planner = Planner::new(cfg)
522            .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name));
523        let plan = planner.plan(&ast);
524        // Parent may be ArgParallel under these thresholds
525        assert!(matches!(
526            plan.root.strategy,
527            ExecStrategy::ArgParallel | ExecStrategy::Sequential
528        ));
529        // First child corresponds to AND(...) and must be Sequential due to SHORT_CIRCUIT
530        assert!(!plan.root.children.is_empty());
531        assert!(matches!(
532            plan.root.children[0].strategy,
533            ExecStrategy::Sequential
534        ));
535    }
536
537    #[test]
538    fn repeated_identical_ranges_defaults_to_sequential() {
539        // Repeated A:A references with tiny dims should not trigger chunking and stay Sequential by default thresholds
540        let t = Tokenizer::new(r#"=SUM(A:A, A:A, A:A)"#).unwrap();
541        let mut parser = formualizer_parse::parser::Parser::new(t.items, false);
542        let ast = parser.parse().unwrap();
543        let mut planner = Planner::new(PlanConfig::default())
544            .with_function_lookup(&|ns, name| crate::function_registry::get(ns, name))
545            .with_range_probe(&|r: &ReferenceType| match r {
546                ReferenceType::Range {
547                    start_row: None,
548                    end_row: None,
549                    ..
550                } => Some((3, 1)),
551                _ => None,
552            });
553        let plan = planner.plan(&ast);
554        assert!(matches!(plan.root.strategy, ExecStrategy::Sequential));
555        assert_eq!(plan.root.children.len(), 3);
556    }
557}
formualizer_eval/planner.rs

formualizer_eval/
planner.rs
