linkrs_query/optimizer/cost/node_estimators/
data_processing.rs

1//! Data processing node estimator
2//!
3//! Provide cost estimates for data processing nodes:
4//! - Filter
5//! - Project
6//! - Unwind
7//! - DataCollect
8//! - Start
9
10use super::{get_input_rows, NodeEstimator};
11use crate::optimizer::cost::config::CostModelConfig;
12use crate::optimizer::cost::estimate::NodeCostEstimate;
13use crate::optimizer::cost::expression_parser::ExpressionParser;
14use crate::optimizer::cost::selectivity::SelectivityEstimator;
15use crate::optimizer::cost::CostCalculator;
16use crate::optimizer::error::CostError;
17use crate::optimizer::stats::StatsView;
18use crate::planning::plan::core::nodes::graph_operations::UnwindNode;
19use crate::planning::plan::PlanNodeEnum;
20use linkrs_core::types::BinaryOperator;
21use linkrs_core::Expression;
22
23/// Data processing node estimator
24pub struct DataProcessingEstimator<'a> {
25    cost_calculator: &'a CostCalculator,
26    selectivity_estimator: &'a SelectivityEstimator,
27    expression_parser: ExpressionParser,
28}
29
30impl<'a> DataProcessingEstimator<'a> {
31    /// Create a new data processing estimator.
32    pub fn new(
33        cost_calculator: &'a CostCalculator,
34        selectivity_estimator: &'a SelectivityEstimator,
35        config: CostModelConfig,
36    ) -> Self {
37        let expression_parser = ExpressionParser::new(config);
38        Self {
39            cost_calculator,
40            selectivity_estimator,
41            expression_parser,
42        }
43    }
44
45    /// Calculate the number of filter conditions.
46    pub fn count_filter_conditions(&self, condition: &Expression) -> usize {
47        match condition {
48            Expression::Binary { op, left, right } => match op {
49                BinaryOperator::And => {
50                    self.count_filter_conditions(left) + self.count_filter_conditions(right)
51                }
52                BinaryOperator::Or => (self.count_filter_conditions(left)
53                    + self.count_filter_conditions(right))
54                .max(1),
55                _ => 1,
56            },
57            Expression::Unary { .. } => 1,
58            Expression::Function { args, .. } => args.iter().map(|_| 1).sum::<usize>().max(1),
59            _ => 1,
60        }
61    }
62
63    /// Estimate the list size of the Unwind nodes
64    fn estimate_unwind_list_size(&self, node: &UnwindNode) -> f64 {
65        let list_expr = node.list_expression();
66
67        // Try to parse the expression to determine the size of the list.
68        let expr_str = list_expr.to_expression_string();
69        if let Some(size) = self.expression_parser.parse_list_size(&expr_str) {
70            return size;
71        }
72
73        // Use the default configuration values.
74        self.expression_parser.config().default_unwind_list_size
75    }
76}
77
78impl<'a> NodeEstimator for DataProcessingEstimator<'a> {
79    fn estimate(
80        &self,
81        stats: &StatsView,
82        node: &PlanNodeEnum,
83        child_estimates: &[NodeCostEstimate],
84    ) -> Result<(f64, u64), CostError> {
85        match node {
86            PlanNodeEnum::Filter(n) => {
87                let input_rows_val = get_input_rows(child_estimates, 0);
88                let condition_expr = match n.condition().expression() {
89                    Some(meta) => meta.inner().clone(),
90                    None => return Ok((0.0, input_rows_val)),
91                };
92                let condition_count = self.count_filter_conditions(&condition_expr);
93                // Estimate the number of rows after filtering.
94                let selectivity = self.selectivity_estimator.estimate_from_expression(
95                    stats.space(),
96                    &condition_expr,
97                    None,
98                );
99                let output_rows = (input_rows_val as f64 * selectivity).max(1.0) as u64;
100                let cost = self
101                    .cost_calculator
102                    .calculate_filter_cost(input_rows_val, condition_count);
103                Ok((cost, output_rows))
104            }
105            PlanNodeEnum::Project(n) => {
106                let input_rows_val = get_input_rows(child_estimates, 0);
107                let columns = n.columns().len();
108                let cost = self
109                    .cost_calculator
110                    .calculate_project_cost(input_rows_val, columns);
111                // The “Project” does not change the number of lines.
112                Ok((cost, input_rows_val))
113            }
114            PlanNodeEnum::Unwind(n) => {
115                let input_rows_val = get_input_rows(child_estimates, 0);
116                let list_size = self.estimate_unwind_list_size(n);
117                let cost = self
118                    .cost_calculator
119                    .calculate_unwind_cost(input_rows_val, list_size);
120                // “Unwind” will expand each line into a line of the list size.
121                let output_rows = (input_rows_val as f64 * list_size) as u64;
122                Ok((cost, output_rows.max(1)))
123            }
124            PlanNodeEnum::DataCollect(_) => {
125                let input_rows_val = get_input_rows(child_estimates, 0);
126                let cost = self
127                    .cost_calculator
128                    .calculate_data_collect_cost(input_rows_val);
129                Ok((cost, input_rows_val))
130            }
131            PlanNodeEnum::Start(_) => Ok((0.0, 0)),
132            _ => Err(CostError::UnsupportedNodeType(format!(
133                "Data Processing Estimator does not support node types: {:?}",
134                std::mem::discriminant(node)
135            ))),
136        }
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143    use crate::optimizer::cost::config::CostModelConfig;
144    use crate::planning::plan::core::nodes::control_flow::start_node::StartNode;
145    use crate::planning::plan::core::nodes::graph_operations::graph_operations_node::*;
146    use crate::planning::plan::core::nodes::operation::filter_node::FilterNode;
147    use crate::planning::plan::core::nodes::operation::project_node::ProjectNode;
148    use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
149    use linkrs_core::types::expr::ExpressionMeta;
150    use linkrs_core::YieldColumn;
151    use linkrs_core::{Expression, Value};
152    use std::sync::Arc;
153
154    fn create_test_expression() -> linkrs_core::types::ContextualExpression {
155        let ctx = Arc::new(ExpressionAnalysisContext::new());
156        let expr_meta = ExpressionMeta::new(Expression::Variable("condition".to_string()));
157        let id = ctx.register_expression(expr_meta);
158        linkrs_core::types::ContextualExpression::new(id, ctx)
159    }
160
161    fn create_test_calculator_with_selectivity() -> (CostCalculator, SelectivityEstimator) {
162        let stats_manager = Arc::new(crate::optimizer::stats::StatisticsManager::new());
163        let config = CostModelConfig::default();
164        let calculator = CostCalculator::with_config(stats_manager.clone(), config);
165        let selectivity_estimator = SelectivityEstimator::new(stats_manager);
166        (calculator, selectivity_estimator)
167    }
168
169    #[test]
170    fn test_filter_estimation() {
171        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
172        let config = CostModelConfig::default();
173        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
174
175        let input = PlanNodeEnum::Start(StartNode::new());
176        let condition = create_test_expression();
177        let node = FilterNode::new(input, condition).expect("Node creation should succeed");
178        let plan_node = PlanNodeEnum::Filter(node);
179
180        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 100)];
181        let result = estimator.estimate(
182            &calculator.stats_view(Some("test")),
183            &plan_node,
184            &child_estimates,
185        );
186
187        assert!(result.is_ok());
188        let (cost, output_rows) = result.expect("Estimation should succeed");
189        assert!(cost > 0.0);
190        assert!(output_rows >= 1);
191    }
192
193    #[test]
194    fn test_project_estimation() {
195        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
196        let config = CostModelConfig::default();
197        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
198
199        let input = PlanNodeEnum::Start(StartNode::new());
200        let ctx = Arc::new(ExpressionAnalysisContext::new());
201        let expr_meta = ExpressionMeta::new(Expression::Variable("col".to_string()));
202        let id = ctx.register_expression(expr_meta);
203        let ctx_expr = linkrs_core::types::ContextualExpression::new(id, ctx);
204        let columns = vec![YieldColumn {
205            expression: ctx_expr,
206            alias: "col".to_string(),
207        }];
208        let node = ProjectNode::new(input, columns).expect("Node creation should succeed");
209        let plan_node = PlanNodeEnum::Project(node);
210
211        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 100)];
212        let result = estimator.estimate(
213            &calculator.stats_view(Some("test")),
214            &plan_node,
215            &child_estimates,
216        );
217
218        assert!(result.is_ok());
219        let (cost, output_rows) = result.expect("Estimation should succeed");
220        assert!(cost > 0.0);
221        assert_eq!(output_rows, 100);
222    }
223
224    #[test]
225    fn test_unwind_estimation() {
226        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
227        let config = CostModelConfig::default();
228        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
229
230        let input = PlanNodeEnum::Start(StartNode::new());
231        let list_expr = create_test_expression();
232        let node = UnwindNode::new(input, "item", list_expr).expect("Node creation should succeed");
233        let plan_node = PlanNodeEnum::Unwind(node);
234
235        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 100)];
236        let result = estimator.estimate(
237            &calculator.stats_view(Some("test")),
238            &plan_node,
239            &child_estimates,
240        );
241
242        assert!(result.is_ok());
243        let (cost, output_rows) = result.expect("Estimation should succeed");
244        assert!(cost > 0.0);
245        assert!(output_rows >= 1);
246    }
247
248    #[test]
249    fn test_data_collect_estimation() {
250        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
251        let config = CostModelConfig::default();
252        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
253
254        let input = PlanNodeEnum::Start(StartNode::new());
255        let node = DataCollectNode::new(input, "ROW").expect("Node creation should succeed");
256        let plan_node = PlanNodeEnum::DataCollect(node);
257
258        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 100)];
259        let result = estimator.estimate(
260            &calculator.stats_view(Some("test")),
261            &plan_node,
262            &child_estimates,
263        );
264
265        assert!(result.is_ok());
266        let (cost, output_rows) = result.expect("Estimation should succeed");
267        assert!(cost > 0.0);
268        assert_eq!(output_rows, 100);
269    }
270
271    #[test]
272    fn test_start_estimation() {
273        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
274        let config = CostModelConfig::default();
275        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
276
277        let node = StartNode::new();
278        let plan_node = PlanNodeEnum::Start(node);
279
280        let child_estimates = vec![];
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_eq!(cost, 0.0);
290        assert_eq!(output_rows, 0);
291    }
292
293    #[test]
294    fn test_unsupported_node_type() {
295        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
296        let config = CostModelConfig::default();
297        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
298
299        let node = PlanNodeEnum::ScanVertices(
300            crate::planning::plan::core::nodes::access::graph_scan_node::ScanVerticesNode::new(
301                1,
302                "test_space",
303            ),
304        );
305        let child_estimates = vec![];
306        let result = estimator.estimate(
307            &calculator.stats_view(Some("test")),
308            &node,
309            &child_estimates,
310        );
311
312        assert!(result.is_err());
313    }
314
315    #[test]
316    fn test_count_filter_conditions_simple() {
317        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
318        let config = CostModelConfig::default();
319        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
320
321        let condition = Expression::Binary {
322            op: BinaryOperator::Equal,
323            left: Box::new(Expression::Variable("a".to_string())),
324            right: Box::new(Expression::Literal(Value::string("1"))),
325        };
326        let count = estimator.count_filter_conditions(&condition);
327        assert_eq!(count, 1);
328    }
329
330    #[test]
331    fn test_count_filter_conditions_and() {
332        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
333        let config = CostModelConfig::default();
334        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
335
336        let condition = Expression::Binary {
337            op: BinaryOperator::And,
338            left: Box::new(Expression::Variable("a".to_string())),
339            right: Box::new(Expression::Variable("b".to_string())),
340        };
341        let count = estimator.count_filter_conditions(&condition);
342        assert_eq!(count, 2);
343    }
344
345    #[test]
346    fn test_count_filter_conditions_or() {
347        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
348        let config = CostModelConfig::default();
349        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
350
351        let condition = Expression::Binary {
352            op: BinaryOperator::Or,
353            left: Box::new(Expression::Variable("a".to_string())),
354            right: Box::new(Expression::Variable("b".to_string())),
355        };
356        let count = estimator.count_filter_conditions(&condition);
357        assert_eq!(count, 2);
358    }
359
360    #[test]
361    fn test_count_filter_conditions_complex() {
362        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
363        let config = CostModelConfig::default();
364        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
365
366        let condition = Expression::Binary {
367            op: BinaryOperator::And,
368            left: Box::new(Expression::Binary {
369                op: BinaryOperator::Equal,
370                left: Box::new(Expression::Variable("a".to_string())),
371                right: Box::new(Expression::Literal(Value::string("1"))),
372            }),
373            right: Box::new(Expression::Binary {
374                op: BinaryOperator::Or,
375                left: Box::new(Expression::Variable("b".to_string())),
376                right: Box::new(Expression::Variable("c".to_string())),
377            }),
378        };
379        let count = estimator.count_filter_conditions(&condition);
380        assert_eq!(count, 3);
381    }
382
383    #[test]
384    fn test_filter_with_zero_input() {
385        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
386        let config = CostModelConfig::default();
387        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
388
389        let input = PlanNodeEnum::Start(StartNode::new());
390        let condition = create_test_expression();
391        let node = FilterNode::new(input, condition).expect("Node creation should succeed");
392        let plan_node = PlanNodeEnum::Filter(node);
393
394        let child_estimates = vec![NodeCostEstimate::new(0.0, 0.0, 0)];
395        let result = estimator.estimate(
396            &calculator.stats_view(Some("test")),
397            &plan_node,
398            &child_estimates,
399        );
400
401        assert!(result.is_ok());
402        let (cost, output_rows) = result.expect("Estimation should succeed");
403        assert!(cost >= 0.0);
404        assert_eq!(output_rows, 1);
405    }
406
407    #[test]
408    fn test_project_with_multiple_columns() {
409        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
410        let config = CostModelConfig::default();
411        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
412
413        let input = PlanNodeEnum::Start(StartNode::new());
414        let ctx = Arc::new(ExpressionAnalysisContext::new());
415        let columns = vec![
416            YieldColumn {
417                expression: {
418                    let expr_meta = ExpressionMeta::new(Expression::Variable("a".to_string()));
419                    let id = ctx.register_expression(expr_meta);
420                    linkrs_core::types::ContextualExpression::new(id, ctx.clone())
421                },
422                alias: "a".to_string(),
423            },
424            YieldColumn {
425                expression: {
426                    let expr_meta = ExpressionMeta::new(Expression::Variable("b".to_string()));
427                    let id = ctx.register_expression(expr_meta);
428                    linkrs_core::types::ContextualExpression::new(id, ctx.clone())
429                },
430                alias: "b".to_string(),
431            },
432            YieldColumn {
433                expression: {
434                    let expr_meta = ExpressionMeta::new(Expression::Variable("c".to_string()));
435                    let id = ctx.register_expression(expr_meta);
436                    linkrs_core::types::ContextualExpression::new(id, ctx.clone())
437                },
438                alias: "c".to_string(),
439            },
440        ];
441        let node = ProjectNode::new(input, columns).expect("Node creation should succeed");
442        let plan_node = PlanNodeEnum::Project(node);
443
444        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 100)];
445        let result = estimator.estimate(
446            &calculator.stats_view(Some("test")),
447            &plan_node,
448            &child_estimates,
449        );
450
451        assert!(result.is_ok());
452        let (cost, output_rows) = result.expect("Estimation should succeed");
453        assert!(cost > 0.0);
454        assert_eq!(output_rows, 100);
455    }
456
457    #[test]
458    fn test_unwind_with_large_list() {
459        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
460        let config = CostModelConfig::default();
461        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
462
463        let input = PlanNodeEnum::Start(StartNode::new());
464        let list_expr = create_test_expression();
465        let node = UnwindNode::new(input, "item", list_expr).expect("Node creation should succeed");
466        let plan_node = PlanNodeEnum::Unwind(node);
467
468        let child_estimates = vec![NodeCostEstimate::new(10.0, 10.0, 1000)];
469        let result = estimator.estimate(
470            &calculator.stats_view(Some("test")),
471            &plan_node,
472            &child_estimates,
473        );
474
475        assert!(result.is_ok());
476        let (cost, output_rows) = result.expect("Estimation should succeed");
477        assert!(cost > 0.0);
478        assert!(output_rows >= 1);
479    }
480
481    #[test]
482    fn test_data_collect_with_large_input() {
483        let (calculator, selectivity_estimator) = create_test_calculator_with_selectivity();
484        let config = CostModelConfig::default();
485        let estimator = DataProcessingEstimator::new(&calculator, &selectivity_estimator, config);
486
487        let input = PlanNodeEnum::Start(StartNode::new());
488        let node = DataCollectNode::new(input, "ROW").expect("Node creation should succeed");
489        let plan_node = PlanNodeEnum::DataCollect(node);
490
491        let child_estimates = vec![NodeCostEstimate::new(1000.0, 1000.0, 1_000_000)];
492        let result = estimator.estimate(
493            &calculator.stats_view(Some("test")),
494            &plan_node,
495            &child_estimates,
496        );
497
498        assert!(result.is_ok());
499        let (cost, output_rows) = result.expect("Estimation should succeed");
500        assert!(cost > 0.0);
501        assert_eq!(output_rows, 1_000_000);
502    }
503}
linkrs_query/optimizer/cost/node_estimators/data_processing.rs

linkrs_query/optimizer/cost/node_estimators/
data_processing.rs
