linkrs_query/optimizer/cost/node_estimators/
control_flow.rs

1//! Control Flow Node Estimator
2//!
3//! Provide cost estimates for flow control nodes:
4//! - Loop
5//! - Select
6//! - PassThrough
7//! - Argument
8
9use super::{get_input_rows, NodeEstimator};
10use crate::optimizer::cost::config::CostModelConfig;
11use crate::optimizer::cost::estimate::NodeCostEstimate;
12use crate::optimizer::cost::expression_parser::ExpressionParser;
13use crate::optimizer::cost::CostCalculator;
14use crate::optimizer::error::CostError;
15use crate::optimizer::stats::StatsView;
16use crate::planning::plan::core::nodes::control_flow::control_flow_node::{LoopNode, SelectNode};
17use crate::planning::plan::PlanNodeEnum;
18
19/// Control Flow Node Estimator
20pub struct ControlFlowEstimator<'a> {
21    cost_calculator: &'a CostCalculator,
22    config: CostModelConfig,
23    expression_parser: ExpressionParser,
24}
25
26impl<'a> ControlFlowEstimator<'a> {
27    /// Create a new control flow estimator.
28    pub fn new(cost_calculator: &'a CostCalculator, config: CostModelConfig) -> Self {
29        let expression_parser = ExpressionParser::new(config);
30        Self {
31            cost_calculator,
32            config,
33            expression_parser,
34        }
35    }
36
37    /// Estimating the number of iterations of the Loop node
38    fn estimate_loop_iterations(&self, node: &LoopNode) -> u32 {
39        let condition = node.condition().to_expression_string();
40
41        // Try to parse the number of iterations using an expression parser.
42        if let Some(iterations) = self.expression_parser.parse_loop_iterations(&condition) {
43            return iterations;
44        }
45
46        // The default configuration values will be used.
47        self.config.default_loop_iterations
48    }
49
50    /// Estimate the number of branches of the selected node.
51    fn estimate_select_branch_count(&self, node: &SelectNode) -> usize {
52        let mut count = 0;
53        if node.if_branch().is_some() {
54            count += 1;
55        }
56        if node.else_branch().is_some() {
57            count += 1;
58        }
59
60        if count == 0 {
61            self.config.default_select_branches
62        } else {
63            count
64        }
65    }
66}
67
68impl<'a> NodeEstimator for ControlFlowEstimator<'a> {
69    fn estimate(
70        &self,
71        _stats: &StatsView,
72        node: &PlanNodeEnum,
73        child_estimates: &[NodeCostEstimate],
74    ) -> Result<(f64, u64), CostError> {
75        match node {
76            PlanNodeEnum::Loop(n) => {
77                let body_estimate = child_estimates
78                    .first()
79                    .copied()
80                    .unwrap_or(NodeCostEstimate::new(0.0, 0.0, 0));
81                let iterations = self.estimate_loop_iterations(n);
82                let cost = self
83                    .cost_calculator
84                    .calculate_loop_cost(body_estimate.total_cost, iterations);
85                // The number of lines output by the `Loop` function is equal to the number of lines output by the loop body multiplied by the number of iterations.
86                let output_rows = body_estimate.output_rows.saturating_mul(iterations as u64);
87                Ok((cost, output_rows))
88            }
89            PlanNodeEnum::RecursiveCte(n) => {
90                // anchor cost once + step cost per iteration. Child
91                // estimates are [anchor, step?].
92                let anchor_estimate = child_estimates
93                    .first()
94                    .copied()
95                    .unwrap_or(NodeCostEstimate::new(0.0, 0.0, 0));
96                let step_estimate = child_estimates
97                    .get(1)
98                    .copied()
99                    .unwrap_or(NodeCostEstimate::new(0.0, 0.0, 0));
100                let iterations = n.max_iterations().min(u64::from(u32::MAX));
101                let step_cost = self
102                    .cost_calculator
103                    .calculate_loop_cost(step_estimate.total_cost, iterations as u32);
104                let cost = anchor_estimate.total_cost + step_cost;
105                let output_rows = anchor_estimate
106                    .output_rows
107                    .saturating_add(step_estimate.output_rows.saturating_mul(iterations));
108                Ok((cost, output_rows))
109            }
110            PlanNodeEnum::Select(n) => {
111                let input_rows_val = get_input_rows(child_estimates, 0);
112                let branch_count = self.estimate_select_branch_count(n);
113                let cost = self
114                    .cost_calculator
115                    .calculate_select_cost(input_rows_val, branch_count);
116                // The number of output lines should be equal to the number of input lines (assuming that on average one branch is selected).
117                Ok((cost, input_rows_val))
118            }
119            PlanNodeEnum::PassThrough(_) => {
120                let input_rows_val = get_input_rows(child_estimates, 0);
121                let cost = self
122                    .cost_calculator
123                    .calculate_pass_through_cost(input_rows_val);
124                Ok((cost, input_rows_val))
125            }
126            PlanNodeEnum::Argument(_) => Ok((0.0, 1)),
127            _ => Err(CostError::UnsupportedNodeType(format!(
128                "The control flow estimator does not support the node type: {:?}",
129                std::mem::discriminant(node)
130            ))),
131        }
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138    use crate::optimizer::cost::config::CostModelConfig;
139    use crate::planning::plan::core::nodes::control_flow::control_flow_node::*;
140    use crate::planning::plan::core::nodes::control_flow::start_node::StartNode;
141    use std::sync::Arc;
142
143    fn create_test_calculator() -> CostCalculator {
144        let stats_manager = Arc::new(crate::optimizer::stats::StatisticsManager::new());
145        let config = CostModelConfig::default();
146        CostCalculator::with_config(stats_manager, config)
147    }
148
149    fn create_test_expression() -> linkrs_core::types::ContextualExpression {
150        use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
151        use linkrs_core::types::expr::ExpressionMeta;
152        use linkrs_core::Expression;
153
154        let ctx = Arc::new(ExpressionAnalysisContext::new());
155        let expr_meta = ExpressionMeta::new(Expression::Variable("condition".to_string()));
156        let id = ctx.register_expression(expr_meta);
157        linkrs_core::types::ContextualExpression::new(id, ctx)
158    }
159
160    #[test]
161    fn test_loop_estimation() {
162        let calculator = create_test_calculator();
163        let config = CostModelConfig::default();
164        let estimator = ControlFlowEstimator::new(&calculator, config);
165
166        let condition = create_test_expression();
167        let mut node = LoopNode::new(1, condition);
168        node.set_body(PlanNodeEnum::Start(StartNode::new()));
169        let plan_node = PlanNodeEnum::Loop(node);
170
171        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 100)];
172        let result = estimator.estimate(
173            &calculator.stats_view(Some("test")),
174            &plan_node,
175            &child_estimates,
176        );
177
178        assert!(result.is_ok());
179        let (cost, output_rows) = result.expect("Estimation should succeed");
180        assert!(cost > 0.0);
181        assert!(output_rows >= 1);
182    }
183
184    #[test]
185    fn test_select_estimation() {
186        let calculator = create_test_calculator();
187        let config = CostModelConfig::default();
188        let estimator = ControlFlowEstimator::new(&calculator, config);
189
190        let condition = create_test_expression();
191        let mut node = SelectNode::new(1, condition);
192        node.set_if_branch(PlanNodeEnum::Start(StartNode::new()));
193        node.set_else_branch(PlanNodeEnum::Start(StartNode::new()));
194        let plan_node = PlanNodeEnum::Select(node);
195
196        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 100)];
197        let result = estimator.estimate(
198            &calculator.stats_view(Some("test")),
199            &plan_node,
200            &child_estimates,
201        );
202
203        assert!(result.is_ok());
204        let (cost, output_rows) = result.expect("Estimation should succeed");
205        assert!(cost > 0.0);
206        assert_eq!(output_rows, 100);
207    }
208
209    #[test]
210    fn test_pass_through_estimation() {
211        let calculator = create_test_calculator();
212        let config = CostModelConfig::default();
213        let estimator = ControlFlowEstimator::new(&calculator, config);
214
215        let node = PassThroughNode::new(1);
216        let plan_node = PlanNodeEnum::PassThrough(node);
217
218        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 100)];
219        let result = estimator.estimate(
220            &calculator.stats_view(Some("test")),
221            &plan_node,
222            &child_estimates,
223        );
224
225        assert!(result.is_ok());
226        let (cost, output_rows) = result.expect("Estimation should succeed");
227        assert!(cost > 0.0);
228        assert_eq!(output_rows, 100);
229    }
230
231    #[test]
232    fn test_argument_estimation() {
233        let calculator = create_test_calculator();
234        let config = CostModelConfig::default();
235        let estimator = ControlFlowEstimator::new(&calculator, config);
236
237        let node = ArgumentNode::new(1, "var_name");
238        let plan_node = PlanNodeEnum::Argument(node);
239
240        let child_estimates = vec![];
241        let result = estimator.estimate(
242            &calculator.stats_view(Some("test")),
243            &plan_node,
244            &child_estimates,
245        );
246
247        assert!(result.is_ok());
248        let (cost, output_rows) = result.expect("Estimation should succeed");
249        assert_eq!(cost, 0.0);
250        assert_eq!(output_rows, 1);
251    }
252
253    #[test]
254    fn test_unsupported_node_type() {
255        let calculator = create_test_calculator();
256        let config = CostModelConfig::default();
257        let estimator = ControlFlowEstimator::new(&calculator, config);
258
259        let node = PlanNodeEnum::Start(StartNode::new());
260        let child_estimates = vec![];
261        let result = estimator.estimate(
262            &calculator.stats_view(Some("test")),
263            &node,
264            &child_estimates,
265        );
266
267        assert!(result.is_err());
268    }
269
270    #[test]
271    fn test_select_without_branches() {
272        let calculator = create_test_calculator();
273        let config = CostModelConfig::default();
274        let estimator = ControlFlowEstimator::new(&calculator, config);
275
276        let condition = create_test_expression();
277        let node = SelectNode::new(1, condition);
278        let plan_node = PlanNodeEnum::Select(node);
279
280        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 100)];
281        let result = estimator.estimate(
282            &calculator.stats_view(Some("test")),
283            &plan_node,
284            &child_estimates,
285        );
286
287        assert!(result.is_ok());
288        let (cost, output_rows) = result.expect("Estimation should succeed");
289        assert!(cost > 0.0);
290        assert_eq!(output_rows, 100);
291    }
292
293    #[test]
294    fn test_select_with_only_if_branch() {
295        let calculator = create_test_calculator();
296        let config = CostModelConfig::default();
297        let estimator = ControlFlowEstimator::new(&calculator, config);
298
299        let condition = create_test_expression();
300        let mut node = SelectNode::new(1, condition);
301        node.set_if_branch(PlanNodeEnum::Start(StartNode::new()));
302        let plan_node = PlanNodeEnum::Select(node);
303
304        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 100)];
305        let result = estimator.estimate(
306            &calculator.stats_view(Some("test")),
307            &plan_node,
308            &child_estimates,
309        );
310
311        assert!(result.is_ok());
312        let (cost, output_rows) = result.expect("Estimation should succeed");
313        assert!(cost > 0.0);
314        assert_eq!(output_rows, 100);
315    }
316
317    #[test]
318    fn test_select_with_only_else_branch() {
319        let calculator = create_test_calculator();
320        let config = CostModelConfig::default();
321        let estimator = ControlFlowEstimator::new(&calculator, config);
322
323        let condition = create_test_expression();
324        let mut node = SelectNode::new(1, condition);
325        node.set_else_branch(PlanNodeEnum::Start(StartNode::new()));
326        let plan_node = PlanNodeEnum::Select(node);
327
328        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 100)];
329        let result = estimator.estimate(
330            &calculator.stats_view(Some("test")),
331            &plan_node,
332            &child_estimates,
333        );
334
335        assert!(result.is_ok());
336        let (cost, output_rows) = result.expect("Estimation should succeed");
337        assert!(cost > 0.0);
338        assert_eq!(output_rows, 100);
339    }
340
341    #[test]
342    fn test_loop_with_different_iterations() {
343        let calculator = create_test_calculator();
344        let config = CostModelConfig::default();
345        let estimator = ControlFlowEstimator::new(&calculator, config);
346
347        let condition = create_test_expression();
348        let mut node = LoopNode::new(1, condition);
349        node.set_body(PlanNodeEnum::Start(StartNode::new()));
350        let plan_node = PlanNodeEnum::Loop(node);
351
352        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 100)];
353        let result = estimator.estimate(
354            &calculator.stats_view(Some("test")),
355            &plan_node,
356            &child_estimates,
357        );
358
359        assert!(result.is_ok());
360        let (cost, output_rows) = result.expect("Estimation should succeed");
361        assert!(cost > 0.0);
362        assert!(output_rows >= 1);
363    }
364
365    #[test]
366    fn test_pass_through_with_zero_input() {
367        let calculator = create_test_calculator();
368        let config = CostModelConfig::default();
369        let estimator = ControlFlowEstimator::new(&calculator, config);
370
371        let node = PassThroughNode::new(1);
372        let plan_node = PlanNodeEnum::PassThrough(node);
373
374        let child_estimates = vec![NodeCostEstimate::new(0.0, 0.0, 0)];
375        let result = estimator.estimate(
376            &calculator.stats_view(Some("test")),
377            &plan_node,
378            &child_estimates,
379        );
380
381        assert!(result.is_ok());
382        let (cost, output_rows) = result.expect("Estimation should succeed");
383        assert!(cost >= 0.0);
384        assert_eq!(output_rows, 0);
385    }
386
387    #[test]
388    fn test_pass_through_with_large_input() {
389        let calculator = create_test_calculator();
390        let config = CostModelConfig::default();
391        let estimator = ControlFlowEstimator::new(&calculator, config);
392
393        let node = PassThroughNode::new(1);
394        let plan_node = PlanNodeEnum::PassThrough(node);
395
396        let child_estimates = vec![NodeCostEstimate::new(1000.0, 1000.0, 1_000_000)];
397        let result = estimator.estimate(
398            &calculator.stats_view(Some("test")),
399            &plan_node,
400            &child_estimates,
401        );
402
403        assert!(result.is_ok());
404        let (cost, output_rows) = result.expect("Estimation should succeed");
405        assert!(cost > 0.0);
406        assert_eq!(output_rows, 1_000_000);
407    }
408
409    #[test]
410    fn test_loop_without_child_estimates() {
411        let calculator = create_test_calculator();
412        let config = CostModelConfig::default();
413        let estimator = ControlFlowEstimator::new(&calculator, config);
414
415        let condition = create_test_expression();
416        let mut node = LoopNode::new(1, condition);
417        node.set_body(PlanNodeEnum::Start(StartNode::new()));
418        let plan_node = PlanNodeEnum::Loop(node);
419
420        let child_estimates = vec![];
421        let result = estimator.estimate(
422            &calculator.stats_view(Some("test")),
423            &plan_node,
424            &child_estimates,
425        );
426
427        assert!(result.is_ok());
428        let (cost, output_rows) = result.expect("Estimation should succeed");
429        assert!(cost >= 0.0);
430        assert_eq!(output_rows, 0);
431    }
432}
linkrs_query/optimizer/cost/node_estimators/control_flow.rs

linkrs_query/optimizer/cost/node_estimators/
control_flow.rs
