From 4f8bcd89d09dc3690c015673013bebdd53f1ead7 Mon Sep 17 00:00:00 2001 From: JoshuaTang <1240604020@qq.com> Date: Sat, 18 Oct 2025 18:42:01 -0700 Subject: [PATCH 1/2] feat(graph): support cypher skip keyword --- rust/lance-graph/src/ast.rs | 2 + rust/lance-graph/src/datafusion_planner.rs | 8 + rust/lance-graph/src/logical_plan.rs | 65 ++++++++ rust/lance-graph/src/parser.rs | 83 ++++++++++- rust/lance-graph/src/query.rs | 164 ++++++++++++++++++++- rust/lance-graph/src/query/clauses.rs | 34 ++++- 6 files changed, 342 insertions(+), 14 deletions(-) diff --git a/rust/lance-graph/src/ast.rs b/rust/lance-graph/src/ast.rs index 4f7e8596d..ff795c303 100644 --- a/rust/lance-graph/src/ast.rs +++ b/rust/lance-graph/src/ast.rs @@ -23,6 +23,8 @@ pub struct CypherQuery { pub limit: Option, /// ORDER BY clause (optional) pub order_by: Option, + /// SKIP/OFFSET clause (optional) + pub skip: Option, } impl CypherQuery { diff --git a/rust/lance-graph/src/datafusion_planner.rs b/rust/lance-graph/src/datafusion_planner.rs index 789b005ff..4b058edd8 100644 --- a/rust/lance-graph/src/datafusion_planner.rs +++ b/rust/lance-graph/src/datafusion_planner.rs @@ -146,6 +146,14 @@ impl DataFusionPlanner { .build() .unwrap()) } + LogicalOperator::Offset { input, offset } => { + let input_plan = self.plan_operator_with_ctx(input, var_labels)?; + Ok(LogicalPlanBuilder::from(input_plan) + .limit((*offset) as usize, None) + .unwrap() + .build() + .unwrap()) + } LogicalOperator::Expand { input, source_variable, diff --git a/rust/lance-graph/src/logical_plan.rs b/rust/lance-graph/src/logical_plan.rs index 6397fcd27..754c53310 100644 --- a/rust/lance-graph/src/logical_plan.rs +++ b/rust/lance-graph/src/logical_plan.rs @@ -72,6 +72,12 @@ pub enum LogicalOperator { sort_items: Vec, }, + /// Apply SKIP/OFFSET + Offset { + input: Box, + offset: u64, + }, + /// Apply LIMIT Limit { input: Box, @@ -147,6 +153,14 @@ impl LogicalPlanner { }; } + // Apply SKIP/OFFSET if present + if let Some(skip) = query.skip { + plan = LogicalOperator::Offset { + input: Box::new(plan), + offset: skip, + }; + } + // Apply LIMIT if present if let Some(limit) = query.limit { plan = LogicalOperator::Limit { @@ -338,6 +352,7 @@ impl LogicalPlanner { LogicalOperator::Project { input, .. } => self.extract_variable_from_plan(input), LogicalOperator::Distinct { input } => self.extract_variable_from_plan(input), LogicalOperator::Sort { input, .. } => self.extract_variable_from_plan(input), + LogicalOperator::Offset { input, .. } => self.extract_variable_from_plan(input), LogicalOperator::Limit { input, .. } => self.extract_variable_from_plan(input), LogicalOperator::Join { left, right, .. } => { // Prefer the right branch's tail variable, else fall back to left @@ -796,6 +811,56 @@ mod tests { } } + #[test] + fn test_order_skip_limit_wrapping() { + // ORDER BY + SKIP + LIMIT should be Limit(Offset(Sort(Project(..)))) + let q = "MATCH (n:Person) RETURN n.name ORDER BY n.name SKIP 5 LIMIT 10"; + let ast = parse_cypher_query(q).unwrap(); + let mut planner = LogicalPlanner::new(); + let logical = planner.plan(&ast).unwrap(); + match logical { + LogicalOperator::Limit { input, count } => { + assert_eq!(count, 10); + match *input { + LogicalOperator::Offset { + input: inner, + offset, + } => { + assert_eq!(offset, 5); + match *inner { + LogicalOperator::Sort { input: inner2, .. } => match *inner2 { + LogicalOperator::Project { .. } => {} + _ => panic!("Expected Project under Sort"), + }, + _ => panic!("Expected Sort under Offset"), + } + } + _ => panic!("Expected Offset under Limit"), + } + } + _ => panic!("Expected Limit at top level"), + } + } + + #[test] + fn test_skip_only_wrapping() { + // SKIP only should be Offset(Project(..)) + let q = "MATCH (n:Person) RETURN n.name SKIP 3"; + let ast = parse_cypher_query(q).unwrap(); + let mut planner = LogicalPlanner::new(); + let logical = planner.plan(&ast).unwrap(); + match logical { + LogicalOperator::Offset { input, offset } => { + assert_eq!(offset, 3); + match *input { + LogicalOperator::Project { .. } => {} + _ => panic!("Expected Project under Offset"), + } + } + _ => panic!("Expected Offset at top level"), + } + } + #[test] fn test_relationship_properties_pushed_into_expand() { let q = "MATCH (a)-[:KNOWS {since: 2020}]->(b) RETURN b"; diff --git a/rust/lance-graph/src/parser.rs b/rust/lance-graph/src/parser.rs index e40dca5bb..0b45bb282 100644 --- a/rust/lance-graph/src/parser.rs +++ b/rust/lance-graph/src/parser.rs @@ -45,7 +45,7 @@ fn cypher_query(input: &str) -> IResult<&str, CypherQuery> { let (input, where_clause) = opt(where_clause)(input)?; let (input, return_clause) = return_clause(input)?; let (input, order_by) = opt(order_by_clause)(input)?; - let (input, limit) = opt(limit_clause)(input)?; + let (input, (skip, limit)) = pagination_clauses(input)?; let (input, _) = multispace0(input)?; Ok(( @@ -56,6 +56,7 @@ fn cypher_query(input: &str) -> IResult<&str, CypherQuery> { return_clause, limit, order_by, + skip, }, )) } @@ -389,6 +390,49 @@ fn limit_clause(input: &str) -> IResult<&str, u64> { Ok((input, limit as u64)) } +// Parse a SKIP clause +fn skip_clause(input: &str) -> IResult<&str, u64> { + let (input, _) = multispace0(input)?; + let (input, _) = tag_no_case("SKIP")(input)?; + let (input, _) = multispace1(input)?; + let (input, skip) = integer_literal(input)?; + + Ok((input, skip as u64)) +} + +// Parse pagination clauses (SKIP and LIMIT) +fn pagination_clauses(input: &str) -> IResult<&str, (Option, Option)> { + let (mut remaining, _) = multispace0(input)?; + let mut skip: Option = None; + let mut limit: Option = None; + + loop { + let before = remaining; + + if skip.is_none() { + if let Ok((i, s)) = skip_clause(remaining) { + skip = Some(s); + remaining = i; + continue; + } + } + + if limit.is_none() { + if let Ok((i, l)) = limit_clause(remaining) { + limit = Some(l); + remaining = i; + continue; + } + } + + if before == remaining { + break; + } + } + + Ok((remaining, (skip, limit))) +} + // Helper parsers // Parse an identifier @@ -572,4 +616,41 @@ mod tests { assert_eq!(result.limit, Some(10)); } + + #[test] + fn test_parse_query_with_skip() { + let query = "MATCH (n:Person) RETURN n.name SKIP 5"; + let result = parse_cypher_query(query).unwrap(); + + assert_eq!(result.skip, Some(5)); + assert_eq!(result.limit, None); + } + + #[test] + fn test_parse_query_with_skip_and_limit() { + let query = "MATCH (n:Person) RETURN n.name SKIP 5 LIMIT 10"; + let result = parse_cypher_query(query).unwrap(); + + assert_eq!(result.skip, Some(5)); + assert_eq!(result.limit, Some(10)); + } + + #[test] + fn test_parse_query_with_skip_and_order_by() { + let query = "MATCH (n:Person) RETURN n.name ORDER BY n.age SKIP 5"; + let result = parse_cypher_query(query).unwrap(); + + assert_eq!(result.skip, Some(5)); + assert!(result.order_by.is_some()); + } + + #[test] + fn test_parse_query_with_skip_order_by_and_limit() { + let query = "MATCH (n:Person) RETURN n.name ORDER BY n.age SKIP 5 LIMIT 10"; + let result = parse_cypher_query(query).unwrap(); + + assert_eq!(result.skip, Some(5)); + assert_eq!(result.limit, Some(10)); + assert!(result.order_by.is_some()); + } } diff --git a/rust/lance-graph/src/query.rs b/rust/lance-graph/src/query.rs index 404811df4..339938c77 100644 --- a/rust/lance-graph/src/query.rs +++ b/rust/lance-graph/src/query.rs @@ -192,14 +192,23 @@ impl CypherQuery { })?; } - // Apply LIMIT if present - if let Some(limit) = self.ast.limit { - df = df - .limit(0, Some(limit as usize)) - .map_err(|e| GraphError::PlanError { - message: format!("Failed to apply LIMIT: {}", e), - location: snafu::Location::new(file!(), line!(), column!()), - })?; + // Apply ORDER BY if present + if let Some(order_by) = &self.ast.order_by { + let sort_expr = to_df_order_by_expr_simple(&order_by.items); + df = df.sort(sort_expr).map_err(|e| GraphError::PlanError { + message: format!("Failed to apply ORDER BY: {}", e), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + } + + // Apply SKIP/OFFSET and LIMIT if present + if self.ast.skip.is_some() || self.ast.limit.is_some() { + let offset = self.ast.skip.unwrap_or(0) as usize; + let fetch = self.ast.limit.map(|l| l as usize); + df = df.limit(offset, fetch).map_err(|e| GraphError::PlanError { + message: format!("Failed to apply SKIP/LIMIT: {}", e), + location: snafu::Location::new(file!(), line!(), column!()), + })?; } // Collect results and concat into a single RecordBatch @@ -1022,6 +1031,7 @@ pub struct CypherQueryBuilder { order_by_items: Vec, limit: Option, distinct: bool, + skip: Option, config: Option, parameters: HashMap, } @@ -1078,6 +1088,12 @@ impl CypherQueryBuilder { self } + /// Add a SKIP clause + pub fn skip(mut self, skip: u64) -> Self { + self.skip = Some(skip); + self + } + /// Build the final CypherQuery pub fn build(self) -> Result { if self.match_clauses.is_empty() { @@ -1111,6 +1127,7 @@ impl CypherQueryBuilder { }) }, limit: self.limit, + skip: self.skip, }; // Generate query text from AST (simplified) @@ -1190,6 +1207,26 @@ fn to_df_boolean_expr_simple( } } +/// Build ORDER BY expressions for simple queries (single table) +fn to_df_order_by_expr_simple( + items: &[crate::ast::OrderByItem], +) -> Vec { + use datafusion::logical_expr::SortExpr; + + items + .iter() + .map(|item| { + let expr = to_df_value_expr_simple(&item.expression); + let asc = matches!(item.direction, crate::ast::SortDirection::Ascending); + SortExpr { + expr, + asc, + nulls_first: false, + } + }) + .collect() +} + fn to_df_value_expr_simple(expr: &crate::ast::ValueExpression) -> datafusion::logical_expr::Expr { use crate::ast::ValueExpression as VE; use datafusion::logical_expr::{col, lit}; @@ -1370,4 +1407,115 @@ mod tests { assert!(got.contains(&"Bob".to_string())); assert!(got.contains(&"Carol".to_string())); } + + #[tokio::test] + async fn test_execute_order_by_asc() { + use arrow_array::{Int64Array, RecordBatch, StringArray}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + + // name, age (int) + let schema = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, true), + Field::new("age", DataType::Int64, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(vec!["Bob", "Alice", "David", "Carol"])), + Arc::new(Int64Array::from(vec![34, 28, 42, 29])), + ], + ) + .unwrap(); + + let cfg = GraphConfig::builder() + .with_node_label("Person", "id") + .build() + .unwrap(); + + // Order ascending by age + let q = CypherQuery::new("MATCH (p:Person) RETURN p.name, p.age ORDER BY p.age ASC") + .unwrap() + .with_config(cfg); + + let mut data = HashMap::new(); + data.insert("people".to_string(), batch); + + let out = q.execute(data).await.unwrap(); + let ages = out.column(1).as_any().downcast_ref::().unwrap(); + let collected: Vec = (0..out.num_rows()).map(|i| ages.value(i)).collect(); + assert_eq!(collected, vec![28, 29, 34, 42]); + } + + #[tokio::test] + async fn test_execute_order_by_desc_with_skip_limit() { + use arrow_array::{Int64Array, RecordBatch, StringArray}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + + let schema = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, true), + Field::new("age", DataType::Int64, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(vec!["Bob", "Alice", "David", "Carol"])), + Arc::new(Int64Array::from(vec![34, 28, 42, 29])), + ], + ) + .unwrap(); + + let cfg = GraphConfig::builder() + .with_node_label("Person", "id") + .build() + .unwrap(); + + // Desc by age, skip 1 (drop 42), take 2 -> [34, 29] + let q = + CypherQuery::new("MATCH (p:Person) RETURN p.age ORDER BY p.age DESC SKIP 1 LIMIT 2") + .unwrap() + .with_config(cfg); + + let mut data = HashMap::new(); + data.insert("people".to_string(), batch); + + let out = q.execute(data).await.unwrap(); + assert_eq!(out.num_rows(), 2); + let ages = out.column(0).as_any().downcast_ref::().unwrap(); + let collected: Vec = (0..out.num_rows()).map(|i| ages.value(i)).collect(); + assert_eq!(collected, vec![34, 29]); + } + + #[tokio::test] + async fn test_execute_skip_without_limit() { + use arrow_array::{Int64Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + + let schema = Arc::new(Schema::new(vec![Field::new("age", DataType::Int64, true)])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(Int64Array::from(vec![10, 20, 30, 40]))], + ) + .unwrap(); + + let cfg = GraphConfig::builder() + .with_node_label("Person", "id") + .build() + .unwrap(); + + let q = CypherQuery::new("MATCH (p:Person) RETURN p.age ORDER BY p.age ASC SKIP 2") + .unwrap() + .with_config(cfg); + + let mut data = HashMap::new(); + data.insert("people".to_string(), batch); + + let out = q.execute(data).await.unwrap(); + assert_eq!(out.num_rows(), 2); + let ages = out.column(0).as_any().downcast_ref::().unwrap(); + let collected: Vec = (0..out.num_rows()).map(|i| ages.value(i)).collect(); + assert_eq!(collected, vec![30, 40]); + } } diff --git a/rust/lance-graph/src/query/clauses.rs b/rust/lance-graph/src/query/clauses.rs index 79d769536..82dd6d5a8 100644 --- a/rust/lance-graph/src/query/clauses.rs +++ b/rust/lance-graph/src/query/clauses.rs @@ -50,13 +50,37 @@ pub(super) fn apply_return_with_qualifier( location: snafu::Location::new(file!(), line!(), column!()), })?; } - if let Some(limit) = ast.limit { - df = df - .limit(0, Some(limit as usize)) - .map_err(|e| GraphError::PlanError { - message: format!("Failed to apply LIMIT: {}", e), + // ORDER BY + if let Some(order_by) = &ast.order_by { + use datafusion::logical_expr::SortExpr; + let mut sorts: Vec = Vec::new(); + for item in &order_by.items { + if let crate::ast::ValueExpression::Property(prop) = &item.expression { + let col_name = qualify(&prop.variable, &prop.property); + let col = datafusion::logical_expr::col(col_name); + let asc = matches!(item.direction, crate::ast::SortDirection::Ascending); + sorts.push(SortExpr { + expr: col, + asc, + nulls_first: false, + }); + } + } + if !sorts.is_empty() { + df = df.sort(sorts).map_err(|e| GraphError::PlanError { + message: format!("Failed to apply ORDER BY: {}", e), location: snafu::Location::new(file!(), line!(), column!()), })?; + } + } + // SKIP/OFFSET and LIMIT + if ast.skip.is_some() || ast.limit.is_some() { + let offset = ast.skip.unwrap_or(0) as usize; + let fetch = ast.limit.map(|l| l as usize); + df = df.limit(offset, fetch).map_err(|e| GraphError::PlanError { + message: format!("Failed to apply SKIP/LIMIT: {}", e), + location: snafu::Location::new(file!(), line!(), column!()), + })?; } Ok(df) } From 85b30efea0086a3e872262d5ccfff7d4acda1fea Mon Sep 17 00:00:00 2001 From: JoshuaTang <1240604020@qq.com> Date: Sat, 18 Oct 2025 18:42:15 -0700 Subject: [PATCH 2/2] docs(readme): update doc with skip support --- rust/lance-graph/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rust/lance-graph/README.md b/rust/lance-graph/README.md index 83b462da2..f405959f7 100644 --- a/rust/lance-graph/README.md +++ b/rust/lance-graph/README.md @@ -98,10 +98,10 @@ A builder (`CypherQueryBuilder`) is also available for constructing queries prog - Node patterns `(:Label)` with optional variables. - Relationship patterns with fixed direction and type, including multi-hop paths. - Property comparisons against literal values with `AND`/`OR`/`NOT`/`EXISTS`. -- RETURN lists of property accesses, optional `DISTINCT`, and `LIMIT`. +- RETURN lists of property accesses, optional `DISTINCT`, `ORDER BY`, `SKIP` (offset), and `LIMIT`. - Positional and named parameters (e.g. `$min_age`). -Features such as ORDER BY, aggregations, optional matches, and subqueries are parsed but not executed yet. +Features such as aggregations, optional matches, and subqueries are parsed but not executed yet. ## Crate Layout