Skip to main content

linkrs_query/optimizer/cost/
expression_parser.rs

1//! Expression parser
2//!
3//! Provide an expression parsing function for:
4//! Estimate the size of the list of Unwind nodes.
5//! Estimate the number of iterations of the Loop node.
6//! Analyze various expression patterns
7//! Expression constant folding optimization
8
9use crate::optimizer::cost::config::CostModelConfig;
10use linkrs_core::types::{BinaryOperator, Expression, UnaryOperator};
11use linkrs_core::value::Value;
12
13/// Expression parser
14///
15/// Used to parse useful information from expression strings, such as the size of a list, the number of iterations, etc.
16#[derive(Debug, Clone)]
17pub struct ExpressionParser {
18    /// Configuration (for internal use)
19    config: CostModelConfig,
20}
21
22impl ExpressionParser {
23    /// Obtain the configuration.
24    pub fn config(&self) -> &CostModelConfig {
25        &self.config
26    }
27}
28
29impl ExpressionParser {
30    /// Create a new expression parser.
31    pub fn new(config: CostModelConfig) -> Self {
32        Self { config }
33    }
34
35    /// Try to parse the list size from the expression string.
36    ///
37    /// The following modes are supported:
38    /// - Array literals: [a, b, c] -> 3
39    /// - range function: range(1, 10) -> 9, range(1, 10, 2) -> 5
40    /// - Range expressions: 1..10 -> 9, 0..=5 -> 6
41    /// - Collection functions: keys(map), values(map), nodes(path), relationships(path)
42    /// - String split: split(str, ",") (estimated)
43    /// - Collection operations: collect(set) (estimated)
44    pub fn parse_list_size(&self, expr: &str) -> Option<f64> {
45        let expr = expr.trim();
46
47        // Attempt to parse array literals [a, b, c]
48        if expr.starts_with('[') && expr.ends_with(']') {
49            return self.parse_array_literal(expr);
50        }
51
52        // Try to parse range(start, end) or range(start, end, step)
53        if expr.starts_with("range(") && expr.ends_with(')') {
54            return self.parse_range_function(expr);
55        }
56
57        // Try to parse the range expression: 1..10 or 0... =5
58        if expr.contains("..") {
59            return self.parse_range_expression(expr);
60        }
61
62        // Try to parse collection functions: keys(), values(), nodes(), relationships()
63        if let Some(size) = self.parse_collection_function(expr) {
64            return Some(size);
65        }
66
67        // Trying to parse a string splitting function
68        if let Some(size) = self.parse_split_function(expr) {
69            return Some(size);
70        }
71
72        // Try to parse the collect function (usually used for aggregation)
73        let expr_lower = expr.to_lowercase();
74        if expr_lower.starts_with("collect(") || expr_lower.contains(".collect()") {
75            // The size of the result of the collect function depends on the input data, using a conservative estimate
76            return Some(self.config.default_unwind_list_size * 2.0);
77        }
78
79        None
80    }
81
82    // ==================== Constant folding optimization ====================
83
84    /// Try to fold the constants within the expression.
85    ///
86    /// Perform a recursive traversal of the expression, and replace all constant expressions that can be evaluated with their literal values.
87    /// For example: 1 + 2 -> 3, “hello” + “world” -> “helloworld”.
88    ///
89    /// # Parameters
90    /// `expr`: The input expression
91    ///
92    /// # Return value
93    /// The folded expression
94    pub fn fold_constants(&self, expr: &Expression) -> Expression {
95        match expr {
96            Expression::Binary { left, op, right } => {
97                let folded_left = self.fold_constants(left);
98                let folded_right = self.fold_constants(right);
99
100                // If both sides are constants, calculate the result directly
101                if let (Expression::Literal(l), Expression::Literal(r)) =
102                    (&folded_left, &folded_right)
103                {
104                    if let Some(result) = self.evaluate_binary_op(op, l, r) {
105                        return Expression::Literal(result);
106                    }
107                }
108
109                Expression::Binary {
110                    left: Box::new(folded_left),
111                    op: *op,
112                    right: Box::new(folded_right),
113                }
114            }
115            Expression::Unary { op, operand } => {
116                let folded_operand = self.fold_constants(operand);
117
118                // If the operand is a constant, compute the result directly
119                if let Expression::Literal(v) = &folded_operand {
120                    if let Some(result) = self.evaluate_unary_op(op, v) {
121                        return Expression::Literal(result);
122                    }
123                }
124
125                Expression::Unary {
126                    op: *op,
127                    operand: Box::new(folded_operand),
128                }
129            }
130            Expression::Function { name, args } => {
131                let folded_args: Vec<Expression> = args
132                    .iter()
133                    .map(|arg| self.fold_constants(arg.as_expr()))
134                    .collect();
135
136                // If all arguments are constants, try to compute the function
137                if folded_args
138                    .iter()
139                    .all(|arg| matches!(arg, Expression::Literal(_)))
140                {
141                    let arg_values: Vec<&Value> = folded_args
142                        .iter()
143                        .filter_map(|arg| match arg {
144                            Expression::Literal(v) => Some(v),
145                            _ => None,
146                        })
147                        .collect();
148
149                    if let Some(result) = self.evaluate_function(name, &arg_values) {
150                        return Expression::Literal(result);
151                    }
152                }
153
154                Expression::function(name.clone(), folded_args)
155            }
156            Expression::List(items) => {
157                let folded_items: Vec<Expression> =
158                    items.iter().map(|item| self.fold_constants(item)).collect();
159                Expression::List(folded_items)
160            }
161            Expression::Map(entries) => {
162                let folded_entries: Vec<(String, Expression)> = entries
163                    .iter()
164                    .map(|(k, v)| (k.clone(), self.fold_constants(v)))
165                    .collect();
166                Expression::Map(folded_entries)
167            }
168            _ => expr.clone(),
169        }
170    }
171
172    /// Evaluating binary operations
173    fn evaluate_binary_op(
174        &self,
175        op: &BinaryOperator,
176        left: &Value,
177        right: &Value,
178    ) -> Option<Value> {
179        match op {
180            BinaryOperator::Add => self.add_values(left, right),
181            BinaryOperator::Subtract => self.subtract_values(left, right),
182            BinaryOperator::Multiply => self.multiply_values(left, right),
183            BinaryOperator::Divide => self.divide_values(left, right),
184            BinaryOperator::Modulo => self.modulo_values(left, right),
185            BinaryOperator::Equal => Some(Value::Bool(
186                self.compare_values(left, right) == Some(std::cmp::Ordering::Equal),
187            )),
188            BinaryOperator::NotEqual => Some(Value::Bool(
189                self.compare_values(left, right) != Some(std::cmp::Ordering::Equal),
190            )),
191            BinaryOperator::LessThan => Some(Value::Bool(
192                self.compare_values(left, right) == Some(std::cmp::Ordering::Less),
193            )),
194            BinaryOperator::GreaterThan => Some(Value::Bool(
195                self.compare_values(left, right) == Some(std::cmp::Ordering::Greater),
196            )),
197            BinaryOperator::LessThanOrEqual => {
198                let cmp = self.compare_values(left, right);
199                Some(Value::Bool(
200                    cmp == Some(std::cmp::Ordering::Less) || cmp == Some(std::cmp::Ordering::Equal),
201                ))
202            }
203            BinaryOperator::GreaterThanOrEqual => {
204                let cmp = self.compare_values(left, right);
205                Some(Value::Bool(
206                    cmp == Some(std::cmp::Ordering::Greater)
207                        || cmp == Some(std::cmp::Ordering::Equal),
208                ))
209            }
210            BinaryOperator::And => self.logical_and(left, right),
211            BinaryOperator::Or => self.logical_or(left, right),
212            BinaryOperator::StringConcat => self.concat_values(left, right),
213            _ => None,
214        }
215    }
216
217    /// Evaluating a unary operation
218    fn evaluate_unary_op(&self, op: &UnaryOperator, operand: &Value) -> Option<Value> {
219        match op {
220            UnaryOperator::Not => match operand {
221                Value::Bool(b) => Some(Value::Bool(!b)),
222                _ => None,
223            },
224            UnaryOperator::Minus => match operand {
225                Value::Int(i) => Some(Value::Int(-i)),
226                Value::Float(f) => Some(Value::Float(-f)),
227                _ => None,
228            },
229            UnaryOperator::IsNull => Some(Value::Bool(operand.is_null())),
230            UnaryOperator::IsNotNull => Some(Value::Bool(!operand.is_null())),
231            _ => None,
232        }
233    }
234
235    /// Evaluation function
236    fn evaluate_function(&self, name: &str, args: &[&Value]) -> Option<Value> {
237        let name_lower = name.to_lowercase();
238
239        match name_lower.as_str() {
240            "abs" if args.len() == 1 => match args[0] {
241                Value::SmallInt(i) => Some(Value::SmallInt(i.abs())),
242                Value::Int(i) => Some(Value::Int(i.abs())),
243                Value::BigInt(i) => Some(Value::BigInt(i.abs())),
244                Value::Float(f) => Some(Value::Float(f.abs())),
245                Value::Double(f) => Some(Value::Double(f.abs())),
246                _ => None,
247            },
248            "length" | "size" if args.len() == 1 => match args[0] {
249                Value::String(s) => Some(Value::BigInt(s.len() as i64)),
250                Value::List(list) => Some(Value::BigInt(list.len() as i64)),
251                _ => None,
252            },
253            "upper" | "toupper" if args.len() == 1 => match args[0] {
254                Value::String(s) => Some(Value::string(s.to_uppercase())),
255                _ => None,
256            },
257            "lower" | "tolower" if args.len() == 1 => match args[0] {
258                Value::String(s) => Some(Value::string(s.to_lowercase())),
259                _ => None,
260            },
261            "substring" | "substr" if args.len() >= 2 => match (args[0], args.get(1).copied()) {
262                (Value::String(s), Some(Value::BigInt(start))) => {
263                    let start_idx = if *start >= 0 { *start as usize } else { 0 };
264                    let end_idx = if args.len() >= 3 {
265                        match args[2] {
266                            Value::BigInt(len) => start_idx + (*len as usize),
267                            _ => s.len(),
268                        }
269                    } else {
270                        s.len()
271                    };
272                    Some(Value::string(
273                        s.chars()
274                            .skip(start_idx)
275                            .take(end_idx - start_idx)
276                            .collect::<String>(),
277                    ))
278                }
279                _ => None,
280            },
281            "trim" if args.len() == 1 => match args[0] {
282                Value::String(s) => Some(Value::string(s.trim())),
283                _ => None,
284            },
285            _ => None,
286        }
287    }
288
289    // ==================== Value operation helper methods ====================
290
291    fn add_values(&self, left: &Value, right: &Value) -> Option<Value> {
292        match (left, right) {
293            (Value::SmallInt(l), Value::SmallInt(r)) => Some(Value::SmallInt(l + r)),
294            (Value::Int(l), Value::Int(r)) => Some(Value::Int(l + r)),
295            (Value::BigInt(l), Value::BigInt(r)) => Some(Value::BigInt(l + r)),
296            (Value::Float(l), Value::Float(r)) => Some(Value::Float(l + r)),
297            (Value::Double(l), Value::Double(r)) => Some(Value::Double(l + r)),
298            (Value::String(l), Value::String(r)) => Some(Value::string(format!("{}{}", l, r))),
299            _ => None,
300        }
301    }
302
303    fn subtract_values(&self, left: &Value, right: &Value) -> Option<Value> {
304        match (left, right) {
305            (Value::SmallInt(l), Value::SmallInt(r)) => Some(Value::SmallInt(l - r)),
306            (Value::Int(l), Value::Int(r)) => Some(Value::Int(l - r)),
307            (Value::BigInt(l), Value::BigInt(r)) => Some(Value::BigInt(l - r)),
308            (Value::Float(l), Value::Float(r)) => Some(Value::Float(l - r)),
309            (Value::Double(l), Value::Double(r)) => Some(Value::Double(l - r)),
310            _ => None,
311        }
312    }
313
314    fn multiply_values(&self, left: &Value, right: &Value) -> Option<Value> {
315        match (left, right) {
316            (Value::SmallInt(l), Value::SmallInt(r)) => Some(Value::SmallInt(l * r)),
317            (Value::Int(l), Value::Int(r)) => Some(Value::Int(l * r)),
318            (Value::BigInt(l), Value::BigInt(r)) => Some(Value::BigInt(l * r)),
319            (Value::Float(l), Value::Float(r)) => Some(Value::Float(l * r)),
320            (Value::Double(l), Value::Double(r)) => Some(Value::Double(l * r)),
321            _ => None,
322        }
323    }
324
325    fn divide_values(&self, left: &Value, right: &Value) -> Option<Value> {
326        match (left, right) {
327            (Value::SmallInt(l), Value::SmallInt(r)) if *r != 0 => Some(Value::SmallInt(l / r)),
328            (Value::Int(l), Value::Int(r)) if *r != 0 => Some(Value::Int(l / r)),
329            (Value::BigInt(l), Value::BigInt(r)) if *r != 0 => Some(Value::BigInt(l / r)),
330            (Value::Float(l), Value::Float(r)) if *r != 0.0 => Some(Value::Float(l / r)),
331            (Value::Double(l), Value::Double(r)) if *r != 0.0 => Some(Value::Double(l / r)),
332            _ => None,
333        }
334    }
335
336    fn modulo_values(&self, left: &Value, right: &Value) -> Option<Value> {
337        match (left, right) {
338            (Value::SmallInt(l), Value::SmallInt(r)) if *r != 0 => Some(Value::SmallInt(l % r)),
339            (Value::Int(l), Value::Int(r)) if *r != 0 => Some(Value::Int(l % r)),
340            (Value::BigInt(l), Value::BigInt(r)) if *r != 0 => Some(Value::BigInt(l % r)),
341            _ => None,
342        }
343    }
344
345    fn compare_values(&self, left: &Value, right: &Value) -> Option<std::cmp::Ordering> {
346        match (left, right) {
347            (Value::SmallInt(l), Value::SmallInt(r)) => Some(l.cmp(r)),
348            (Value::Int(l), Value::Int(r)) => Some(l.cmp(r)),
349            (Value::BigInt(l), Value::BigInt(r)) => Some(l.cmp(r)),
350            (Value::Float(l), Value::Float(r)) => l.partial_cmp(r),
351            (Value::Double(l), Value::Double(r)) => l.partial_cmp(r),
352            (Value::String(l), Value::String(r)) => Some(l.cmp(r)),
353            (Value::Bool(l), Value::Bool(r)) => Some(l.cmp(r)),
354            _ => None,
355        }
356    }
357
358    fn logical_and(&self, left: &Value, right: &Value) -> Option<Value> {
359        match (left, right) {
360            (Value::Bool(l), Value::Bool(r)) => Some(Value::Bool(*l && *r)),
361            _ => None,
362        }
363    }
364
365    fn logical_or(&self, left: &Value, right: &Value) -> Option<Value> {
366        match (left, right) {
367            (Value::Bool(l), Value::Bool(r)) => Some(Value::Bool(*l || *r)),
368            _ => None,
369        }
370    }
371
372    fn concat_values(&self, left: &Value, right: &Value) -> Option<Value> {
373        match (left, right) {
374            (Value::String(l), Value::String(r)) => Some(Value::string(format!("{}{}", l, r))),
375            _ => None,
376        }
377    }
378
379    /// Estimating the number of iterations of the Loop node
380    ///
381    /// Try to parse the number of iterations from the conditional string. The following patterns are supported:
382    /// - Digital direct value: "10" → 10
383    /// - Range expressions: "1..10" or "range(1, 10)" -> 9
384    /// - Comparison expressions: "i < 10", "i <= 10" → 10
385    /// - Set size: "Items" – Use the set size to make an estimate.
386    ///
387    /// If the content cannot be parsed, use the default values configured in the system.
388    pub fn parse_loop_iterations(&self, condition: &str) -> Option<u32> {
389        let condition = condition.trim();
390
391        // Trying to parse numbers directly
392        if let Ok(num) = condition.parse::<u32>() {
393            return Some(num.max(1));
394        }
395
396        // Try to parse the range expression: 1..10 or 1... =10
397        if condition.contains("..") {
398            return self.parse_range_expression_u32(condition);
399        }
400
401        // Try to parse range(start, end) or range(start, end, step)
402        if condition.starts_with("range(") && condition.ends_with(")") {
403            return self.parse_range_function(condition).map(|v| v as u32);
404        }
405
406        // Try to parse comparison expressions: i < 10, i <= 10, count > 5, etc.
407        if let Some(iterations) = self.parse_comparison_expression(condition) {
408            return Some(iterations as u32);
409        }
410
411        // Try to parse list/set size: [a,b,c] or {a,b,c}
412        if condition.starts_with('[') && condition.ends_with(']') {
413            return Some(self.parse_collection_size(condition));
414        }
415
416        None
417    }
418
419    /// Analyzing array literals
420    fn parse_array_literal(&self, expr: &str) -> Option<f64> {
421        let inner = &expr[1..expr.len() - 1];
422        if inner.trim().is_empty() {
423            return Some(0.0);
424        }
425        // Working with nested arrays and complex expressions
426        // Handling nested arrays and complex expressions
427        let count = self.count_top_level_commas(inner) as f64;
428        Some(count)
429    }
430
431    /// Count the number of commas at the top level (used for processing nested structures)
432    fn count_top_level_commas(&self, s: &str) -> usize {
433        let mut count = 0;
434        let mut depth = 0;
435        let chars = s.chars().peekable();
436
437        for c in chars {
438            match c {
439                '[' | '(' | '{' => depth += 1,
440                ']' | ')' | '}' => depth -= 1,
441                ',' if depth == 0 => count += 1,
442                _ => {}
443            }
444        }
445
446        count + 1 // Number of commas + 1 = Number of elements
447    }
448
449    /// Analyzing the range function
450    fn parse_range_function(&self, expr: &str) -> Option<f64> {
451        let args_str = &expr[6..expr.len() - 1];
452        let args: Vec<&str> = args_str.split(',').map(|s| s.trim()).collect();
453
454        if args.len() >= 2 {
455            let start: i64 = args[0].parse().ok()?;
456            let end: i64 = args[1].parse().ok()?;
457            let step: i64 = if args.len() >= 3 {
458                args[2].parse().ok()?
459            } else {
460                1
461            };
462
463            if step != 0 {
464                let count = ((end - start) / step).abs() as f64;
465                return Some(count.max(0.0));
466            }
467        }
468        None
469    }
470
471    /// Analyzing the range expression
472    fn parse_range_expression(&self, expr: &str) -> Option<f64> {
473        if let Some(pos) = expr.find("..") {
474            let start_str = expr[..pos].trim();
475            let end_part = &expr[pos + 2..];
476
477            let (end_str, inclusive) = if let Some(stripped) = end_part.strip_prefix('=') {
478                (stripped, true)
479            } else {
480                (end_part, false)
481            };
482            if let (Ok(start), Ok(end)) = (start_str.parse::<i64>(), end_str.trim().parse::<i64>())
483            {
484                if end > start {
485                    let count = if inclusive {
486                        end - start + 1
487                    } else {
488                        end - start
489                    };
490                    return Some(count.max(0) as f64);
491                }
492            }
493        }
494        None
495    }
496
497    /// Parse the range expression (return an u32 value)
498    fn parse_range_expression_u32(&self, expr: &str) -> Option<u32> {
499        // Processing 1...10 format (without end)
500        if let Some(pos) = expr.find("..") {
501            let start_str = expr[..pos].trim();
502            let end_part = &expr[pos + 2..];
503
504            // Check for inclusion of the equal sign (1... =10 indicates end of inclusion)
505            let (end_str, inclusive) = if let Some(stripped) = end_part.strip_prefix('=') {
506                (stripped, true)
507            } else {
508                (end_part, false)
509            };
510
511            if let (Ok(start), Ok(end)) = (start_str.parse::<i64>(), end_str.trim().parse::<i64>())
512            {
513                if end > start {
514                    let count = if inclusive {
515                        end - start + 1
516                    } else {
517                        end - start
518                    };
519                    return Some(count.max(1) as u32);
520                }
521            }
522        }
523        None
524    }
525
526    /// Analyzing comparative expressions (such as "i < 10", "count <= 100")
527    fn parse_comparison_expression(&self, expr: &str) -> Option<f64> {
528        // Matching patterns: var < num, var <= num, var > num, var >= num
529        let operators = [("<", 0u32), ("<=", 0u32), (">", 0u32), (">=", 0u32)];
530
531        for (op, _) in &operators {
532            if let Some(pos) = expr.find(op) {
533                let right_side = &expr[pos + op.len()..];
534                if let Ok(num) = right_side.trim().parse::<i64>() {
535                    // For the < operator, the actual number of iterations is num (if num > 0)
536                    let iterations = if num > 0 { num as f64 } else { 1.0 };
537                    // Unable to determine starting value, use conservative estimate
538                    let iterations = iterations + 10.0;
539                    return Some(iterations.max(1.0));
540                }
541            }
542        }
543        None
544    }
545
546    /// Analyzing set functions (keys, values, nodes, relationships)
547    fn parse_collection_function(&self, expr: &str) -> Option<f64> {
548        let expr_lower = expr.to_lowercase();
549
550        // keys(map) or map.keys() - returns the key list of the map
551        if expr_lower.contains(".keys()") {
552            // Unable to determine map size, use default estimate
553            return Some(self.config.default_unwind_list_size);
554        }
555
556        // values(map) or map.values()
557        if expr_lower.starts_with("values(") || expr_lower.contains(".values()") {
558            return Some(self.config.default_unwind_list_size);
559        }
560
561        // nodes(path) - returns the node list in the path
562        if expr_lower.contains(".nodes()") {
563            // Path length unknown, use default estimate
564            return Some(self.config.default_unwind_list_size);
565        }
566
567        // relationships(path) or rels(path) - returns the relationship list in the path
568        if expr_lower.starts_with("relationships(") || expr_lower.starts_with("rels(") {
569            return Some(self.config.default_unwind_list_size - 1.0); // The number of relations is usually 1 less than the number of nodes
570        }
571
572        // labels(node) - returns the label list (usually very small)
573        if expr_lower.starts_with("labels(") {
574            return Some(1.0); // Usually a node has only 1-2 tags
575        }
576
577        None
578    }
579
580    /// Analysis of the string splitting function
581    fn parse_split_function(&self, expr: &str) -> Option<f64> {
582        let expr_lower = expr.to_lowercase();
583
584        // split(string, delimiter) or string.split(delimiter)
585        if expr_lower.contains("split(") || expr_lower.contains(".split(") {
586            // Assume that the average length of each element is 5 characters for string length estimation.
587            // Assuming an average length of 5 characters per element
588            return Some(self.config.default_unwind_list_size);
589        }
590
591        None
592    }
593
594    /// Resolve set size (e.g. "[a, b, c]")
595    fn parse_collection_size(&self, expr: &str) -> u32 {
596        let inner = &expr[1..expr.len() - 1];
597        if inner.trim().is_empty() {
598            return 0;
599        }
600        // Perform a simple calculation: count the number of commas and then add 1.
601        // Simple calculation of the number of commas + 1
602        let count = inner.split(',').count() as u32;
603        count.max(1)
604    }
605}
606
607#[cfg(test)]
608mod tests {
609    use super::*;
610
611    #[test]
612    fn test_parse_array_literal() {
613        let parser = ExpressionParser::new(CostModelConfig::default());
614        assert_eq!(parser.parse_list_size("[1, 2, 3]"), Some(3.0));
615        assert_eq!(parser.parse_list_size("[]"), Some(0.0));
616        assert_eq!(parser.parse_list_size("[a]"), Some(1.0));
617    }
618
619    #[test]
620    fn test_parse_range_function() {
621        let parser = ExpressionParser::new(CostModelConfig::default());
622        assert_eq!(parser.parse_list_size("range(1, 10)"), Some(9.0));
623        assert_eq!(parser.parse_list_size("range(1, 10, 2)"), Some(4.0));
624    }
625
626    #[test]
627    fn test_parse_range_expression() {
628        let parser = ExpressionParser::new(CostModelConfig::default());
629        assert_eq!(parser.parse_list_size("1..10"), Some(9.0));
630        assert_eq!(parser.parse_list_size("0..=5"), Some(6.0));
631    }
632
633    #[test]
634    fn test_parse_loop_iterations_number() {
635        let parser = ExpressionParser::new(CostModelConfig::default());
636        assert_eq!(parser.parse_loop_iterations("10"), Some(10)); // At least once
637        assert_eq!(parser.parse_loop_iterations("0"), Some(1)); // At least 1
638    }
639
640    #[test]
641    fn test_parse_loop_iterations_comparison() {
642        let parser = ExpressionParser::new(CostModelConfig::default());
643        assert_eq!(parser.parse_loop_iterations("i < 10"), Some(20));
644        assert_eq!(parser.parse_loop_iterations("i <= 5"), Some(15));
645    }
646}