From 1c622f28f14e85861fe4d79dbb31e755a2893c12 Mon Sep 17 00:00:00 2001 From: Jay Pokale Date: Sun, 25 Oct 2026 10:59:24 +0530 Subject: [PATCH] feat(parser): parse arithmetic operators in value expressions The AST, semantic analysis and the DataFusion planner already support ValueExpression::Arithmetic, but the parser never produced it, so queries such as RETURN n.age + 1 or WHERE n.age * 2 > 50 failed to parse. Parse +, -, *, / and % with the usual precedence and left associativity, and allow parenthesized sub-expressions. --- crates/lance-graph/src/parser.rs | 174 +++++++++++++++++- .../tests/test_arithmetic_expressions.rs | 88 +++++++++ 2 files changed, 260 insertions(+), 2 deletions(-) create mode 100644 crates/lance-graph/tests/test_arithmetic_expressions.rs diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index 03a79e58d..cfbcd3078 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -491,9 +491,67 @@ fn basic_value_expression(input: &str) -> IResult<&str, ValueExpression> { ))(input) } -// Parse a value expression -// Optimization: Use peek to avoid expensive backtracking for non-vector queries +// Parse a value expression: `+` and `-` bind looser than `*`, `/` and `%`, and all of them are +// left-associative. fn value_expression(input: &str) -> IResult<&str, ValueExpression> { + let (mut input, mut left) = multiplicative_expression(input)?; + while let Ok((rest, (_, op, _, right))) = tuple(( + multispace0, + one_of("+-"), + multispace0, + multiplicative_expression, + ))(input) + { + let operator = match op { + '+' => ArithmeticOperator::Add, + _ => ArithmeticOperator::Subtract, + }; + left = ValueExpression::Arithmetic { + left: Box::new(left), + operator, + right: Box::new(right), + }; + input = rest; + } + Ok((input, left)) +} + +fn multiplicative_expression(input: &str) -> IResult<&str, ValueExpression> { + let (mut input, mut left) = primary_value_expression(input)?; + while let Ok((rest, (_, op, _, right))) = tuple(( + multispace0, + one_of("*/%"), + multispace0, + primary_value_expression, + ))(input) + { + let operator = match op { + '*' => ArithmeticOperator::Multiply, + '/' => ArithmeticOperator::Divide, + _ => ArithmeticOperator::Modulo, + }; + left = ValueExpression::Arithmetic { + left: Box::new(left), + operator, + right: Box::new(right), + }; + input = rest; + } + Ok((input, left)) +} + +// Parse an operand of an arithmetic expression +// Optimization: Use peek to avoid expensive backtracking for non-vector queries +fn primary_value_expression(input: &str) -> IResult<&str, ValueExpression> { + if let Ok(result) = delimited( + tuple((char('('), multispace0)), + value_expression, + tuple((multispace0, char(')'))), + )(input) + { + return Ok(result); + } + // Peek at first identifier to dispatch to correct parser // This eliminates failed parser attempts for every non-vector expression if let Ok((_, first_ident)) = peek(identifier)(input) { @@ -1277,6 +1335,118 @@ mod tests { assert!(result.order_by.is_some()); } + fn return_expression(query: &str) -> ValueExpression { + let query = parse_cypher_query(query).unwrap(); + query.return_clause.items[0].expression.clone() + } + + fn arith( + left: ValueExpression, + operator: ArithmeticOperator, + right: ValueExpression, + ) -> ValueExpression { + ValueExpression::Arithmetic { + left: Box::new(left), + operator, + right: Box::new(right), + } + } + + fn int(value: i64) -> ValueExpression { + ValueExpression::Literal(PropertyValue::Integer(value)) + } + + #[test] + fn test_parse_arithmetic_precedence() { + // `*` binds tighter than `+`: 1 + (2 * 3) + assert_eq!( + return_expression("MATCH (n:Person) RETURN 1 + 2 * 3"), + arith( + int(1), + ArithmeticOperator::Add, + arith(int(2), ArithmeticOperator::Multiply, int(3)) + ) + ); + } + + #[test] + fn test_parse_arithmetic_left_associative() { + // (10 - 4) - 3, and (8 / 2) % 3 + assert_eq!( + return_expression("MATCH (n:Person) RETURN 10 - 4 - 3"), + arith( + arith(int(10), ArithmeticOperator::Subtract, int(4)), + ArithmeticOperator::Subtract, + int(3) + ) + ); + assert_eq!( + return_expression("MATCH (n:Person) RETURN 8 / 2 % 3"), + arith( + arith(int(8), ArithmeticOperator::Divide, int(2)), + ArithmeticOperator::Modulo, + int(3) + ) + ); + } + + #[test] + fn test_parse_arithmetic_parentheses_and_properties() { + let age = ValueExpression::Property(PropertyRef { + variable: "n".to_string(), + property: "age".to_string(), + }); + assert_eq!( + return_expression("MATCH (n:Person) RETURN (n.age + 1) * 2"), + arith( + arith(age.clone(), ArithmeticOperator::Add, int(1)), + ArithmeticOperator::Multiply, + int(2) + ) + ); + // Without spaces, `-1` is still a subtraction of the literal 1 + assert_eq!( + return_expression("MATCH (n:Person) RETURN n.age-1"), + arith(age, ArithmeticOperator::Subtract, int(1)) + ); + } + + #[test] + fn test_parse_arithmetic_in_where_comparison() { + let query = + parse_cypher_query("MATCH (n:Person) WHERE n.age + 1 > 30 RETURN n.name").unwrap(); + let where_clause = query.where_clause.unwrap(); + match where_clause.expression { + BooleanExpression::Comparison { + left, + operator, + right, + } => { + assert!(matches!( + left, + ValueExpression::Arithmetic { + operator: ArithmeticOperator::Add, + .. + } + )); + assert_eq!(operator, ComparisonOperator::GreaterThan); + assert_eq!(right, int(30)); + } + other => panic!("expected a comparison, got {other:?}"), + } + } + + #[test] + fn test_parse_count_star_still_parses_with_arithmetic() { + assert!(matches!( + return_expression("MATCH (n:Person) RETURN count(*) * 2"), + ValueExpression::Arithmetic { + operator: ArithmeticOperator::Multiply, + .. + } + )); + } + #[test] fn test_parse_count_star() { let query = "MATCH (n:Person) RETURN count(*) AS total"; diff --git a/crates/lance-graph/tests/test_arithmetic_expressions.rs b/crates/lance-graph/tests/test_arithmetic_expressions.rs new file mode 100644 index 000000000..b835082d6 --- /dev/null +++ b/crates/lance-graph/tests/test_arithmetic_expressions.rs @@ -0,0 +1,88 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! End-to-end tests for arithmetic expressions (`+`, `-`, `*`, `/`, `%`) in RETURN and WHERE. + +use arrow_array::{Int64Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use lance_graph::config::GraphConfig; +use lance_graph::{CypherQuery, ExecutionStrategy, NodeMapping}; +use std::collections::HashMap; +use std::sync::Arc; + +// | id | name | age | +// |----|---------|-----| +// | 1 | Bob | 25 | +// | 2 | Alice | 30 | +// | 3 | Charlie | 35 | +fn person_datasets() -> HashMap { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + Field::new("age", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["Bob", "Alice", "Charlie"])), + Arc::new(Int64Array::from(vec![25, 30, 35])), + ], + ) + .unwrap(); + HashMap::from([("Person".to_string(), batch)]) +} + +fn config() -> GraphConfig { + GraphConfig::builder() + .with_node_mapping(NodeMapping { + label: "Person".to_string(), + id_field: "id".to_string(), + property_fields: vec!["name".to_string(), "age".to_string()], + filter_conditions: None, + }) + .build() + .unwrap() +} + +async fn run(query: &str) -> RecordBatch { + CypherQuery::new(query) + .unwrap() + .with_config(config()) + .execute(person_datasets(), Some(ExecutionStrategy::DataFusion)) + .await + .unwrap() +} + +fn int_column(batch: &RecordBatch, name: &str) -> Vec { + batch + .column_by_name(name) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() +} + +#[tokio::test] +async fn test_return_arithmetic_with_precedence() { + let result = + run("MATCH (p:Person) RETURN p.age + p.id * 10 AS score, (p.age - 5) % 7 AS bucket ORDER BY p.id").await; + assert_eq!(int_column(&result, "score"), vec![35, 50, 65]); + assert_eq!(int_column(&result, "bucket"), vec![6, 4, 2]); +} + +#[tokio::test] +async fn test_where_arithmetic_comparison() { + let result = + run("MATCH (p:Person) WHERE p.age / 5 - 1 >= 5 RETURN p.name ORDER BY p.name").await; + let names = result + .column_by_name("p.name") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let names: Vec<&str> = names.iter().map(|n| n.unwrap()).collect(); + assert_eq!(names, vec!["Alice", "Charlie"]); +}