feat(graph): add WITH clause with query chaining support

Allen Cheng · Allen Cheng · commit 6c523e23398f · 2026-01-13T22:24:45.000-08:00
Adds WITH clause for intermediate projections, aggregations, and
query chaining in Cypher queries.

Supported syntax:
- WITH projection: WITH p.name AS name
- WITH aggregation: WITH city, count(*) AS total
- WITH ORDER BY/LIMIT: WITH p ORDER BY p.age DESC LIMIT 10
- Post-WITH WHERE: WITH ... WHERE total &gt; 1
- Post-WITH MATCH: WITH ... MATCH (p2:Person) ...

Changes:
- Add WithClause to AST with items, order_by, limit fields
- Add with_clause, post_with_match_clauses, post_with_where_clause to CypherQuery
- Parse WITH clause and post-WITH MATCH/WHERE in parser
- Add semantic analysis for WITH scope
- Add plan_with_clause in logical planner
- Chain post-WITH MATCH using plan_match_clause_with_base
- Add 5 comprehensive tests for WITH functionality

Note: WITH p (passing whole node) then MATCH (p)-[]-&gt;(f) requires
explicit property projection. Use WITH p.id AS id ... instead.
diff --git a/rust/lance-graph/src/ast.rs b/rust/lance-graph/src/ast.rs
@@ -15,8 +15,14 @@ use std::collections::HashMap;
 pub struct CypherQuery {
     /// MATCH clauses
     pub match_clauses: Vec<MatchClause>,
-    /// WHERE clause (optional)
+    /// WHERE clause (optional, before WITH if present)
     pub where_clause: Option<WhereClause>,
+    /// WITH clause (optional) - intermediate projection/aggregation
+    pub with_clause: Option<WithClause>,
+    /// MATCH clauses after WITH (optional) - query chaining
+    pub post_with_match_clauses: Vec<MatchClause>,
+    /// WHERE clause after WITH (optional) - filters the WITH results
+    pub post_with_where_clause: Option<WhereClause>,
     /// RETURN clause
     pub return_clause: ReturnClause,
     /// LIMIT clause (optional)
@@ -323,6 +329,20 @@ pub enum ArithmeticOperator {
     Modulo,
 }
 
+/// WITH clause for intermediate projections/aggregations
+///
+/// WITH acts as a query stage boundary, projecting results that become
+/// the input for subsequent clauses. 
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct WithClause {
+    /// Items to project (similar to RETURN)
+    pub items: Vec<ReturnItem>,
+    /// Optional ORDER BY within WITH
+    pub order_by: Option<OrderByClause>,
+    /// Optional LIMIT within WITH
+    pub limit: Option<u64>,
+}
+
 /// RETURN clause specifying what to return
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct ReturnClause {
diff --git a/rust/lance-graph/src/logical_plan.rs b/rust/lance-graph/src/logical_plan.rs
@@ -156,14 +156,32 @@ impl LogicalPlanner {
         // Start with the MATCH clause(s)
         let mut plan = self.plan_match_clauses(&query.match_clauses)?;
 
-        // Apply WHERE clause if present
+        // Apply WHERE clause if present (before WITH)
         if let Some(where_clause) = &query.where_clause {
             plan = LogicalOperator::Filter {
                 input: Box::new(plan),
                 predicate: where_clause.expression.clone(),
             };
         }
 
+        // Apply WITH clause if present (intermediate projection/aggregation)
+        if let Some(with_clause) = &query.with_clause {
+            plan = self.plan_with_clause(with_clause, plan)?;
+        }
+
+        // Apply post-WITH MATCH clauses if present (query chaining)
+        for match_clause in &query.post_with_match_clauses {
+            plan = self.plan_match_clause_with_base(Some(plan), match_clause)?;
+        }
+
+        // Apply post-WITH WHERE clause if present
+        if let Some(post_where) = &query.post_with_where_clause {
+            plan = LogicalOperator::Filter {
+                input: Box::new(plan),
+                predicate: post_where.expression.clone(),
+            };
+        }
+
         // Apply RETURN clause
         plan = self.plan_return_clause(&query.return_clause, plan)?;
 
@@ -429,6 +447,53 @@ impl LogicalPlanner {
 
         Ok(plan)
     }
+
+    /// Plan WITH clause - intermediate projection/aggregation with optional ORDER BY and LIMIT
+    fn plan_with_clause(
+        &self,
+        with_clause: &WithClause,
+        input: LogicalOperator,
+    ) -> Result<LogicalOperator> {
+        // WITH creates a projection (like RETURN)
+        let projections = with_clause
+            .items
+            .iter()
+            .map(|item| ProjectionItem {
+                expression: item.expression.clone(),
+                alias: item.alias.clone(),
+            })
+            .collect();
+
+        let mut plan = LogicalOperator::Project {
+            input: Box::new(input),
+            projections,
+        };
+
+        // Apply ORDER BY within WITH if present
+        if let Some(order_by) = &with_clause.order_by {
+            plan = LogicalOperator::Sort {
+                input: Box::new(plan),
+                sort_items: order_by
+                    .items
+                    .iter()
+                    .map(|item| SortItem {
+                        expression: item.expression.clone(),
+                        direction: item.direction.clone(),
+                    })
+                    .collect(),
+            };
+        }
+
+        // Apply LIMIT within WITH if present
+        if let Some(limit) = with_clause.limit {
+            plan = LogicalOperator::Limit {
+                input: Box::new(plan),
+                count: limit,
+            };
+        }
+
+        Ok(plan)
+    }
 }
 
 impl Default for LogicalPlanner {
diff --git a/rust/lance-graph/src/parser.rs b/rust/lance-graph/src/parser.rs
@@ -42,7 +42,20 @@ pub fn parse_cypher_query(input: &str) -> Result<CypherQuery> {
 fn cypher_query(input: &str) -> IResult<&str, CypherQuery> {
     let (input, _) = multispace0(input)?;
     let (input, match_clauses) = many0(match_clause)(input)?;
-    let (input, where_clause) = opt(where_clause)(input)?;
+    let (input, pre_with_where) = opt(where_clause)(input)?;
+
+    // Optional WITH clause with optional post-WITH MATCH and WHERE
+    let (input, with_result) = opt(with_clause)(input)?;
+    // Only try to parse post-WITH clauses if we have a WITH clause
+    let (input, post_with_matches, post_with_where) = match with_result {
+        Some(_) => {
+            let (input, matches) = many0(match_clause)(input)?;
+            let (input, where_cl) = opt(where_clause)(input)?;
+            (input, matches, where_cl)
+        }
+        None => (input, vec![], None),
+    };
+
     let (input, return_clause) = return_clause(input)?;
     let (input, order_by) = opt(order_by_clause)(input)?;
     let (input, (skip, limit)) = pagination_clauses(input)?;
@@ -52,7 +65,10 @@ fn cypher_query(input: &str) -> IResult<&str, CypherQuery> {
         input,
         CypherQuery {
             match_clauses,
-            where_clause,
+            where_clause: pre_with_where,
+            with_clause: with_result,
+            post_with_match_clauses: post_with_matches,
+            post_with_where_clause: post_with_where,
             return_clause,
             limit,
             order_by,
@@ -657,6 +673,26 @@ fn property_reference(input: &str) -> IResult<&str, PropertyRef> {
     ))
 }
 
+
+// Parse a WITH clause (intermediate projection/aggregation)
+fn with_clause(input: &str) -> IResult<&str, WithClause> {
+    let (input, _) = multispace0(input)?;
+    let (input, _) = tag_no_case("WITH")(input)?;
+    let (input, _) = multispace1(input)?;
+    let (input, items) = separated_list0(comma_ws, return_item)(input)?;
+    let (input, order_by) = opt(order_by_clause)(input)?;
+    let (input, limit) = opt(limit_clause)(input)?;
+
+    Ok((
+        input,
+        WithClause {
+            items,
+            order_by,
+            limit,
+        },
+    ))
+}
+
 // Parse a RETURN clause
 fn return_clause(input: &str) -> IResult<&str, ReturnClause> {
     let (input, _) = multispace0(input)?;
diff --git a/rust/lance-graph/src/query.rs b/rust/lance-graph/src/query.rs
@@ -1178,6 +1178,9 @@ impl CypherQueryBuilder {
             where_clause: self
                 .where_expression
                 .map(|expr| crate::ast::WhereClause { expression: expr }),
+            with_clause: None, // WITH not supported via builder yet
+            post_with_match_clauses: vec![],
+            post_with_where_clause: None,
             return_clause: crate::ast::ReturnClause {
                 distinct: self.distinct,
                 items: self.return_items,
diff --git a/rust/lance-graph/src/semantic.rs b/rust/lance-graph/src/semantic.rs
@@ -44,6 +44,8 @@ pub enum VariableType {
 pub enum ScopeType {
     Match,
     Where,
+    With,
+    PostWithWhere,
     Return,
     OrderBy,
 }
@@ -79,32 +81,48 @@ impl SemanticAnalyzer {
             }
         }
 
-        // Phase 2: Validate WHERE clause
+        // Phase 2: Validate WHERE clause (before WITH)
         if let Some(where_clause) = &query.where_clause {
             self.current_scope = ScopeType::Where;
             if let Err(e) = self.analyze_where_clause(where_clause) {
                 errors.push(format!("WHERE clause error: {}", e));
             }
         }
 
-        // Phase 3: Validate RETURN clause
+        // Phase 3: Validate WITH clause if present
+        if let Some(with_clause) = &query.with_clause {
+            self.current_scope = ScopeType::With;
+            if let Err(e) = self.analyze_with_clause(with_clause) {
+                errors.push(format!("WITH clause error: {}", e));
+            }
+        }
+
+        // Phase 4: Validate post-WITH WHERE clause if present
+        if let Some(post_where) = &query.post_with_where_clause {
+            self.current_scope = ScopeType::PostWithWhere;
+            if let Err(e) = self.analyze_where_clause(post_where) {
+                errors.push(format!("Post-WITH WHERE clause error: {}", e));
+            }
+        }
+
+        // Phase 5: Validate RETURN clause
         self.current_scope = ScopeType::Return;
         if let Err(e) = self.analyze_return_clause(&query.return_clause) {
             errors.push(format!("RETURN clause error: {}", e));
         }
 
-        // Phase 4: Validate ORDER BY clause
+        // Phase 6: Validate ORDER BY clause
         if let Some(order_by) = &query.order_by {
             self.current_scope = ScopeType::OrderBy;
             if let Err(e) = self.analyze_order_by_clause(order_by) {
                 errors.push(format!("ORDER BY clause error: {}", e));
             }
         }
 
-        // Phase 5: Schema validation
+        // Phase 7: Schema validation
         self.validate_schema(&mut warnings);
 
-        // Phase 6: Type checking
+        // Phase 8: Type checking
         self.validate_types(&mut errors);
 
         Ok(SemanticResult {
@@ -416,6 +434,21 @@ impl SemanticAnalyzer {
         Ok(())
     }
 
+    /// Analyze WITH clause
+    fn analyze_with_clause(&mut self, with_clause: &WithClause) -> Result<()> {
+        // Validate WITH item expressions (similar to RETURN)
+        for item in &with_clause.items {
+            self.analyze_value_expression(&item.expression)?;
+        }
+        // Validate ORDER BY within WITH if present
+        if let Some(order_by) = &with_clause.order_by {
+            for item in &order_by.items {
+                self.analyze_value_expression(&item.expression)?;
+            }
+        }
+        Ok(())
+    }
+
     /// Analyze ORDER BY clause
     fn analyze_order_by_clause(&mut self, order_by: &OrderByClause) -> Result<()> {
         for item in &order_by.items {
@@ -558,6 +591,9 @@ mod tests {
         let query = CypherQuery {
             match_clauses: vec![],
             where_clause: None,
+            with_clause: None,
+            post_with_match_clauses: vec![],
+            post_with_where_clause: None,
             return_clause: ReturnClause {
                 distinct: false,
                 items: vec![ReturnItem {
@@ -585,6 +621,9 @@ mod tests {
                 patterns: vec![GraphPattern::Node(node)],
             }],
             where_clause: None,
+            with_clause: None,
+            post_with_match_clauses: vec![],
+            post_with_where_clause: None,
             return_clause: ReturnClause {
                 distinct: false,
                 items: vec![ReturnItem {
@@ -615,6 +654,9 @@ mod tests {
                 patterns: vec![GraphPattern::Node(node1), GraphPattern::Node(node2)],
             }],
             where_clause: None,
+            with_clause: None,
+            post_with_match_clauses: vec![],
+            post_with_where_clause: None,
             return_clause: ReturnClause {
                 distinct: false,
                 items: vec![],
@@ -661,6 +703,9 @@ mod tests {
                 patterns: vec![GraphPattern::Path(path)],
             }],
             where_clause: None,
+            with_clause: None,
+            post_with_match_clauses: vec![],
+            post_with_where_clause: None,
             return_clause: ReturnClause {
                 distinct: false,
                 items: vec![],
@@ -690,6 +735,9 @@ mod tests {
                 patterns: vec![GraphPattern::Node(node)],
             }],
             where_clause: Some(where_clause),
+            with_clause: None,
+            post_with_match_clauses: vec![],
+            post_with_where_clause: None,
             return_clause: ReturnClause {
                 distinct: false,
                 items: vec![],
@@ -729,6 +777,9 @@ mod tests {
                 patterns: vec![GraphPattern::Path(path)],
             }],
             where_clause: None,
+            with_clause: None,
+            post_with_match_clauses: vec![],
+            post_with_where_clause: None,
             return_clause: ReturnClause {
                 distinct: false,
                 items: vec![],
@@ -755,6 +806,9 @@ mod tests {
                 patterns: vec![GraphPattern::Node(node)],
             }],
             where_clause: None,
+            with_clause: None,
+            post_with_match_clauses: vec![],
+            post_with_where_clause: None,
             return_clause: ReturnClause {
                 distinct: false,
                 items: vec![],
@@ -792,6 +846,9 @@ mod tests {
                 patterns: vec![GraphPattern::Node(node)],
             }],
             where_clause: None,
+            with_clause: None,
+            post_with_match_clauses: vec![],
+            post_with_where_clause: None,
             return_clause: ReturnClause {
                 distinct: false,
                 items: vec![],
@@ -834,6 +891,9 @@ mod tests {
                 patterns: vec![GraphPattern::Path(path)],
             }],
             where_clause: None,
+            with_clause: None,
+            post_with_match_clauses: vec![],
+            post_with_where_clause: None,
             return_clause: ReturnClause {
                 distinct: false,
                 items: vec![],
@@ -898,6 +958,9 @@ mod tests {
                 patterns: vec![GraphPattern::Path(path)],
             }],
             where_clause: None,
+            with_clause: None,
+            post_with_match_clauses: vec![],
+            post_with_where_clause: None,
             return_clause: ReturnClause {
                 distinct: false,
                 items: vec![],
diff --git a/rust/lance-graph/tests/test_datafusion_pipeline.rs b/rust/lance-graph/tests/test_datafusion_pipeline.rs
