From 23e99d568fc8fcebeb9788087557200fdc787eed Mon Sep 17 00:00:00 2001 From: prrao87 Date: Sat, 24 Oct 2026 09:11:32 -0400 Subject: [PATCH 1/5] feat(cypher): nodes as values, node lists, UNWIND, parameters and inferred labels Run more of the official LDBC IC Cypher unmodified: node comparison and IN on ids, collect(node) and UNWIND back to nodes, IN over any list expression, chained comparisons, inline properties that read variables, labels inferred from relationship endpoints, mixed-type list literals, aggregates nested in expressions, and $parameters on CypherEngine. --- crates/lance-graph-python/src/graph.rs | 50 +++- crates/lance-graph/src/ast.rs | 9 + .../src/datafusion_planner/analysis.rs | 11 +- .../builder/aggregate_ops.rs | 71 +++++- .../datafusion_planner/builder/basic_ops.rs | 43 +++- .../datafusion_planner/builder/expand_ops.rs | 10 +- .../src/datafusion_planner/builder/helpers.rs | 10 +- .../src/datafusion_planner/builder/mod.rs | 3 +- .../src/datafusion_planner/cypher_udf.rs | 64 +++++ .../src/datafusion_planner/expression.rs | 11 +- .../lance-graph/src/datafusion_planner/mod.rs | 1 + crates/lance-graph/src/logical_plan.rs | 240 +++++++++++++++--- .../lance-graph/src/parameter_substitution.rs | 4 + crates/lance-graph/src/parser.rs | 119 ++++++++- .../lance-graph/src/pattern_alternatives.rs | 117 +++++++++ crates/lance-graph/src/pattern_order.rs | 25 +- crates/lance-graph/src/semantic.rs | 47 +++- python/python/tests/test_engine_values.py | 135 ++++++++++ 18 files changed, 888 insertions(+), 82 deletions(-) create mode 100644 crates/lance-graph/src/datafusion_planner/cypher_udf.rs create mode 100644 python/python/tests/test_engine_values.py diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 653dbc6..9c973c9 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -824,6 +824,12 @@ fn python_to_json(value: &Bound<'_, PyAny>) -> PyResult { )) } else if let Ok(s) = value.extract::() { Ok(JsonValue::String(s)) + } else if value.hasattr("isoformat")? { + // datetime and date: ISO 8601 text, which DataFusion casts when comparing with + // timestamp and date columns. + Ok(JsonValue::String( + value.call_method0("isoformat")?.extract::()?, + )) } else { Err(PyValueError::new_err("Unsupported parameter type")) } @@ -1877,9 +1883,19 @@ impl CypherEngine { .clone() } - fn parse(&self, query: &str, state: &EngineState) -> PyResult { + fn parse( + &self, + query: &str, + params: Option<&Bound<'_, PyDict>>, + state: &EngineState, + ) -> PyResult { + let mut parameters = HashMap::new(); + for (key, value) in params.into_iter().flat_map(|p| p.iter()) { + parameters.insert(key.extract::()?, python_to_json(&value)?); + } let query = RustCypherQuery::new(query) .map_err(graph_error_to_pyerr)? + .with_parameters(parameters) .with_config(self.config.clone()) .with_fan_out(state.fan_out()); Ok(if self.use_csr { @@ -1945,9 +1961,15 @@ impl CypherEngine { /// If the query is invalid /// RuntimeError /// If query execution fails - fn execute(&self, py: Python, query: &str) -> PyResult> { + #[pyo3(signature = (query, params = None))] + fn execute( + &self, + py: Python, + query: &str, + params: Option<&Bound<'_, PyDict>>, + ) -> PyResult> { let state = self.snapshot(); - let cypher_query = self.parse(query, &state)?; + let cypher_query = self.parse(query, params, &state)?; let result_batch = RT .block_on( Some(py), @@ -1965,11 +1987,17 @@ impl CypherEngine { /// stream with bounded memory; joins, sorts and aggregations may still buffer. /// The reader keeps its snapshot alive, so a concurrent ``refresh()`` does not /// affect it. Planning errors raise immediately; execution errors raise while reading. - fn execute_batches(&self, py: Python, query: &str) -> PyResult> { + #[pyo3(signature = (query, params = None))] + fn execute_batches( + &self, + py: Python, + query: &str, + params: Option<&Bound<'_, PyDict>>, + ) -> PyResult> { use arrow::pyarrow::IntoPyArrow; let state = self.snapshot(); - let cypher_query = self.parse(query, &state)?; + let cypher_query = self.parse(query, params, &state)?; let stream = RT .block_on( Some(py), @@ -1987,9 +2015,15 @@ impl CypherEngine { } /// Return the query plan, including the physical plan with pushed-down scans. - fn explain(&self, py: Python, query: &str) -> PyResult { + #[pyo3(signature = (query, params = None))] + fn explain( + &self, + py: Python, + query: &str, + params: Option<&Bound<'_, PyDict>>, + ) -> PyResult { let state = self.snapshot(); - let cypher_query = self.parse(query, &state)?; + let cypher_query = self.parse(query, params, &state)?; RT.block_on( Some(py), cypher_query @@ -2111,7 +2145,7 @@ impl CypherEngine { vector_search: &VectorSearch, ) -> PyResult> { let state = self.snapshot(); - let cypher_query = self.parse(query, &state)?; + let cypher_query = self.parse(query, None, &state)?; let vs = vector_search.inner.clone(); // Execute query to get candidates, then apply vector reranking diff --git a/crates/lance-graph/src/ast.rs b/crates/lance-graph/src/ast.rs index 83641e6..b0de0e4 100644 --- a/crates/lance-graph/src/ast.rs +++ b/crates/lance-graph/src/ast.rs @@ -238,6 +238,9 @@ pub enum PropertyValue { Parameter(String), /// Property reference (e.g., node.property) Property(PropertyRef), + /// A variable from an earlier clause, as an inline property value (`{id: tagId}`). + /// Like `Property` values, it is moved into the clause's WHERE before planning. + Variable(String), } /// Reference to a property of a node or relationship @@ -284,6 +287,12 @@ pub enum BooleanExpression { expression: ValueExpression, list: Vec, }, + /// `x IN list` where the list is any expression, such as a collected alias + /// (`friend IN friends`) rather than a literal `[...]`. + Member { + expression: ValueExpression, + list: ValueExpression, + }, /// LIKE pattern matching (case-sensitive) Like { expression: ValueExpression, diff --git a/crates/lance-graph/src/datafusion_planner/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index 0d2a7b7..0da47f7 100644 --- a/crates/lance-graph/src/datafusion_planner/analysis.rs +++ b/crates/lance-graph/src/datafusion_planner/analysis.rs @@ -206,10 +206,19 @@ fn analyze_operator( analyze_operator(left, analysis, rel_counter)?; analyze_operator(right, analysis, rel_counter)?; } - LogicalOperator::Unwind { input, .. } => { + LogicalOperator::Unwind { + input, + alias, + label, + .. + } => { if let Some(op) = input { analyze_operator(op, analysis, rel_counter)?; } + if let Some(label) = label { + analysis.var_to_label.insert(alias.clone(), label.clone()); + analysis.required_datasets.insert(label.clone()); + } } LogicalOperator::ShortestPath { input, .. } => { analyze_operator(input, analysis, rel_counter)?; diff --git a/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs index 9528106..781e7fa 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs @@ -3,11 +3,55 @@ //! Aggregation operations: Projection with aggregates and grouping +use crate::ast::ValueExpression; use crate::datafusion_planner::DataFusionPlanner; use crate::error::Result; use crate::logical_plan::*; use datafusion::logical_expr::{col, LogicalPlan, LogicalPlanBuilder}; +/// Replace each aggregate inside `expr` with a reference to a new column `__agg_`, +/// pushing the aggregate to `aggregates`, so `head(collect(x))` or `count(x) + 1` aggregate +/// first and apply the outer expression after. Cypher's `collect` gives `[]`, not null, +/// when nothing was collected. +fn split_aggregates(expr: &mut ValueExpression, aggregates: &mut Vec<(String, ValueExpression)>) { + match expr { + ValueExpression::AggregateFunction { name, .. } => { + let column = format!("__agg_{}", aggregates.len()); + let is_collect = name.eq_ignore_ascii_case("collect"); + let aggregate = std::mem::replace(expr, ValueExpression::Variable(column.clone())); + aggregates.push((column.clone(), aggregate)); + if is_collect { + *expr = ValueExpression::ScalarFunction { + name: "coalesce".to_string(), + args: vec![ + ValueExpression::Variable(column), + ValueExpression::ScalarFunction { + name: "__list".to_string(), + args: vec![], + }, + ], + }; + } + } + ValueExpression::ScalarFunction { args, .. } => args + .iter_mut() + .for_each(|a| split_aggregates(a, aggregates)), + ValueExpression::Arithmetic { left, right, .. } => { + split_aggregates(left, aggregates); + split_aggregates(right, aggregates); + } + ValueExpression::Case { otherwise, whens } => { + for (_, then) in whens { + split_aggregates(then, aggregates); + } + if let Some(e) = otherwise { + split_aggregates(e, aggregates); + } + } + _ => {} + } +} + impl DataFusionPlanner { pub(crate) fn build_project_with_aggregates( &self, @@ -17,8 +61,9 @@ impl DataFusionPlanner { // Separate group expressions (non-aggregates) from aggregate expressions let mut group_exprs = Vec::new(); let mut agg_exprs = Vec::new(); - // Store computed aliases for aggregates to reuse in final projection - let mut agg_aliases = Vec::new(); + // Each aggregating item, with its aggregates replaced by `__agg_` columns. + let mut outer = Vec::new(); + let mut aggregates = Vec::new(); let schema = input_plan.schema().clone(); for p in projections { @@ -30,22 +75,25 @@ impl DataFusionPlanner { })); continue; } - let expr = super::super::expression::to_df_value_expr(&p.expression); - if super::super::expression::contains_aggregate(&p.expression) { - // Aggregate expressions get aliased. Lowercase, like plain projections: the - // final projection refers to it with `col`, which normalizes identifiers, so a - // camelCase alias (`count(*) AS postCount`) failed to plan. + // Aliases are lowercase, like plain projections: the final projection refers + // to columns with `col`, which normalizes identifiers. let alias = if let Some(alias) = &p.alias { alias.to_lowercase() } else { super::super::expression::to_cypher_column_name(&p.expression) }; - agg_exprs.push(expr.alias(&alias)); - agg_aliases.push(alias); + let mut rewritten = p.expression.clone(); + let first = aggregates.len(); + split_aggregates(&mut rewritten, &mut aggregates); + for (column, aggregate) in &aggregates[first..] { + let aggregate = super::super::expression::to_df_value_expr(aggregate); + agg_exprs.push(aggregate.alias(column)); + } + outer.push(super::super::expression::to_df_value_expr(&rewritten).alias(alias)); } else { // Group expressions: use raw expression for grouping, no alias - group_exprs.push(expr); + group_exprs.push(super::super::expression::to_df_value_expr(&p.expression)); } } @@ -77,8 +125,7 @@ impl DataFusionPlanner { }; final_projection.push(aliased); } else { - // For aggregates, reference the column using the same alias we computed earlier - final_projection.push(col(&agg_aliases[agg_idx])); + final_projection.push(outer[agg_idx].clone()); agg_idx += 1; } } diff --git a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs index 886be8e..6697c8b 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs @@ -256,6 +256,7 @@ impl DataFusionPlanner { input: &Option>, expression: &crate::ast::ValueExpression, alias: &str, + label: Option<&str>, ) -> Result { let input_plan = if let Some(input_op) = input { self.build_operator(ctx, input_op)? @@ -268,6 +269,13 @@ impl DataFusionPlanner { // Convert expression to DataFusion Expr let df_expr = super::super::expression::to_df_value_expr(expression); + // A list of node ids: unnest under a temporary name, then join each id back to the + // node table so the alias carries the node's columns like any matched node. + let unnested = match label { + Some(_) => format!("__unwind_{}", alias.to_lowercase()), + None => alias.to_string(), + }; + let alias = unnested.as_str(); // We project the list expression first (aliased as the target alias temporarily) // DataFusion unnest takes a column name. @@ -284,7 +292,40 @@ impl DataFusionPlanner { .unnest_column(alias) .map_err(|e| self.plan_error("Failed to build unnest", e))?; - builder + let Some(label) = label else { + return builder + .build() + .map_err(|e| self.plan_error("Failed to build unwind plan", e)); + }; + let variable = &alias["__unwind_".len()..]; + let mapping = self.config.get_node_mapping(label).ok_or_else(|| { + crate::error::GraphError::PlanError { + message: format!("No node mapping for '{label}'"), + location: snafu::Location::new(file!(), line!(), column!()), + } + })?; + let nodes = self.build_scan(ctx, variable, label, &Default::default())?; + let id = crate::case_insensitive::qualify_column(variable, &mapping.id_field); + let joined = builder + .join( + nodes, + datafusion::logical_expr::JoinType::Inner, + (vec![alias.to_string()], vec![id]), + None, + ) + .map_err(|e| self.plan_error("Failed to join unwound ids to nodes", e))? + .build() + .map_err(|e| self.plan_error("Failed to build unwind plan", e))?; + let kept: Vec = joined + .schema() + .columns() + .into_iter() + .filter(|c| c.name != alias) + .map(datafusion::logical_expr::Expr::Column) + .collect(); + LogicalPlanBuilder::from(joined) + .project(kept) + .map_err(|e| self.plan_error("Failed to project unwound nodes", e))? .build() .map_err(|e| self.plan_error("Failed to build unwind plan", e)) } diff --git a/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs index ee37d37..3e1f981 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs @@ -454,8 +454,16 @@ impl DataFusionPlanner { // Derive expected column names from source and target node schemas // This ensures we only project columns that actually belong to source/target nodes - let expected_columns = + let mut expected_columns = self.get_expected_varlength_columns(ctx, source_variable, target_variable)?; + // Everything the input carries (other nodes, scalars from a WITH) passes through. + expected_columns.extend( + input_plan + .schema() + .fields() + .iter() + .map(|f| f.name().clone()), + ); // Generate a plan for each hop count and UNION them let mut plans = Vec::new(); diff --git a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs index bf34c26..f813d3c 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs @@ -153,10 +153,18 @@ impl DataFusionPlanner { Self::collect_labels(left, labels); Self::collect_labels(right, labels); } - LogicalOperator::Unwind { input, .. } => { + LogicalOperator::Unwind { + input, + alias, + label, + .. + } => { if let Some(op) = input { Self::collect_labels(op, labels); } + if let Some(label) = label { + labels.insert(alias.clone(), label.clone()); + } } LogicalOperator::OptionalMatch { optional, .. } => { Self::collect_labels(optional, labels) diff --git a/crates/lance-graph/src/datafusion_planner/builder/mod.rs b/crates/lance-graph/src/datafusion_planner/builder/mod.rs index ddc93af..2068a1c 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/mod.rs @@ -101,7 +101,8 @@ impl DataFusionPlanner { input, expression, alias, - } => self.build_unwind(ctx, input, expression, alias), + label, + } => self.build_unwind(ctx, input, expression, alias, label.as_deref()), LogicalOperator::Union { inputs } => self.build_union(ctx, inputs), LogicalOperator::OptionalMatch { input, diff --git a/crates/lance-graph/src/datafusion_planner/cypher_udf.rs b/crates/lance-graph/src/datafusion_planner/cypher_udf.rs new file mode 100644 index 0000000..fec7d6a --- /dev/null +++ b/crates/lance-graph/src/datafusion_planner/cypher_udf.rs @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Scalar functions with Cypher semantics that DataFusion's built-ins don't match. + +use std::sync::{Arc, LazyLock}; + +use arrow_schema::DataType; +use datafusion::error::Result; +use datafusion::functions_nested::make_array::make_array_udf; +use datafusion::logical_expr::{ + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; + +/// A list literal `[a, b, ...]`. Like `make_array`, except that elements of different +/// kinds become strings instead of failing to cast: Cypher lists may mix types (LDBC IC1's +/// `[uni.name, studyAt.classYear, uniCity.name]`), Arrow lists hold one type. +#[derive(Debug, PartialEq, Eq, Hash)] +struct CypherList { + signature: Signature, +} + +impl ScalarUDFImpl for CypherList { + fn name(&self) -> &str { + "cypher_list" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + make_array_udf().return_type(arg_types) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + let kinds: Vec<&DataType> = arg_types.iter().filter(|t| !t.is_null()).collect(); + let text = + |t: &&DataType| matches!(t, DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View); + if kinds.windows(2).any(|w| w[0] != w[1]) && kinds.iter().any(text) { + Ok(vec![DataType::Utf8; arg_types.len()]) + } else { + make_array_udf().coerce_types(arg_types) + } + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_array_udf().invoke_with_args(args) + } +} + +static CYPHER_LIST: LazyLock> = LazyLock::new(|| { + Arc::new(ScalarUDF::new_from_impl(CypherList { + signature: Signature::user_defined(Volatility::Immutable), + })) +}); + +/// `[items...]` as a Cypher list. +pub(crate) fn list(items: Vec) -> Expr { + if items.is_empty() { + return datafusion::functions_nested::expr_fn::make_array(items); + } + CYPHER_LIST.call(items) +} diff --git a/crates/lance-graph/src/datafusion_planner/expression.rs b/crates/lance-graph/src/datafusion_planner/expression.rs index ace45c4..73cb5f4 100644 --- a/crates/lance-graph/src/datafusion_planner/expression.rs +++ b/crates/lance-graph/src/datafusion_planner/expression.rs @@ -64,6 +64,10 @@ pub(crate) fn to_df_boolean_expr(expr: &BooleanExpression) -> Expr { let list_exprs = list.iter().map(to_df_value_expr).collect::>(); Expr::InList(DFInList::new(Box::new(expr), list_exprs, false)) } + BE::Member { expression, list } => datafusion::functions_nested::expr_fn::array_has( + to_df_value_expr(list), + to_df_value_expr(expression), + ), BE::And(l, r) => Expr::BinaryExpr(BinaryExpr { left: Box::new(to_df_boolean_expr(l)), op: Operator::And, @@ -125,6 +129,7 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { VE::Literal(PV::Null) => { datafusion::logical_expr::Expr::Literal(datafusion::scalar::ScalarValue::Null, None) } + VE::Literal(PV::Variable(v)) => col(v.to_lowercase()), VE::Literal(PV::Parameter(name)) => { panic!( "Parameter ${} should have been substituted during semantic analysis", @@ -158,9 +163,7 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { "coalesce" => datafusion::functions::core::expr_fn::coalesce( args.iter().map(to_df_value_expr).collect(), ), - "__list" => datafusion::functions_nested::expr_fn::make_array( - args.iter().map(to_df_value_expr).collect(), - ), + "__list" => super::cypher_udf::list(args.iter().map(to_df_value_expr).collect()), "tostring" => Expr::Cast(datafusion::logical_expr::Cast::new( Box::new(to_df_value_expr(&args[0])), arrow_schema::DataType::Utf8, @@ -436,7 +439,7 @@ pub(crate) fn to_cypher_column_name(expr: &ValueExpression) -> String { // Convert to Cypher dot notation: variable.property format!("{}.{}", prop.variable, prop.property) } - VE::Variable(v) => v.clone(), + VE::Variable(v) => v.to_lowercase(), VE::Literal(PropertyValue::Property(prop)) => { // Handle nested property references format!("{}.{}", prop.variable, prop.property) diff --git a/crates/lance-graph/src/datafusion_planner/mod.rs b/crates/lance-graph/src/datafusion_planner/mod.rs index 9cbb661..6c8d8ab 100644 --- a/crates/lance-graph/src/datafusion_planner/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/mod.rs @@ -17,6 +17,7 @@ pub mod analysis; mod builder; mod config_helpers; +mod cypher_udf; mod expression; mod join_ops; mod scan_ops; diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index bbdcd51..731c00f 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -33,6 +33,9 @@ pub enum LogicalOperator { expression: ValueExpression, /// The alias for the unwound element alias: String, + /// Set when the list holds node ids of this label (`collect(node)`): each unwound + /// id is joined back to its node table, so `alias` is a node again. + label: Option, }, /// Apply a filter predicate (WHERE clause) @@ -189,6 +192,8 @@ pub struct LogicalPlanner<'a> { fan_out: Option<&'a FanOut>, /// Mark columns created for EXISTS subqueries so far (names `__exists_`). exists_marks: usize, + /// Aliases holding lists of node ids (`collect(node) AS friends`), with the label. + node_lists: HashMap, } impl<'a> LogicalPlanner<'a> { @@ -198,6 +203,7 @@ impl<'a> LogicalPlanner<'a> { config, fan_out: None, exists_marks: 0, + node_lists: HashMap::new(), } } @@ -316,14 +322,7 @@ impl<'a> LogicalPlanner<'a> { if let Some(order_by) = &query.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(), + sort_items: self.sort_items(order_by), }; } @@ -401,17 +400,117 @@ impl<'a> LogicalPlanner<'a> { base: Option, unwind_clause: &UnwindClause, ) -> Result { - // Register the alias variable - self.variables - .insert(unwind_clause.alias.clone(), "Unwound".to_string()); + let label = match &unwind_clause.expression { + ValueExpression::Variable(list) => self.node_lists.get(list).cloned(), + _ => None, + }; + // A list of node ids unwinds to nodes; anything else to plain values. + let kind = label.clone().unwrap_or_else(|| "Unwound".to_string()); + self.variables.insert(unwind_clause.alias.clone(), kind); + let mut expression = unwind_clause.expression.clone(); + self.resolve_value(&mut expression); Ok(LogicalOperator::Unwind { input: base.map(Box::new), - expression: unwind_clause.expression.clone(), + expression, alias: unwind_clause.alias.clone(), + label, }) } + /// Node and path variables used as values: a node is its id (so `a = b`, `a IN list`, + /// `collect(a)`, `count(DISTINCT a)` and `a IS NULL` work on ids) and a path is its + /// length (so `path IS NULL` is true when no path was found). Bare items of WITH and + /// RETURN, which carry whole nodes, are not resolved. + fn resolve_value(&self, expr: &mut ValueExpression) { + match expr { + ValueExpression::Variable(var) => { + let property = match self.variables.get(var).map(String::as_str) { + Some("Path") => Some("length".to_string()), + Some(label) => self + .config + .get_node_mapping(label) + .map(|m| m.id_field.clone()), + None => None, + }; + if let Some(property) = property { + *expr = ValueExpression::Property(PropertyRef { + variable: var.clone(), + property, + }); + } + } + // Path functions take the path variable itself. + ValueExpression::ScalarFunction { name, args } + if matches!(name.to_lowercase().as_str(), "length" | "nodes") => {} + ValueExpression::ScalarFunction { args, .. } + | ValueExpression::AggregateFunction { args, .. } => { + args.iter_mut().for_each(|a| self.resolve_value(a)) + } + ValueExpression::Arithmetic { left, right, .. } => { + self.resolve_value(left); + self.resolve_value(right); + } + ValueExpression::Case { whens, otherwise } => { + for (condition, then) in whens { + self.resolve_boolean(condition); + self.resolve_value(then); + } + if let Some(e) = otherwise { + self.resolve_value(e); + } + } + _ => {} + } + } + + /// [`Self::resolve_value`] for every value in a predicate. + fn resolve_boolean(&self, expr: &mut BooleanExpression) { + match expr { + BooleanExpression::Comparison { left, right, .. } => { + self.resolve_value(left); + self.resolve_value(right); + } + BooleanExpression::And(a, b) | BooleanExpression::Or(a, b) => { + self.resolve_boolean(a); + self.resolve_boolean(b); + } + BooleanExpression::Not(inner) => self.resolve_boolean(inner), + BooleanExpression::In { expression, list } => { + self.resolve_value(expression); + list.iter_mut().for_each(|v| self.resolve_value(v)); + } + BooleanExpression::Member { expression, list } => { + self.resolve_value(expression); + self.resolve_value(list); + } + BooleanExpression::Like { expression, .. } + | BooleanExpression::ILike { expression, .. } + | BooleanExpression::Contains { expression, .. } + | BooleanExpression::StartsWith { expression, .. } + | BooleanExpression::EndsWith { expression, .. } + | BooleanExpression::IsNull(expression) + | BooleanExpression::IsNotNull(expression) => self.resolve_value(expression), + BooleanExpression::Exists(_) | BooleanExpression::PatternExists { .. } => {} + } + } + + /// Sort items with node and path values resolved (see [`Self::resolve_value`]). + fn sort_items(&self, order_by: &OrderByClause) -> Vec { + order_by + .items + .iter() + .map(|item| { + let mut expression = item.expression.clone(); + self.resolve_value(&mut expression); + SortItem { + expression, + direction: item.direction.clone(), + } + }) + .collect() + } + /// Plan a single MATCH clause, optionally starting from an existing base plan. /// /// A clause that starts a plan may be re-rooted at its most selective node, using @@ -429,6 +528,15 @@ impl<'a> LogicalPlanner<'a> { location: snafu::Location::new(file!(), line!(), column!()), }); } + let inferred = MatchClause { + patterns: crate::pattern_alternatives::infer_labels( + &match_clause.patterns, + self.config, + &self.variables, + ), + ..match_clause.clone() + }; + let match_clause = &inferred; if crate::pattern_alternatives::needs_expansion(&match_clause.patterns) { let branches = crate::pattern_alternatives::expand( &match_clause.patterns, @@ -504,6 +612,7 @@ impl<'a> LogicalPlanner<'a> { ) -> Result { let mut predicate = predicate.clone(); let input = self.rewrite_exists(input, &mut predicate)?; + self.resolve_boolean(&mut predicate); Ok(LogicalOperator::Filter { input: Box::new(input), predicate, @@ -721,10 +830,24 @@ impl<'a> LogicalPlanner<'a> { })?; keys.push(qualify_column(var, &mapping.id_field)); } + // The seed also carries the earlier scalars its WHERE reads (`friend IN friends`). + let mut carried = shared.clone(); + if let Some(w) = &clause.where_clause { + let mut names = std::collections::HashSet::new(); + crate::pattern_order::boolean_variables(&w.expression, &mut names); + let local = crate::pattern_order::pattern_variables(&clause.patterns); + let mut scalars: Vec = names + .into_iter() + .filter(|n| !self.variables.contains_key(*n) && !local.contains(*n)) + .map(str::to_string) + .collect(); + scalars.sort(); + carried.extend(scalars); + } let seed = LogicalOperator::Distinct { input: Box::new(LogicalOperator::Project { input: Box::new(base.clone()), - projections: shared + projections: carried .iter() .map(|var| ProjectionItem { expression: ValueExpression::Variable(var.clone()), @@ -843,17 +966,35 @@ impl<'a> LogicalPlanner<'a> { base: Option, path: &PathPattern, ) -> Result { - // Establish a base plan - let mut plan = if let Some(p) = base { - p - } else { - self.plan_node_scan(&path.start_node)? - }; - - // Determine the current source variable for the first hop - let mut current_src = match &path.start_node.variable { - Some(var) => var.clone(), - None => self.extract_variable_from_plan(&plan)?, + // Establish a base plan. A path that continues earlier rows from a node they don't + // bind starts with a scan of that node, crossed with those rows. + let start_bound = path + .start_node + .variable + .as_ref() + .is_some_and(|v| self.variables.contains_key(v)); + let (mut plan, mut current_src) = match base { + Some(p) if start_bound => { + let var = path + .start_node + .variable + .clone() + .expect("bound start is named"); + (p, var) + } + base => { + let scan = self.plan_node_scan(&path.start_node)?; + let var = self.extract_variable_from_plan(&scan)?; + let plan = match base { + Some(p) => LogicalOperator::Join { + left: Box::new(p), + right: Box::new(scan), + join_type: JoinType::Cross, + }, + None => scan, + }; + (plan, var) + } }; // Validate start node label consistency if variable already exists @@ -1027,8 +1168,10 @@ impl<'a> LogicalPlanner<'a> { } } _ => { + let mut expression = item.expression.clone(); + self.resolve_value(&mut expression); projections.push(ProjectionItem { - expression: item.expression.clone(), + expression, alias: alias.clone(), }); } @@ -1083,11 +1226,41 @@ impl<'a> LogicalPlanner<'a> { (None, ValueExpression::Variable(var)) => Some(var.clone()), (None, _) => None, }; - projections.push(ProjectionItem { - expression: item.expression.clone(), - alias, - }); + let mut expression = item.expression.clone(); + if !matches!(expression, ValueExpression::Variable(_)) { + self.resolve_value(&mut expression); + } + projections.push(ProjectionItem { expression, alias }); + } + // Which outputs are lists of node ids: `collect(node)`, or such a list carried on. + let mut node_lists = HashMap::new(); + for item in &items { + let label = match &item.expression { + ValueExpression::AggregateFunction { name, args, .. } + if name.eq_ignore_ascii_case("collect") => + { + match args.as_slice() { + [ValueExpression::Variable(v)] => self + .variables + .get(v) + .filter(|l| self.config.get_node_mapping(l).is_some()) + .cloned(), + _ => None, + } + } + ValueExpression::Variable(v) => self.node_lists.get(v).cloned(), + _ => None, + }; + let name = match (&item.alias, &item.expression) { + (Some(alias), _) => Some(alias.clone()), + (None, ValueExpression::Variable(v)) => Some(v.clone()), + _ => None, + }; + if let (Some(label), Some(name)) = (label, name) { + node_lists.insert(name, label); + } } + self.node_lists = node_lists; // WITH ends the scope of everything it does not project: a later MATCH reusing a // dropped name binds a new variable instead of expecting columns that are gone. let kept: Vec<&str> = with_clause @@ -1115,14 +1288,7 @@ impl<'a> LogicalPlanner<'a> { 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(), + sort_items: self.sort_items(order_by), }; } diff --git a/crates/lance-graph/src/parameter_substitution.rs b/crates/lance-graph/src/parameter_substitution.rs index cdc1183..d0fdaae 100644 --- a/crates/lance-graph/src/parameter_substitution.rs +++ b/crates/lance-graph/src/parameter_substitution.rs @@ -209,6 +209,10 @@ fn substitute_in_boolean_expression( substitute_in_value_expression(item, parameters)?; } } + BooleanExpression::Member { expression, list } => { + substitute_in_value_expression(expression, parameters)?; + substitute_in_value_expression(list, parameters)?; + } BooleanExpression::Like { expression, .. } | BooleanExpression::ILike { expression, .. } | BooleanExpression::Contains { expression, .. } diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index a3883fe..b6b4d3a 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -35,9 +35,75 @@ pub fn parse_cypher_query(input: &str) -> Result { }); } + let mut query = query; + let stages = query + .further_with_stages + .iter_mut() + .flat_map(|s| &mut s.reading_clauses); + let mut anonymous = 0; + for clause in query + .reading_clauses + .iter_mut() + .chain(&mut query.post_with_reading_clauses) + .chain(stages) + { + if let ReadingClause::Match(m) = clause { + move_variable_properties_to_where(m, &mut anonymous); + } + } Ok(query) } +/// Inline node properties whose value is a variable or property from an earlier clause +/// (`(t:Tag {id: tagId})`) become `t.id = tagId` in the clause's own WHERE, since scans +/// filter only on constants. An unnamed node gets a name for the condition. +fn move_variable_properties_to_where(clause: &mut MatchClause, anonymous: &mut usize) { + let mut conditions = Vec::new(); + for pattern in &mut clause.patterns { + let nodes: Vec<&mut NodePattern> = match pattern { + GraphPattern::Node(n) => vec![n], + GraphPattern::Path(p) => std::iter::once(&mut p.start_node) + .chain(p.segments.iter_mut().map(|s| &mut s.end_node)) + .collect(), + }; + for node in nodes { + let keys: Vec = node + .properties + .iter() + .filter(|(_, v)| { + matches!(v, PropertyValue::Variable(_) | PropertyValue::Property(_)) + }) + .map(|(k, _)| k.clone()) + .collect(); + for key in keys { + let value = match node.properties.remove(&key) { + Some(PropertyValue::Variable(v)) => ValueExpression::Variable(v), + Some(PropertyValue::Property(p)) => ValueExpression::Property(p), + _ => unreachable!("filtered above"), + }; + let variable = node.variable.get_or_insert_with(|| { + *anonymous += 1; + format!("__inline_{anonymous}") + }); + conditions.push(BooleanExpression::Comparison { + left: ValueExpression::Property(PropertyRef { + variable: variable.clone(), + property: key, + }), + operator: ComparisonOperator::Equal, + right: value, + }); + } + } + } + let existing = clause.where_clause.take().map(|w| w.expression); + clause.where_clause = conditions + .into_iter() + .chain(existing) + .reduce(|a, b| BooleanExpression::And(Box::new(a), Box::new(b))) + .map(|expression| WhereClause { expression }); +} + // Top-level parser for a complete Cypher query fn cypher_query(input: &str) -> IResult<&str, CypherQuery> { let (input, _) = multispace0(input)?; @@ -327,7 +393,11 @@ fn property_pair(input: &str) -> IResult<&str, (String, PropertyValue)> { let (input, _) = multispace0(input)?; let (input, _) = char(':')(input)?; let (input, _) = multispace0(input)?; - let (input, value) = property_value(input)?; + let (input, value) = alt(( + property_value, + map(property_reference, PropertyValue::Property), + map(identifier, |v| PropertyValue::Variable(v.to_string())), + ))(input)?; Ok((input, (key.to_string(), value))) } @@ -458,6 +528,17 @@ fn comparison_expression(input: &str) -> IResult<&str, BooleanExpression> { }, )); } + if let Ok((input_after_in, (_, _, list))) = + tuple((tag_no_case("IN"), multispace1, value_expression))(input) + { + return Ok(( + input_after_in, + BooleanExpression::Member { + expression: left, + list, + }, + )); + } // Match LIKE pattern if let Ok((input_after_like, (_, _, pattern))) = tuple((tag_no_case("LIKE"), multispace0, string_literal))(input) @@ -539,16 +620,32 @@ fn comparison_expression(input: &str) -> IResult<&str, BooleanExpression> { let (input, operator) = comparison_operator(input)?; let (input, _) = multispace0(input)?; - let (input, right) = value_expression(input)?; - - Ok(( - input, - BooleanExpression::Comparison { - left: left_clone, - operator, - right, - }, - )) + let (mut input, mut right) = value_expression(input)?; + let mut expr = BooleanExpression::Comparison { + left: left_clone, + operator, + right: right.clone(), + }; + // A chain `a < b <= c` means `a < b AND b <= c`. + while let Ok((rest, (_, operator, _, next))) = tuple(( + multispace0, + comparison_operator, + multispace0, + value_expression, + ))(input) + { + expr = BooleanExpression::And( + Box::new(expr), + Box::new(BooleanExpression::Comparison { + left: right, + operator, + right: next.clone(), + }), + ); + right = next; + input = rest; + } + Ok((input, expr)) } // Parse a comparison operator diff --git a/crates/lance-graph/src/pattern_alternatives.rs b/crates/lance-graph/src/pattern_alternatives.rs index af1adf2..4169a75 100644 --- a/crates/lance-graph/src/pattern_alternatives.rs +++ b/crates/lance-graph/src/pattern_alternatives.rs @@ -31,6 +31,123 @@ use crate::error::{GraphError, Result}; /// Branch limit, so a pattern full of alternatives can't explode the plan. const MAX_BRANCHES: usize = 64; +/// `patterns` with labels filled in for unlabeled nodes that the configured endpoint labels +/// pin down: `(friend)<-[:postHasCreator]-(post)` makes `post` a Post when +/// `postHasCreator` goes from Post to Person. Labels flow along paths and between +/// patterns through shared variables; `bound` holds labels from earlier clauses. Nodes +/// next to alternatives whose labels differ are left for [`expand`] to branch on. +pub fn infer_labels( + patterns: &[GraphPattern], + config: &GraphConfig, + bound: &HashMap, +) -> Vec { + let mut patterns = patterns.to_vec(); + let mut known: HashMap = bound + .iter() + .filter(|(_, l)| config.get_node_mapping(l).is_some()) + .map(|(v, l)| (v.clone(), l.clone())) + .collect(); + // The label one end of `rel` has, if every type agrees: `source` side or target side. + let end_label = |rel: &RelationshipPattern, source: bool| -> Option { + let mut labels = rel.types.iter().map(|t| { + let m = config.get_relationship_mapping(t)?; + if source { + m.source_label.clone() + } else { + m.target_label.clone() + } + }); + let first = labels.next()??; + labels.all(|l| l.as_ref() == Some(&first)).then_some(first) + }; + loop { + let mut changed = false; + for pattern in &mut patterns { + let GraphPattern::Path(path) = pattern else { + continue; + }; + let rels: Vec = path + .segments + .iter() + .map(|s| s.relationship.clone()) + .collect(); + let mut nodes: Vec<&mut NodePattern> = std::iter::once(&mut path.start_node) + .chain(path.segments.iter_mut().map(|s| &mut s.end_node)) + .collect(); + let label_of = + |n: &NodePattern, known: &HashMap| match n.labels.as_slice() { + [l] => Some(l.clone()), + [] => n.variable.as_ref().and_then(|v| known.get(v).cloned()), + _ => None, + }; + for (i, rel) in rels.iter().enumerate() { + if rel.types.is_empty() { + continue; + } + // (left side, right side) as written. + let (left, right) = match rel.direction { + RelationshipDirection::Outgoing => { + (end_label(rel, true), end_label(rel, false)) + } + RelationshipDirection::Incoming => { + (end_label(rel, false), end_label(rel, true)) + } + RelationshipDirection::Undirected => { + let (s, t) = (end_label(rel, true), end_label(rel, false)); + if s == t { + (s.clone(), t) + } else { + (None, None) + } + } + }; + for (slot, label) in [(i, left), (i + 1, right)] { + let Some(label) = label else { continue }; + let node = &mut nodes[slot]; + if !node.labels.is_empty() || label_of(node, &known).is_some() { + continue; + } + match &node.variable { + Some(v) => { + known.insert(v.clone(), label); + } + None => node.labels = vec![label], + } + changed = true; + } + } + // Written labels of named nodes inform the other patterns. + for node in nodes { + if let (Some(v), [l]) = (&node.variable, node.labels.as_slice()) { + changed |= known.insert(v.clone(), l.clone()).is_none(); + } + } + } + if !changed { + break; + } + } + // Write inferred labels onto the unlabeled occurrences of named, unbound variables. + for pattern in &mut patterns { + let nodes: Vec<&mut NodePattern> = match pattern { + GraphPattern::Node(n) => vec![n], + GraphPattern::Path(p) => std::iter::once(&mut p.start_node) + .chain(p.segments.iter_mut().map(|s| &mut s.end_node)) + .collect(), + }; + for node in nodes { + if let Some(v) = &node.variable { + if node.labels.is_empty() && !bound.contains_key(v) { + if let Some(l) = known.get(v) { + node.labels = vec![l.clone()]; + } + } + } + } + } + patterns +} + /// True when any relationship has several types or any node several labels. pub fn needs_expansion(patterns: &[GraphPattern]) -> bool { let nodes_multi = |n: &NodePattern| n.labels.len() > 1; diff --git a/crates/lance-graph/src/pattern_order.rs b/crates/lance-graph/src/pattern_order.rs index ca51280..8915797 100644 --- a/crates/lance-graph/src/pattern_order.rs +++ b/crates/lance-graph/src/pattern_order.rs @@ -148,6 +148,7 @@ fn score_conjunct( } => Selectivity::Weak, BooleanExpression::Comparison { .. } | BooleanExpression::In { .. } + | BooleanExpression::Member { .. } | BooleanExpression::Like { .. } | BooleanExpression::ILike { .. } | BooleanExpression::Contains { .. } @@ -169,7 +170,7 @@ fn is_constant(expr: &ValueExpression) -> bool { ) } -fn boolean_variables<'a>(expr: &'a BooleanExpression, out: &mut HashSet<&'a str>) { +pub(crate) fn boolean_variables<'a>(expr: &'a BooleanExpression, out: &mut HashSet<&'a str>) { match expr { // A subquery says nothing about how selective one variable is. BooleanExpression::PatternExists { .. } => {} @@ -189,6 +190,10 @@ fn boolean_variables<'a>(expr: &'a BooleanExpression, out: &mut HashSet<&'a str> value_variables(expression, out); list.iter().for_each(|v| value_variables(v, out)); } + BooleanExpression::Member { expression, list } => { + value_variables(expression, out); + value_variables(list, out); + } BooleanExpression::Like { expression, .. } | BooleanExpression::ILike { expression, .. } | BooleanExpression::Contains { expression, .. } @@ -232,6 +237,24 @@ fn value_variables<'a>(expr: &'a ValueExpression, out: &mut HashSet<&'a str>) { } } +/// The node and relationship variables written in `patterns`. +pub(crate) fn pattern_variables(patterns: &[GraphPattern]) -> HashSet<&str> { + let mut out = HashSet::new(); + for pattern in patterns { + match pattern { + GraphPattern::Node(n) => out.extend(n.variable.as_deref()), + GraphPattern::Path(p) => { + out.extend(p.start_node.variable.as_deref()); + for s in &p.segments { + out.extend(s.relationship.variable.as_deref()); + out.extend(s.end_node.variable.as_deref()); + } + } + } + } + out +} + /// A MATCH clause's patterns as nodes (in first-mention order) and edges. struct PatternGraph { /// One entry per variable, merging labels and inline properties from every mention. diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index 6d51d7a..47818fd 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -18,6 +18,8 @@ use std::collections::{HashMap, HashSet}; pub struct SemanticAnalyzer { config: GraphConfig, variables: HashMap, + /// Variables whose scope a WITH ended; still type-checked at the end. + retired: Vec, current_scope: ScopeType, } @@ -66,6 +68,7 @@ impl SemanticAnalyzer { Self { config, variables: HashMap::new(), + retired: Vec::new(), current_scope: ScopeType::Match, } } @@ -277,6 +280,10 @@ impl SemanticAnalyzer { // Normalize to lowercase for case-insensitive behavior let var_name_lower = var_name.to_lowercase(); if let Some(existing) = self.variables.get_mut(&var_name_lower) { + // An UNWIND alias over a list of nodes is matched as a node. + if existing.variable_type == VariableType::Property { + existing.variable_type = VariableType::Node; + } if existing.variable_type != VariableType::Node { return Err(GraphError::PlanError { message: format!("Variable '{}' redefined with different type", var_name), @@ -380,6 +387,10 @@ impl SemanticAnalyzer { self.analyze_value_expression(item)?; } } + BooleanExpression::Member { expression, list } => { + self.analyze_value_expression(expression)?; + self.analyze_value_expression(list)?; + } BooleanExpression::Like { expression, .. } => { self.analyze_value_expression(expression)?; } @@ -711,10 +722,38 @@ impl SemanticAnalyzer { // Validate WITH item expressions (similar to RETURN) for item in &with_clause.items { self.analyze_value_expression(&item.expression)?; - if let Some(alias) = &item.alias { - self.register_projection_alias(alias); - } } + // WITH ends the scope of everything it does not carry; `WITH b AS x` renames b. + let mut kept = HashMap::new(); + for item in &with_clause.items { + let carried = match &item.expression { + ValueExpression::Variable(v) => self.variables.get(&v.to_lowercase()).cloned(), + _ => None, + }; + let name = match (&item.alias, &item.expression) { + (Some(alias), _) => alias.clone(), + (None, ValueExpression::Variable(v)) => v.clone(), + // An unaliased `WITH b.name` stays addressable as `b.name`. + (None, ValueExpression::Property(p)) => { + if let Some(info) = self.variables.get(&p.variable.to_lowercase()) { + kept.insert(p.variable.to_lowercase(), info.clone()); + } + continue; + } + (None, _) => continue, + }; + let info = carried.unwrap_or_else(|| VariableInfo { + name: name.clone(), + variable_type: VariableType::Property, + labels: vec![], + properties: HashSet::new(), + defined_in: self.current_scope.clone(), + }); + kept.insert(name.to_lowercase(), info); + } + self.retired + .extend(std::mem::take(&mut self.variables).into_values()); + self.variables = kept; // Validate ORDER BY within WITH if present if let Some(order_by) = &with_clause.order_by { for item in &order_by.items { @@ -769,7 +808,7 @@ impl SemanticAnalyzer { // - Check that arithmetic operations are valid // Check that properties referenced in patterns exist in schema when property fields are defined - for var_info in self.variables.values() { + for var_info in self.variables.values().chain(&self.retired) { match var_info.variable_type { VariableType::Node => { // Collect property_fields from all known label mappings that specify properties diff --git a/python/python/tests/test_engine_values.py b/python/python/tests/test_engine_values.py new file mode 100644 index 0000000..23698fb --- /dev/null +++ b/python/python/tests/test_engine_values.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Nodes as values, lists of nodes, UNWIND, parameters and inferred labels. + +The constructs LDBC's official IC queries use: `a = b` on nodes, `x IN list`, `collect(node)` +then `UNWIND`, inline properties that read a variable, chained comparisons, unlabeled +nodes whose label follows from a relationship's endpoints. Answers are written out by hand. +""" + +from datetime import datetime, timezone + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +def edges(pairs): + src, dst = zip(*pairs) + return pa.table( + {"src": pa.array(src, pa.int64()), "dst": pa.array(dst, pa.int64())} + ) + + +def utc(*args): + return datetime(*args, tzinfo=timezone.utc) + + +TABLES = { + "Person": pa.table({"id": [1, 2, 3, 4], "name": ["a", "b", "c", "d"]}), + "City": pa.table({"id": [100, 101], "name": ["X", "Y"]}), + "Post": pa.table( + { + "id": [10, 11, 12], + "created": pa.array( + [utc(2020, 1, 1), utc(2020, 6, 1), utc(2021, 1, 1)], + pa.timestamp("us", "UTC"), + ), + } + ), + "KNOWS": edges([(1, 2), (2, 3), (3, 4), (1, 3)]), + "LIVES_IN": edges([(1, 100), (2, 100), (3, 101), (4, 101)]), + "HAS_CREATOR": edges([(10, 2), (11, 3), (12, 3)]), +} +CONFIG = ( + GraphConfig.builder() + .with_node_label("Person", "id") + .with_node_label("City", "id") + .with_node_label("Post", "id") + .with_relationship("KNOWS", "src", "dst", "Person", "Person") + .with_relationship("LIVES_IN", "src", "dst", "Person", "City") + .with_relationship("HAS_CREATOR", "src", "dst", "Post", "Person") + .build() +) + +CASES = [ + # A node compared with a node compares ids. + ( + "MATCH (a:Person {id: 1})-[:KNOWS*1..2]->(b) WHERE NOT a = b RETURN DISTINCT b.id", + [(2,), (3,), (4,)], + ), + # `p` is a Post: HAS_CREATOR goes from Post to Person. + ("MATCH (a:Person {id: 2})<-[:HAS_CREATOR]-(p) RETURN p.id", [(10,)]), + ( + "MATCH (x:City {name: 'X'}), (p:Person)-[:LIVES_IN]->(c) WHERE c IN [x] RETURN p.id", + [(1,), (2,)], + ), + # A list of nodes, and a pattern after WITH that starts at a new node. + ( + "MATCH (:City {name: 'Y'})<-[:LIVES_IN]-(p:Person) WITH collect(p) AS ys " + "MATCH (a:Person)-[:KNOWS]->(b) WHERE b IN ys RETURN a.id, b.id", + [(1, 3), (2, 3), (3, 4)], + ), + ( + "MATCH (:Person {id: 1})-[:KNOWS]->(b) WITH collect(b) AS bs " + "UNWIND bs AS f RETURN f.id, f.name", + [(2, "b"), (3, "c")], + ), + ( + "MATCH (:Person {id: 1})-[:KNOWS]->(b) WITH collect(DISTINCT b) AS bs " + "UNWIND bs AS f MATCH (f)<-[:HAS_CREATOR]-(post) RETURN f.id, post.id", + [(2, 10), (3, 11), (3, 12)], + ), + # An inline property that reads a variable, on an unnamed node. + ( + "MATCH (c:City {name: 'X'}) WITH c.id AS cid " + "MATCH (p:Person)-[:LIVES_IN]->(:City {id: cid}) RETURN p.id", + [(1,), (2,)], + ), + ("MATCH (a:Person)-[:KNOWS*1..2]->(b) RETURN count(DISTINCT b) AS n", [(3,)]), + # A scalar from WITH survives a variable-length path. + ( + "MATCH (c:City {name: 'Y'}) WITH c.id AS cid " + "MATCH (:Person {id: 1})-[:KNOWS*1..2]->(b)-[:LIVES_IN]->(c2) WHERE c2.id = cid " + "RETURN DISTINCT b.id", + [(3,), (4,)], + ), + # collect() of nothing is an empty list, not null. + ( + "MATCH (a:Person) OPTIONAL MATCH (a)<-[:HAS_CREATOR]-(p:Post) " + "WITH a, collect(p.id) AS posts RETURN a.id, size(posts)", + [(1, 0), (2, 1), (3, 2), (4, 0)], + ), + # An OPTIONAL MATCH whose WHERE reads a list carried by WITH (LDBC IC5's shape). + ( + "MATCH (:Person)-[:KNOWS]->(a:Person)-[:LIVES_IN]->(c:City) WITH DISTINCT c, a " + "WITH c, collect(a) AS known " + "OPTIONAL MATCH (c)<-[:LIVES_IN]-(x:Person)<-[:HAS_CREATOR]-(post:Post) " + "WHERE x IN known RETURN c.name, count(post) AS n", + [("X", 1), ("Y", 2)], + ), + # Mixed element types become strings, as Arrow lists hold one type. + ("MATCH (p:Person {id: 1}) RETURN [p.name, p.id] AS l", [(["a", "1"],)]), +] + + +def rows(table): + return sorted(tuple(r.values()) for r in table.to_pylist()) + + +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize(("query", "expected"), CASES) +def test_values(use_csr, query, expected): + engine = CypherEngine(CONFIG, TABLES, use_csr=use_csr) + assert rows(engine.execute(query)) == expected + + +def test_parameters_and_chained_comparison(): + engine = CypherEngine(CONFIG, TABLES) + query = "MATCH (p:Post) WHERE $start <= p.created < $end RETURN p.id" + params = {"start": utc(2020, 3, 1), "end": utc(2021, 1, 1)} + assert rows(engine.execute(query, params)) == [(11,)] + assert rows( + engine.execute("MATCH (p:Person {id: $id}) RETURN p.name", {"id": 3}) + ) == [("c",)] From 464e7046a1fdfe8a206b82d5b9ce3bae0762edbc Mon Sep 17 00:00:00 2001 From: prrao87 Date: Sat, 24 Oct 2026 09:26:05 -0400 Subject: [PATCH 2/5] feat(cypher): boolean values, maps, field access and ordered collect Predicates as values (path IS NULL, NOT (pattern) AS isNew), map literals whose node values keep their properties, value.field on maps and temporal values, head/toInteger/toFloat/date/datetime and duration.inSeconds, and collect() in the order of a preceding WITH ... ORDER BY. --- crates/lance-graph/src/ast.rs | 84 ++++++++++- .../builder/aggregate_ops.rs | 17 ++- .../datafusion_planner/builder/basic_ops.rs | 65 ++++++--- .../src/datafusion_planner/cypher_udf.rs | 115 ++++++++++++++- .../src/datafusion_planner/expression.rs | 132 ++++++++++++++++-- crates/lance-graph/src/logical_plan.rs | 15 ++ .../lance-graph/src/parameter_substitution.rs | 7 + crates/lance-graph/src/parser.rs | 103 ++++++++++---- crates/lance-graph/src/pattern_order.rs | 3 + crates/lance-graph/src/semantic.rs | 24 ++++ python/python/tests/test_engine_values.py | 41 ++++++ 11 files changed, 539 insertions(+), 67 deletions(-) diff --git a/crates/lance-graph/src/ast.rs b/crates/lance-graph/src/ast.rs index b0de0e4..495ced7 100644 --- a/crates/lance-graph/src/ast.rs +++ b/crates/lance-graph/src/ast.rs @@ -398,6 +398,16 @@ pub enum ValueExpression { }, /// Parameter reference for query parameters (e.g., $query_vector) Parameter(String), + /// A predicate used as a value: `path IS NULL`, `NOT (a)-[:R]-(b) AS isNew`. + Boolean(Box), + /// A map literal `{key: value, ...}`; a node as a value inside it keeps its properties. + Map(Vec<(String, ValueExpression)>), + /// `value.field` on something other than a node or relationship variable: a map key, + /// or a component of a date or timestamp (`birthday.month`). + Field { + base: Box, + field: String, + }, /// Vector literal: [0.1, 0.2, 0.3] /// Represents an inline vector for similarity search VectorLiteral(Vec), @@ -419,7 +429,8 @@ pub fn classify_function(name: &str) -> FunctionType { match name.to_lowercase().as_str() { "count" | "sum" | "avg" | "min" | "max" | "collect" => FunctionType::Aggregate, "tolower" | "lower" | "toupper" | "upper" | "coalesce" | "date_part" | "to_epoch_ms" - | "length" | "nodes" | "size" | "range" | "__index" | "__list" | "tostring" => { + | "length" | "nodes" | "size" | "range" | "__index" | "__list" | "tostring" | "head" + | "tointeger" | "tofloat" | "date" | "datetime" | "duration.inseconds" => { FunctionType::Scalar } // Vector functions are handled separately as special variants @@ -571,6 +582,77 @@ impl PropertyRef { } } +/// Apply `f` to every value inside `expr`, children before parents, including the values +/// inside predicates. Patterns of EXISTS subqueries are left alone. +pub fn transform_value(expr: &mut ValueExpression, f: &mut dyn FnMut(&mut ValueExpression)) { + match expr { + ValueExpression::ScalarFunction { args, .. } + | ValueExpression::AggregateFunction { args, .. } => { + args.iter_mut().for_each(|a| transform_value(a, f)) + } + ValueExpression::Arithmetic { left, right, .. } + | ValueExpression::VectorDistance { left, right, .. } + | ValueExpression::VectorSimilarity { left, right, .. } => { + transform_value(left, f); + transform_value(right, f); + } + ValueExpression::Case { whens, otherwise } => { + for (condition, then) in whens { + transform_boolean_values(condition, f); + transform_value(then, f); + } + if let Some(e) = otherwise { + transform_value(e, f); + } + } + ValueExpression::Boolean(b) => transform_boolean_values(b, f), + ValueExpression::Map(entries) => { + entries.iter_mut().for_each(|(_, v)| transform_value(v, f)) + } + ValueExpression::Field { base, .. } => transform_value(base, f), + ValueExpression::Variable(_) + | ValueExpression::Property(_) + | ValueExpression::Literal(_) + | ValueExpression::Parameter(_) + | ValueExpression::VectorLiteral(_) => {} + } + f(expr); +} + +/// [`transform_value`] over every value in a predicate. +pub fn transform_boolean_values( + expr: &mut BooleanExpression, + f: &mut dyn FnMut(&mut ValueExpression), +) { + match expr { + BooleanExpression::Comparison { left, right, .. } => { + transform_value(left, f); + transform_value(right, f); + } + BooleanExpression::And(a, b) | BooleanExpression::Or(a, b) => { + transform_boolean_values(a, f); + transform_boolean_values(b, f); + } + BooleanExpression::Not(inner) => transform_boolean_values(inner, f), + BooleanExpression::In { expression, list } => { + transform_value(expression, f); + list.iter_mut().for_each(|v| transform_value(v, f)); + } + BooleanExpression::Member { expression, list } => { + transform_value(expression, f); + transform_value(list, f); + } + BooleanExpression::Like { expression, .. } + | BooleanExpression::ILike { expression, .. } + | BooleanExpression::Contains { expression, .. } + | BooleanExpression::StartsWith { expression, .. } + | BooleanExpression::EndsWith { expression, .. } + | BooleanExpression::IsNull(expression) + | BooleanExpression::IsNotNull(expression) => transform_value(expression, f), + BooleanExpression::Exists(_) | BooleanExpression::PatternExists { .. } => {} + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs index 781e7fa..0cf71fc 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs @@ -57,6 +57,7 @@ impl DataFusionPlanner { &self, input_plan: LogicalPlan, projections: &[ProjectionItem], + ordering: &[datafusion::logical_expr::SortExpr], ) -> Result { // Separate group expressions (non-aggregates) from aggregate expressions let mut group_exprs = Vec::new(); @@ -66,6 +67,14 @@ impl DataFusionPlanner { let mut aggregates = Vec::new(); let schema = input_plan.schema().clone(); + let bound: Vec = projections + .iter() + .map(|p| ProjectionItem { + expression: super::super::expression::bind_value(&p.expression, &schema), + alias: p.alias.clone(), + }) + .collect(); + let projections = bound.as_slice(); for p in projections { // Grouping by a node groups by all its columns (`WITH b, count(*) AS n`). if let Some(columns) = Self::entity_columns(&schema, &p.expression, None) { @@ -87,7 +96,13 @@ impl DataFusionPlanner { let first = aggregates.len(); split_aggregates(&mut rewritten, &mut aggregates); for (column, aggregate) in &aggregates[first..] { - let aggregate = super::super::expression::to_df_value_expr(aggregate); + let mut aggregate = super::super::expression::to_df_value_expr(aggregate); + // collect() gathers in the order of a preceding `WITH ... ORDER BY`. + if let datafusion::logical_expr::Expr::AggregateFunction(f) = &mut aggregate { + if f.func.name() == "array_agg" && !f.params.distinct { + f.params.order_by = ordering.to_vec(); + } + } agg_exprs.push(aggregate.alias(column)); } outer.push(super::super::expression::to_df_value_expr(&rewritten).alias(alias)); diff --git a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs index 6697c8b..c10b116 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs @@ -86,7 +86,8 @@ impl DataFusionPlanner { predicate: &crate::ast::BooleanExpression, ) -> Result { let input_plan = self.build_operator(ctx, input)?; - let expr = super::super::expression::to_df_boolean_expr(predicate); + let predicate = super::super::expression::bind_boolean(predicate, input_plan.schema()); + let expr = super::super::expression::to_df_boolean_expr(&predicate); LogicalPlanBuilder::from(input_plan) .filter(expr) .map_err(|e| self.plan_error("Failed to build filter", e))? @@ -100,18 +101,29 @@ impl DataFusionPlanner { input: &LogicalOperator, projections: &[ProjectionItem], ) -> Result { - let input_plan = self.build_operator(ctx, input)?; - // Check if any projection contains an aggregate function let has_aggregates = projections .iter() .any(|p| super::super::expression::contains_aggregate(&p.expression)); - if has_aggregates { - self.build_project_with_aggregates(input_plan, projections) - } else { - self.build_simple_project(input_plan, projections) + if !has_aggregates { + let input_plan = self.build_operator(ctx, input)?; + return self.build_simple_project(input_plan, projections); } + // An ORDER BY right before an aggregation only orders what collect() gathers, as in + // Neo4j (LDBC IC7 takes `head(collect(...))` of sorted rows): it becomes the + // aggregate's own ORDER BY instead of a sort of its input. + if let LogicalOperator::Sort { + input: inner, + sort_items, + } = input + { + let input_plan = self.build_operator(ctx, inner)?; + let ordering = self.sort_exprs(&input_plan, sort_items); + return self.build_project_with_aggregates(input_plan, projections, &ordering); + } + let input_plan = self.build_operator(ctx, input)?; + self.build_project_with_aggregates(input_plan, projections, &[]) } /// The columns of a node or relationship variable (`b` -> `b__id`, `b__name`, ...), @@ -162,7 +174,8 @@ impl DataFusionPlanner { exprs.extend(columns); continue; } - let expr = super::super::expression::to_df_value_expr(&p.expression); + let bound = super::super::expression::bind_value(&p.expression, &schema); + let expr = super::super::expression::to_df_value_expr(&bound); // Apply alias if provided, otherwise use Cypher dot notation // Normalize alias to lowercase for case-insensitive behavior exprs.push(if let Some(alias) = &p.alias { @@ -200,26 +213,33 @@ impl DataFusionPlanner { sort_items: &[SortItem], ) -> Result { let input_plan = self.build_operator(ctx, input)?; + let sort_exprs = self.sort_exprs(&input_plan, sort_items); - // Convert sort items to DataFusion sort expressions - let sort_exprs: Vec = sort_items + LogicalPlanBuilder::from(input_plan) + .sort(sort_exprs) + .map_err(|e| self.plan_error("Failed to build sort", e))? + .build() + .map_err(|e| self.plan_error("Failed to build plan", e)) + } + + /// `sort_items` as DataFusion sort expressions over `input_plan`'s columns. + pub(crate) fn sort_exprs( + &self, + input_plan: &LogicalPlan, + sort_items: &[SortItem], + ) -> Vec { + sort_items .iter() .map(|item| { - let expr = super::super::expression::to_df_value_expr(&item.expression); - let asc = matches!(item.direction, crate::ast::SortDirection::Ascending); + let bound = + super::super::expression::bind_value(&item.expression, input_plan.schema()); SortExpr { - expr, - asc, + expr: super::super::expression::to_df_value_expr(&bound), + asc: matches!(item.direction, crate::ast::SortDirection::Ascending), nulls_first: true, } }) - .collect(); - - LogicalPlanBuilder::from(input_plan) - .sort(sort_exprs) - .map_err(|e| self.plan_error("Failed to build sort", e))? - .build() - .map_err(|e| self.plan_error("Failed to build plan", e)) + .collect() } pub(crate) fn build_limit( @@ -268,7 +288,8 @@ impl DataFusionPlanner { }; // Convert expression to DataFusion Expr - let df_expr = super::super::expression::to_df_value_expr(expression); + let bound = super::super::expression::bind_value(expression, input_plan.schema()); + let df_expr = super::super::expression::to_df_value_expr(&bound); // A list of node ids: unnest under a temporary name, then join each id back to the // node table so the alias carries the node's columns like any matched node. let unnested = match label { diff --git a/crates/lance-graph/src/datafusion_planner/cypher_udf.rs b/crates/lance-graph/src/datafusion_planner/cypher_udf.rs index fec7d6a..e4f47a2 100644 --- a/crates/lance-graph/src/datafusion_planner/cypher_udf.rs +++ b/crates/lance-graph/src/datafusion_planner/cypher_udf.rs @@ -5,11 +5,15 @@ use std::sync::{Arc, LazyLock}; -use arrow_schema::DataType; +use arrow::array::{Array, AsArray}; +use arrow::compute::{cast, date_part, DatePart}; +use arrow_schema::{DataType, Field, FieldRef}; +use datafusion::common::{exec_err, plan_err, ScalarValue}; use datafusion::error::Result; use datafusion::functions_nested::make_array::make_array_udf; use datafusion::logical_expr::{ - ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, + lit, ColumnarValue, Expr, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, + Signature, Volatility, }; /// A list literal `[a, b, ...]`. Like `make_array`, except that elements of different @@ -62,3 +66,110 @@ pub(crate) fn list(items: Vec) -> Expr { } CYPHER_LIST.call(items) } + +/// `value.field` on a value: a map (struct) key, matched case-insensitively like +/// properties, or a component of a date or timestamp (`birthday.month`, `t.year`). +#[derive(Debug, PartialEq, Eq, Hash)] +struct CypherField { + signature: Signature, +} + +/// The date or time component `name` stands for, if any. +fn temporal_part(name: &str) -> Option { + Some(match name.to_lowercase().as_str() { + "year" => DatePart::Year, + "quarter" => DatePart::Quarter, + "month" => DatePart::Month, + "week" => DatePart::Week, + "day" => DatePart::Day, + "dayofyear" | "ordinalday" => DatePart::DayOfYear, + "hour" => DatePart::Hour, + "minute" => DatePart::Minute, + "second" => DatePart::Second, + "millisecond" => DatePart::Millisecond, + "microsecond" => DatePart::Microsecond, + "nanosecond" => DatePart::Nanosecond, + _ => return None, + }) +} + +fn is_temporal(t: &DataType) -> bool { + matches!( + t, + DataType::Date32 | DataType::Date64 | DataType::Timestamp(..) + ) +} + +/// Index of the struct field named `name`, ignoring case. +fn struct_field(fields: &arrow_schema::Fields, name: &str) -> Option { + fields + .iter() + .position(|f| f.name().eq_ignore_ascii_case(name)) +} + +impl ScalarUDFImpl for CypherField { + fn name(&self) -> &str { + "cypher_field" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + unreachable!("return_field_from_args decides the type") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + let Some(Some(ScalarValue::Utf8(Some(name)))) = args.scalar_arguments.get(1) else { + return plan_err!("cypher_field takes a field name literal"); + }; + let data_type = match args.arg_fields[0].data_type() { + DataType::Struct(fields) => match struct_field(fields, name) { + Some(i) => fields[i].data_type().clone(), + None => return plan_err!("no key '{name}' in this map"), + }, + t if is_temporal(t) && temporal_part(name).is_some() => DataType::Int64, + DataType::Null => DataType::Null, + t => return plan_err!("can't read '.{name}' of a {t} value"), + }; + Ok(Arc::new(Field::new(self.name(), data_type, true))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let ColumnarValue::Scalar(ScalarValue::Utf8(Some(name))) = &args.args[1] else { + return exec_err!("cypher_field takes a field name literal"); + }; + let base = args.args[0].to_array(args.number_rows)?; + let result = match base.data_type() { + DataType::Struct(fields) => { + let i = struct_field(fields, name).expect("checked when planning"); + let column = base.as_struct().column(i).clone(); + // A null map gives null keys, whatever its child arrays hold. + let nulls = arrow::buffer::NullBuffer::union( + base.logical_nulls().as_ref(), + column.logical_nulls().as_ref(), + ); + let data = column.to_data().into_builder().nulls(nulls).build()?; + arrow::array::make_array(data) + } + DataType::Null => base, + _ => { + let part = temporal_part(name).expect("checked when planning"); + cast(&date_part(&base, part)?, &DataType::Int64)? + } + }; + Ok(ColumnarValue::Array(result)) + } +} + +static CYPHER_FIELD: LazyLock> = LazyLock::new(|| { + Arc::new(ScalarUDF::new_from_impl(CypherField { + signature: Signature::any(2, Volatility::Immutable), + })) +}); + +/// `base.name` on a map or temporal value. +pub(crate) fn field(base: Expr, name: &str) -> Expr { + CYPHER_FIELD.call(vec![base, lit(name.to_string())]) +} diff --git a/crates/lance-graph/src/datafusion_planner/expression.rs b/crates/lance-graph/src/datafusion_planner/expression.rs index 73cb5f4..1759939 100644 --- a/crates/lance-graph/src/datafusion_planner/expression.rs +++ b/crates/lance-graph/src/datafusion_planner/expression.rs @@ -168,6 +168,26 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { Box::new(to_df_value_expr(&args[0])), arrow_schema::DataType::Utf8, )), + "head" => datafusion::functions_nested::expr_fn::array_element( + to_df_value_expr(&args[0]), + lit(1_i64), + ), + "tointeger" => cast_to(to_df_value_expr(&args[0]), arrow_schema::DataType::Int64), + "tofloat" => cast_to(to_df_value_expr(&args[0]), arrow_schema::DataType::Float64), + "date" => cast_to(to_df_value_expr(&args[0]), arrow_schema::DataType::Date32), + "datetime" => cast_to( + to_df_value_expr(&args[0]), + arrow_schema::DataType::Timestamp( + arrow_schema::TimeUnit::Nanosecond, + Some("UTC".into()), + ), + ), + // A duration between two timestamps, as the whole seconds it spans. + "duration.inseconds" => datafusion::functions::core::expr_fn::named_struct(vec![ + lit("seconds"), + (epoch_ms(to_df_value_expr(&args[1])) - epoch_ms(to_df_value_expr(&args[0]))) + / lit(1000_i64), + ]), // A path's length and node ids are columns of its variable (see shortest_path). "length" | "nodes" => match &args[0] { VE::Variable(path) => col(qualify_column(path, &name.to_lowercase())), @@ -197,17 +217,7 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { ), // Milliseconds since the Unix epoch, as Kuzu's to_epoch_ms: timestamps of any unit // are first cast to millisecond precision, then to their raw Int64 value. - "to_epoch_ms" => { - use arrow_schema::{DataType, TimeUnit}; - let ms = Expr::Cast(datafusion::logical_expr::Cast::new( - Box::new(to_df_value_expr(&args[0])), - DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())), - )); - Expr::Cast(datafusion::logical_expr::Cast::new( - Box::new(ms), - DataType::Int64, - )) - } + "to_epoch_ms" => epoch_ms(to_df_value_expr(&args[0])), _ => { // Unknown scalar function - return NULL Expr::Literal(datafusion::scalar::ScalarValue::Null, None) @@ -294,7 +304,13 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { // Cypher's collect skips nulls; DataFusion's array_agg keeps them. use datafusion::logical_expr::ExprFunctionExt; let arg_expr = to_df_value_expr(&args[0]); - let agg = array_agg(arg_expr.clone()).filter(arg_expr.is_not_null()); + // A map literal is never null; repeating it in a FILTER would make + // DataFusion extract it as a common expression, which loses the + // ordering of `collect` after `WITH ... ORDER BY` (LDBC IC7). + let agg = match &args[0] { + VE::Map(_) => array_agg(arg_expr).filter(lit(true)), + _ => array_agg(arg_expr.clone()).filter(arg_expr.is_not_null()), + }; let agg = if *distinct { agg.distinct() } else { agg }; agg.build().expect("FILTER and DISTINCT apply to array_agg") } else { @@ -400,15 +416,106 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { name ); } + VE::Boolean(b) => to_df_boolean_expr(b), + VE::Map(entries) => datafusion::functions::core::expr_fn::named_struct( + entries + .iter() + .flat_map(|(k, v)| [lit(k.clone()), to_df_value_expr(v)]) + .collect(), + ), + VE::Field { base, field } => super::cypher_udf::field(to_df_value_expr(base), field), } } +/// Milliseconds since the Unix epoch of a timestamp of any unit. +fn epoch_ms(expr: Expr) -> Expr { + use arrow_schema::{DataType, TimeUnit}; + let ms = Expr::Cast(datafusion::logical_expr::Cast::new( + Box::new(expr), + DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())), + )); + Expr::Cast(datafusion::logical_expr::Cast::new( + Box::new(ms), + DataType::Int64, + )) +} + +fn cast_to(expr: Expr, data_type: arrow_schema::DataType) -> Expr { + Expr::Cast(datafusion::logical_expr::Cast::new( + Box::new(expr), + data_type, + )) +} + +/// `expr` with names resolved against `schema`, the input it is evaluated on: +/// `x.y` where `x` is a value column (a map or date from WITH, not a node) becomes a field +/// access, and a node inside a map literal becomes a map of its property columns. +pub(crate) fn bind_value( + expr: &ValueExpression, + schema: &datafusion::common::DFSchema, +) -> ValueExpression { + let has = |name: &str| schema.has_column_with_unqualified_name(name); + let mut expr = expr.clone(); + crate::ast::transform_value(&mut expr, &mut |e| match e { + ValueExpression::Property(p) + if !has(&qualify_column(&p.variable, &p.property)) + && has(&p.variable.to_lowercase()) => + { + *e = ValueExpression::Field { + base: Box::new(ValueExpression::Variable(p.variable.clone())), + field: p.property.clone(), + }; + } + ValueExpression::Map(entries) => { + for (_, value) in entries.iter_mut() { + let ValueExpression::Variable(var) = value else { + continue; + }; + let var = var.to_lowercase(); + let prefix = format!("{var}__"); + if has(&var) { + continue; + } + let properties: Vec<(String, ValueExpression)> = schema + .fields() + .iter() + .filter_map(|f| f.name().strip_prefix(&prefix)) + .map(|property| { + let column = ValueExpression::Property(crate::ast::PropertyRef { + variable: var.clone(), + property: property.to_string(), + }); + (property.to_string(), column) + }) + .collect(); + if !properties.is_empty() { + *value = ValueExpression::Map(properties); + } + } + } + _ => {} + }); + expr +} + +/// [`bind_value`] for every value in a predicate. +pub(crate) fn bind_boolean( + expr: &BooleanExpression, + schema: &datafusion::common::DFSchema, +) -> BooleanExpression { + let mut expr = expr.clone(); + crate::ast::transform_boolean_values(&mut expr, &mut |v| *v = bind_value(v, schema)); + expr +} + /// Check if a ValueExpression contains an aggregate function pub(crate) fn contains_aggregate(expr: &ValueExpression) -> bool { use crate::ast::ValueExpression as VE; match expr { VE::AggregateFunction { .. } => true, VE::ScalarFunction { args, .. } => args.iter().any(contains_aggregate), + VE::Map(entries) => entries.iter().any(|(_, v)| contains_aggregate(v)), + VE::Field { base, .. } => contains_aggregate(base), VE::Arithmetic { left, right, .. } => contains_aggregate(left) || contains_aggregate(right), VE::Case { whens, otherwise } => { whens.iter().any(|(_, then)| contains_aggregate(then)) @@ -465,6 +572,7 @@ pub(crate) fn to_cypher_column_name(expr: &ValueExpression) -> String { name.to_lowercase() } } + VE::Field { base, field } => format!("{}.{}", to_cypher_column_name(base), field), _ => { // For other expressions (literals, arithmetic), use a generic name // In practice, these should always have explicit aliases diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index 731c00f..ab1c0e6 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -460,6 +460,13 @@ impl<'a> LogicalPlanner<'a> { self.resolve_value(e); } } + ValueExpression::Boolean(b) => self.resolve_boolean(b), + // A node inside a map keeps its properties (see `bind_value`), not just its id. + ValueExpression::Map(entries) => entries + .iter_mut() + .filter(|(_, v)| !matches!(v, ValueExpression::Variable(_))) + .for_each(|(_, v)| self.resolve_value(v)), + ValueExpression::Field { base, .. } => self.resolve_value(base), _ => {} } } @@ -724,6 +731,14 @@ impl<'a> LogicalPlanner<'a> { input = self.rewrite_exists_in_value(input, left)?; self.rewrite_exists_in_value(input, right) } + ValueExpression::Boolean(b) => self.rewrite_exists(input, b), + ValueExpression::Map(entries) => { + for (_, v) in entries.iter_mut() { + input = self.rewrite_exists_in_value(input, v)?; + } + Ok(input) + } + ValueExpression::Field { base, .. } => self.rewrite_exists_in_value(input, base), _ => Ok(input), } } diff --git a/crates/lance-graph/src/parameter_substitution.rs b/crates/lance-graph/src/parameter_substitution.rs index d0fdaae..0d2f55a 100644 --- a/crates/lance-graph/src/parameter_substitution.rs +++ b/crates/lance-graph/src/parameter_substitution.rs @@ -288,6 +288,13 @@ fn substitute_in_value_expression( substitute_in_value_expression(otherwise, parameters)?; } } + ValueExpression::Boolean(b) => substitute_in_boolean_expression(b, parameters)?, + ValueExpression::Map(entries) => { + for (_, v) in entries { + substitute_in_value_expression(v, parameters)?; + } + } + ValueExpression::Field { base, .. } => substitute_in_value_expression(base, parameters)?, _ => {} } Ok(()) diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index b6b4d3a..8645c32 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -666,14 +666,52 @@ fn basic_value_expression(input: &str) -> IResult<&str, ValueExpression> { alt(( parse_vector_literal, // Try vector literal first [0.1, 0.2] list_literal, // Then a list of any expressions [a.x, 'y'] (or []) - parse_parameter, // Try $parameter - function_call, // Regular function calls + map_literal, + parse_parameter, // Try $parameter + function_call, // Regular function calls map(property_value, ValueExpression::Literal), // Try literals BEFORE property references map(property_reference, ValueExpression::Property), map(identifier, |id| ValueExpression::Variable(id.to_string())), ))(input) } +// A value where a predicate may also appear (projections, CASE, lists, arguments, maps): +// `path IS NULL`, `NOT (a)-[:R]-(b)`, `x IN list` are boolean values. +fn expression(input: &str) -> IResult<&str, ValueExpression> { + alt(( + map(boolean_expression, |b| { + ValueExpression::Boolean(Box::new(b)) + }), + value_expression, + ))(input) +} + +// `{key: value, ...}` +fn map_literal(input: &str) -> IResult<&str, ValueExpression> { + let entry = tuple(( + multispace0, + identifier, + multispace0, + char(':'), + multispace0, + expression, + )); + let (input, entries) = delimited( + tuple((char('{'), multispace0)), + separated_list0(tuple((multispace0, char(','))), entry), + tuple((multispace0, char('}'))), + )(input)?; + Ok(( + input, + ValueExpression::Map( + entries + .into_iter() + .map(|(_, k, _, _, _, v)| (k.to_string(), v)) + .collect(), + ), + )) +} + // Parse a value expression: arithmetic over primaries, `*` `/` `%` binding tighter than // `+` `-`, all left-associative. fn value_expression(input: &str) -> IResult<&str, ValueExpression> { @@ -741,18 +779,29 @@ fn primary_value_expression(input: &str) -> IResult<&str, ValueExpression> { case_expression, single_value_expression, ))(input)?; - // `list[i]`, with the `[` right after the value (so `x IN [1, 2]` is not an index). - while let Ok((rest, index)) = delimited( - tuple((char('['), multispace0)), - value_expression, - tuple((multispace0, char(']'))), - )(input) - { - value = ValueExpression::ScalarFunction { - name: "__index".to_string(), - args: vec![value, index], - }; - input = rest; + // `list[i]`, with the `[` right after the value (so `x IN [1, 2]` is not an index), and + // `value.field`. + loop { + if let Ok((rest, index)) = delimited( + tuple((char('['), multispace0)), + value_expression, + tuple((multispace0, char(']'))), + )(input) + { + value = ValueExpression::ScalarFunction { + name: "__index".to_string(), + args: vec![value, index], + }; + input = rest; + } else if let Ok((rest, field)) = preceded(char('.'), identifier)(input) { + value = ValueExpression::Field { + base: Box::new(value), + field: field.to_string(), + }; + input = rest; + } else { + break; + } } Ok((input, value)) } @@ -761,16 +810,16 @@ fn primary_value_expression(input: &str) -> IResult<&str, ValueExpression> { fn case_expression(input: &str) -> IResult<&str, ValueExpression> { let keyword = |k: &'static str| tuple((multispace0, tag_no_case(k), multispace1)); let (input, _) = tuple((tag_no_case("CASE"), multispace1))(input)?; - let (input, operand) = opt(terminated(value_expression, peek(keyword("WHEN"))))(input)?; + let (input, operand) = opt(terminated(expression, peek(keyword("WHEN"))))(input)?; let (input, whens) = match &operand { None => many1(tuple(( preceded(keyword("WHEN"), boolean_expression), - preceded(keyword("THEN"), value_expression), + preceded(keyword("THEN"), expression), )))(input)?, Some(operand) => { let (input, pairs) = many1(tuple(( - preceded(keyword("WHEN"), value_expression), - preceded(keyword("THEN"), value_expression), + preceded(keyword("WHEN"), expression), + preceded(keyword("THEN"), expression), )))(input)?; let whens = pairs .into_iter() @@ -786,7 +835,7 @@ fn case_expression(input: &str) -> IResult<&str, ValueExpression> { (input, whens) } }; - let (input, otherwise) = opt(preceded(keyword("ELSE"), value_expression))(input)?; + let (input, otherwise) = opt(preceded(keyword("ELSE"), expression))(input)?; let (input, _) = tuple((multispace0, tag_no_case("END")))(input)?; Ok(( input, @@ -900,7 +949,8 @@ fn parse_parameter(input: &str) -> IResult<&str, ValueExpression> { // Parse a function call: function_name(args) fn function_call(input: &str) -> IResult<&str, ValueExpression> { - let (input, name) = identifier(input)?; + // Namespaced names such as `duration.inSeconds`. + let (input, name) = recognize(pair(identifier, many0(preceded(char('.'), identifier))))(input)?; let (input, _) = multispace0(input)?; let (input, _) = char('(')(input)?; let (input, _) = multispace0(input)?; @@ -935,10 +985,8 @@ fn function_call(input: &str) -> IResult<&str, ValueExpression> { } // Parse regular function arguments - let (input, args) = separated_list0( - tuple((multispace0, char(','), multispace0)), - value_expression, - )(input)?; + let (input, args) = + separated_list0(tuple((multispace0, char(','), multispace0)), expression)(input)?; let (input, _) = multispace0(input)?; let (input, _) = char(')')(input)?; @@ -1042,10 +1090,7 @@ fn parse_vector_literal(input: &str) -> IResult<&str, ValueExpression> { fn list_literal(input: &str) -> IResult<&str, ValueExpression> { let (input, items) = delimited( tuple((char('['), multispace0)), - separated_list0( - tuple((multispace0, char(','), multispace0)), - value_expression, - ), + separated_list0(tuple((multispace0, char(','), multispace0)), expression), tuple((multispace0, char(']'))), )(input)?; Ok(( @@ -1134,7 +1179,7 @@ fn return_clause(input: &str) -> IResult<&str, ReturnClause> { // Parse a return item fn return_item(input: &str) -> IResult<&str, ReturnItem> { - let (input, expression) = value_expression(input)?; + let (input, expression) = expression(input)?; let (input, _) = multispace0(input)?; let (input, alias) = opt(preceded( tuple((tag_no_case("AS"), multispace1)), diff --git a/crates/lance-graph/src/pattern_order.rs b/crates/lance-graph/src/pattern_order.rs index 8915797..6d496f2 100644 --- a/crates/lance-graph/src/pattern_order.rs +++ b/crates/lance-graph/src/pattern_order.rs @@ -231,6 +231,9 @@ fn value_variables<'a>(expr: &'a ValueExpression, out: &mut HashSet<&'a str>) { value_variables(otherwise, out); } } + ValueExpression::Boolean(b) => boolean_variables(b, out), + ValueExpression::Map(entries) => entries.iter().for_each(|(_, v)| value_variables(v, out)), + ValueExpression::Field { base, .. } => value_variables(base, out), ValueExpression::Literal(_) | ValueExpression::Parameter(_) | ValueExpression::VectorLiteral(_) => {} diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index 47818fd..561f1ac 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -469,6 +469,23 @@ impl SemanticAnalyzer { } } "__list" => {} + "head" | "tointeger" | "tofloat" | "date" | "datetime" => { + if args.len() != 1 { + return Err(GraphError::PlanError { + message: format!("{} requires exactly 1 argument", name), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + } + "duration.inseconds" => { + if args.len() != 2 { + return Err(GraphError::PlanError { + message: "duration.inSeconds requires exactly 2 arguments" + .to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + } "tostring" => { if args.len() != 1 { return Err(GraphError::PlanError { @@ -674,6 +691,13 @@ impl SemanticAnalyzer { ValueExpression::Parameter(_) => { // Parameters are always valid (resolved at runtime) } + ValueExpression::Boolean(b) => self.analyze_boolean_expression(b)?, + ValueExpression::Map(entries) => { + for (_, v) in entries { + self.analyze_value_expression(v)?; + } + } + ValueExpression::Field { base, .. } => self.analyze_value_expression(base)?, } Ok(()) } diff --git a/python/python/tests/test_engine_values.py b/python/python/tests/test_engine_values.py index 23698fb..2a2d57d 100644 --- a/python/python/tests/test_engine_values.py +++ b/python/python/tests/test_engine_values.py @@ -113,6 +113,47 @@ def utc(*args): ("MATCH (p:Person {id: 1}) RETURN [p.name, p.id] AS l", [(["a", "1"],)]), ] +# Predicates as values, maps, fields of values, and collect() after ORDER BY. +CASES += [ + ( + "MATCH (a:Person)-[:KNOWS]->(b) " + "RETURN a.id, b.id, (a)-[:KNOWS]->(:Person)-[:KNOWS]->(b) AS viaOne", + [(1, 2, False), (1, 3, True), (2, 3, False), (3, 4, False)], + ), + ( + "MATCH (p:Person) OPTIONAL MATCH (p)<-[:HAS_CREATOR]-(post:Post) " + "RETURN DISTINCT p.id, post IS NULL AS none", + [(1, True), (2, False), (3, False), (4, True)], + ), + ( + "MATCH (p:Person {id: 3})<-[:HAS_CREATOR]-(post:Post) WITH p, post " + "ORDER BY post.created DESC " + "WITH p, head(collect({id: post.id, at: post.created})) AS latest " + "RETURN p.id, latest.id", + [(3, 12)], + ), + ( + "MATCH (p:Person {id: 3})<-[:HAS_CREATOR]-(post:Post) WITH p, post " + "ORDER BY post.created ASC " + "WITH p, head(collect({id: post.id, at: post.created})) AS first RETURN first.id", + [(11,)], + ), + ( + "MATCH (post:Post {id: 11}) RETURN post.created.month, post.created.year", + [(6, 2020)], + ), + ( + "MATCH (a:Post {id: 10}), (b:Post {id: 11}) " + "RETURN toInteger(duration.inSeconds(a.created, b.created).seconds / 86400.0) AS days", + [(152,)], + ), + ( + "MATCH (a:Person {id: 4}), (b:Person {id: 1}), p = shortestPath((a)-[:KNOWS*]-(b)) " + "RETURN CASE p IS NULL WHEN true THEN -1 ELSE length(p) END AS len", + [(2,)], + ), +] + def rows(table): return sorted(tuple(r.values()) for r in table.to_pylist()) From 2c40940ad8bfe5320ebc877ed7e36560427b55e0 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Sat, 24 Oct 2026 09:35:00 -0400 Subject: [PATCH 3/5] feat(cypher): list and pattern comprehensions [x IN list WHERE p | v] unnests each row's list with positions (keeping rows whose list is empty), evaluates the predicate (pattern predicates included) and value, and collects back in order. [(a)-[:R]->(b) WHERE p | v] collects per binding of the row's variables, like an OPTIONAL MATCH, and is [] without matches. --- crates/lance-graph/src/ast.rs | 31 ++++ .../src/datafusion_planner/analysis.rs | 21 +++ .../builder/comprehension_ops.rs | 170 ++++++++++++++++++ .../src/datafusion_planner/builder/helpers.rs | 6 + .../src/datafusion_planner/builder/mod.rs | 25 +++ .../src/datafusion_planner/expression.rs | 3 + crates/lance-graph/src/logical_plan.rs | 164 +++++++++++++++++ .../lance-graph/src/parameter_substitution.rs | 17 ++ crates/lance-graph/src/parser.rs | 49 ++++- crates/lance-graph/src/pattern_order.rs | 24 +++ crates/lance-graph/src/semantic.rs | 39 ++++ python/python/tests/test_engine_values.py | 59 +++++- 12 files changed, 598 insertions(+), 10 deletions(-) create mode 100644 crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs diff --git a/crates/lance-graph/src/ast.rs b/crates/lance-graph/src/ast.rs index 495ced7..52d40c7 100644 --- a/crates/lance-graph/src/ast.rs +++ b/crates/lance-graph/src/ast.rs @@ -408,6 +408,20 @@ pub enum ValueExpression { base: Box, field: String, }, + /// `[(a)-[:R]->(b) WHERE predicate | value]`: the values for each match of the pattern + /// from this row's bound variables. + PatternComprehension { + pattern: Box, + filter: Option>, + map: Box, + }, + /// `[x IN list WHERE predicate | value]`, both parts optional. + ListComprehension { + variable: String, + list: Box, + filter: Option>, + map: Option>, + }, /// Vector literal: [0.1, 0.2, 0.3] /// Represents an inline vector for similarity search VectorLiteral(Vec), @@ -610,6 +624,23 @@ pub fn transform_value(expr: &mut ValueExpression, f: &mut dyn FnMut(&mut ValueE entries.iter_mut().for_each(|(_, v)| transform_value(v, f)) } ValueExpression::Field { base, .. } => transform_value(base, f), + ValueExpression::ListComprehension { + list, filter, map, .. + } => { + transform_value(list, f); + if let Some(b) = filter { + transform_boolean_values(b, f); + } + if let Some(m) = map { + transform_value(m, f); + } + } + ValueExpression::PatternComprehension { filter, map, .. } => { + if let Some(b) = filter { + transform_boolean_values(b, f); + } + transform_value(map, f); + } ValueExpression::Variable(_) | ValueExpression::Property(_) | ValueExpression::Literal(_) diff --git a/crates/lance-graph/src/datafusion_planner/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index 0da47f7..e573d8e 100644 --- a/crates/lance-graph/src/datafusion_planner/analysis.rs +++ b/crates/lance-graph/src/datafusion_planner/analysis.rs @@ -38,6 +38,8 @@ pub struct RelationshipInstance { pub struct PlanningContext<'a> { pub analysis: &'a QueryAnalysis, pub(crate) relationship_instance_idx: HashMap, + /// Element rows of the list comprehensions being built, innermost last. + pub(crate) comprehension_elements: Vec, } impl<'a> PlanningContext<'a> { @@ -45,6 +47,7 @@ impl<'a> PlanningContext<'a> { Self { analysis, relationship_instance_idx: HashMap::new(), + comprehension_elements: Vec::new(), } } @@ -236,6 +239,24 @@ fn analyze_operator( analyze_operator(input, analysis, rel_counter)?; } } + // Same order as `build_list_comprehension`: the rows, then the element plan. + LogicalOperator::ListComprehension { + input, + variable, + label, + element, + .. + } => { + analyze_operator(input, analysis, rel_counter)?; + if let Some(label) = label { + analysis + .var_to_label + .insert(variable.clone(), label.clone()); + analysis.required_datasets.insert(label.clone()); + } + analyze_operator(element, analysis, rel_counter)?; + } + LogicalOperator::ComprehensionElement => {} } Ok(()) } diff --git a/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs new file mode 100644 index 0000000..43b77d2 --- /dev/null +++ b/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs @@ -0,0 +1,170 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! List comprehensions: `[x IN list WHERE keep | map]`, per row. + +use std::collections::HashSet; + +use arrow_schema::DataType; +use datafusion::common::Column; +use datafusion::functions::core::expr_fn::coalesce; +use datafusion::functions_aggregate::array_agg::array_agg; +use datafusion::functions_nested::expr_fn::{array_element, cardinality, gen_series, make_array}; +use datafusion::functions_window::expr_fn::row_number; +use datafusion::logical_expr::{ + col, lit, Cast, Expr, ExprFunctionExt, JoinType, LogicalPlan, LogicalPlanBuilder, +}; + +use super::super::analysis::PlanningContext; +use super::super::expression::{bind_boolean, bind_value, to_df_boolean_expr, to_df_value_expr}; +use crate::ast::{BooleanExpression, ValueExpression}; +use crate::case_insensitive::qualify_column; +use crate::datafusion_planner::DataFusionPlanner; +use crate::error::Result; +use crate::logical_plan::LogicalOperator; + +fn int64(expr: Expr) -> Expr { + Expr::Cast(Cast::new(Box::new(expr), DataType::Int64)) +} + +impl DataFusionPlanner { + /// Number the input rows, unnest each row's list as (position, element) rows (one + /// padding row for an empty or null list, so every input row survives), build + /// `element` over them, then collect the kept values back per input row by position. + #[allow(clippy::too_many_arguments)] + pub(crate) fn build_list_comprehension( + &self, + ctx: &mut PlanningContext, + input: &LogicalOperator, + list: &ValueExpression, + (variable, label): (&str, Option<&str>), + element: &LogicalOperator, + keep: Option<&BooleanExpression>, + map: &ValueExpression, + output: &str, + ) -> Result { + let error = |what: &str| { + let what = what.to_string(); + move |e| self.plan_error(&what, e) + }; + let row = format!("{output}__row"); + let pos = format!("{output}__pos"); + let numbered = LogicalPlanBuilder::from(self.build_operator(ctx, input)?) + .window(vec![row_number().alias(&row)]) + .map_err(error("Failed to number comprehension rows"))? + .build() + .map_err(error("Failed to number comprehension rows"))?; + let list = to_df_value_expr(&bind_value(list, numbered.schema())); + let length = int64(coalesce(vec![cardinality(list.clone()), lit(0_u64)])); + let last = datafusion::functions::core::expr_fn::greatest(vec![length.clone(), lit(1_i64)]); + let positions = gen_series(lit(1_i64), last, lit(1_i64)); + let exploded = LogicalPlanBuilder::from(numbered.clone()) + .project(vec![ + datafusion::logical_expr::wildcard(), + datafusion::logical_expr::select_expr::SelectExpr::Expression( + positions.alias(&pos), + ), + ]) + .and_then(|b| b.unnest_column(pos.as_str())) + .map_err(error("Failed to unnest comprehension list"))?; + let item = array_element(list, col(&pos)); + let element_rows = match label { + None => exploded + .project(vec![ + datafusion::logical_expr::wildcard(), + datafusion::logical_expr::select_expr::SelectExpr::Expression( + item.alias(variable.to_lowercase()), + ), + ]) + .and_then(|b| b.build()) + .map_err(error("Failed to project comprehension elements"))?, + Some(label) => { + // Node ids: each element joins to its node, keeping padding rows. + let id_column = format!("{output}__id"); + let mapping = self.config.get_node_mapping(label).ok_or_else(|| { + crate::error::GraphError::PlanError { + message: format!("No node mapping for '{label}'"), + location: snafu::Location::new(file!(), line!(), column!()), + } + })?; + let nodes = self.build_scan(ctx, variable, label, &Default::default())?; + let joined = exploded + .project(vec![ + datafusion::logical_expr::wildcard(), + datafusion::logical_expr::select_expr::SelectExpr::Expression( + item.alias(&id_column), + ), + ]) + .and_then(|b| { + b.join( + nodes, + JoinType::Left, + ( + vec![id_column.clone()], + vec![qualify_column(variable, &mapping.id_field)], + ), + None, + ) + }) + .and_then(|b| b.build()) + .map_err(error("Failed to join comprehension elements to nodes"))?; + let kept: Vec = joined + .schema() + .columns() + .into_iter() + .filter(|c| c.name != id_column) + .map(Expr::Column) + .collect(); + LogicalPlanBuilder::from(joined) + .project(kept) + .and_then(|b| b.build()) + .map_err(error("Failed to project comprehension elements"))? + } + }; + + ctx.comprehension_elements.push(element_rows); + let elements = self.build_operator(ctx, element); + ctx.comprehension_elements.pop(); + let elements = elements?; + + // Padding rows (position past the list's end) are never collected. + let mut kept = col(&pos).lt_eq(length); + if let Some(keep) = keep { + kept = kept.and(to_df_boolean_expr(&bind_boolean(keep, elements.schema()))); + } + let value = to_df_value_expr(&bind_value(map, elements.schema())); + let collected = array_agg(value) + .filter(kept) + .order_by(vec![col(&pos).sort(true, false)]) + .build() + .map_err(error("Failed to collect comprehension values"))? + .alias(output); + // One group per input row: its number plus every input column, all equal within it. + let names: HashSet = numbered + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect(); + let group: Vec = elements + .schema() + .columns() + .into_iter() + .filter(|c| names.contains(&c.name)) + .map(Expr::Column) + .collect(); + let mut out: Vec = numbered + .schema() + .fields() + .iter() + .filter(|f| *f.name() != row) + .map(|f| Expr::Column(Column::new_unqualified(f.name()))) + .collect(); + out.push(coalesce(vec![col(output), make_array(vec![])]).alias(output)); + LogicalPlanBuilder::from(elements) + .aggregate(group, vec![collected]) + .and_then(|b| b.project(out)) + .and_then(|b| b.build()) + .map_err(error("Failed to build list comprehension")) + } +} diff --git a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs index f813d3c..173cfbf 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs @@ -109,6 +109,10 @@ impl DataFusionPlanner { Self::collect_variables(first, vars); } } + LogicalOperator::ListComprehension { input, .. } => { + Self::collect_variables(input, vars) + } + LogicalOperator::ComprehensionElement => {} } } @@ -171,6 +175,8 @@ impl DataFusionPlanner { } LogicalOperator::ShortestPath { input, .. } => Self::collect_labels(input, labels), LogicalOperator::Union { .. } => {} + LogicalOperator::ListComprehension { input, .. } => Self::collect_labels(input, labels), + LogicalOperator::ComprehensionElement => {} } } } diff --git a/crates/lance-graph/src/datafusion_planner/builder/mod.rs b/crates/lance-graph/src/datafusion_planner/builder/mod.rs index 2068a1c..cb2d705 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/mod.rs @@ -14,6 +14,7 @@ mod aggregate_ops; mod basic_ops; +mod comprehension_ops; mod expand_ops; mod helpers; mod join_builder; @@ -104,6 +105,30 @@ impl DataFusionPlanner { label, } => self.build_unwind(ctx, input, expression, alias, label.as_deref()), LogicalOperator::Union { inputs } => self.build_union(ctx, inputs), + LogicalOperator::ListComprehension { + input, + list, + variable, + label, + element, + keep, + map, + output, + } => self.build_list_comprehension( + ctx, + input, + list, + (variable, label.as_deref()), + element, + keep.as_ref(), + map, + output, + ), + LogicalOperator::ComprehensionElement => Ok(ctx + .comprehension_elements + .last() + .expect("element rows are built inside their comprehension") + .clone()), LogicalOperator::OptionalMatch { input, optional, diff --git a/crates/lance-graph/src/datafusion_planner/expression.rs b/crates/lance-graph/src/datafusion_planner/expression.rs index 1759939..8b2bd0a 100644 --- a/crates/lance-graph/src/datafusion_planner/expression.rs +++ b/crates/lance-graph/src/datafusion_planner/expression.rs @@ -424,6 +424,9 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { .collect(), ), VE::Field { base, field } => super::cypher_udf::field(to_df_value_expr(base), field), + VE::ListComprehension { .. } | VE::PatternComprehension { .. } => { + unreachable!("the logical planner turns comprehensions into operators") + } } } diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index ab1c0e6..e12cae6 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -154,6 +154,26 @@ pub enum LogicalOperator { optional: Box, keys: Vec, }, + + /// `[variable IN list WHERE keep | map]` for each row of `input`, into the column + /// `output`. The list is unnested with element positions (an empty or null list keeps + /// its row), `element` is planned over those element rows (its leaf is + /// [`LogicalOperator::ComprehensionElement`]; EXISTS marks of `keep` join onto it), and + /// the kept values are collected back per row in position order. With `label`, the + /// list holds node ids and each element is joined to its node. + ListComprehension { + input: Box, + list: ValueExpression, + variable: String, + label: Option, + element: Box, + keep: Option, + map: ValueExpression, + output: String, + }, + + /// The element rows of the enclosing [`LogicalOperator::ListComprehension`]. + ComprehensionElement, } /// Projection item for SELECT/RETURN clauses @@ -194,6 +214,10 @@ pub struct LogicalPlanner<'a> { exists_marks: usize, /// Aliases holding lists of node ids (`collect(node) AS friends`), with the label. node_lists: HashMap, + /// List comprehensions planned so far (output columns `__list_`). + comprehensions: usize, + /// The node label of each path variable's nodes, for `nodes(path)`. + path_labels: HashMap, } impl<'a> LogicalPlanner<'a> { @@ -204,6 +228,8 @@ impl<'a> LogicalPlanner<'a> { fan_out: None, exists_marks: 0, node_lists: HashMap::new(), + comprehensions: 0, + path_labels: HashMap::new(), } } @@ -739,6 +765,134 @@ impl<'a> LogicalPlanner<'a> { Ok(input) } ValueExpression::Field { base, .. } => self.rewrite_exists_in_value(input, base), + ValueExpression::PatternComprehension { + pattern, + filter, + map, + } => { + // As an EXISTS mark, but collecting `map` per binding of the shared + // variables instead of marking `true`. + self.comprehensions += 1; + let column = format!("__pattern_{}", self.comprehensions); + let inner = MatchClause { + patterns: vec![GraphPattern::Path((**pattern).clone())], + optional: true, + where_clause: filter.as_ref().map(|w| WhereClause { + expression: (**w).clone(), + }), + }; + let before = self.variables.clone(); + let planned = (|| -> Result<(LogicalOperator, ValueExpression)> { + let planned = self.plan_optional_match(Some(input), &inner)?; + let mut value = (**map).clone(); + let LogicalOperator::OptionalMatch { + input, + optional, + keys, + } = planned + else { + unreachable!("plan_optional_match returns an OptionalMatch") + }; + let optional = self.rewrite_exists_in_value(*optional, &mut value)?; + self.resolve_value(&mut value); + let mut projections: Vec = keys + .iter() + .map(|k| ProjectionItem { + expression: ValueExpression::Variable( + k.split("__").next().unwrap_or(k).to_string(), + ), + alias: None, + }) + .collect(); + projections.push(ProjectionItem { + expression: ValueExpression::AggregateFunction { + name: "collect".to_string(), + args: vec![value], + distinct: false, + }, + alias: Some(column.clone()), + }); + let collected = LogicalOperator::Project { + input: Box::new(optional), + projections, + }; + let joined = LogicalOperator::OptionalMatch { + input, + optional: Box::new(collected), + keys, + }; + // No match: the left join gives null, which is an empty list. + let result = ValueExpression::ScalarFunction { + name: "coalesce".to_string(), + args: vec![ + ValueExpression::Variable(column.clone()), + ValueExpression::ScalarFunction { + name: "__list".to_string(), + args: vec![], + }, + ], + }; + Ok((joined, result)) + })(); + self.variables = before; + let (joined, result) = planned?; + *expr = result; + Ok(joined) + } + ValueExpression::ListComprehension { + variable, + list, + filter, + map, + } => { + input = self.rewrite_exists_in_value(input, list)?; + self.resolve_value(list); + let label = match list.as_ref() { + ValueExpression::Variable(l) => self.node_lists.get(l).cloned(), + ValueExpression::ScalarFunction { name, args } + if name.eq_ignore_ascii_case("nodes") => + { + match args.as_slice() { + [ValueExpression::Variable(p)] => self.path_labels.get(p).cloned(), + _ => None, + } + } + _ => None, + }; + self.comprehensions += 1; + let output = format!("__list_{}", self.comprehensions); + // The variable is local: an element of the list, a node when it has a label. + let before = self.variables.clone(); + let kind = label.clone().unwrap_or_else(|| "Unwound".to_string()); + self.variables.insert(variable.clone(), kind); + let lowered = (|| -> Result { + let mut element = LogicalOperator::ComprehensionElement; + let mut keep = filter.as_deref().cloned(); + if let Some(k) = &mut keep { + element = self.rewrite_exists(element, k)?; + self.resolve_boolean(k); + } + let mut value = map + .as_deref() + .cloned() + .unwrap_or_else(|| ValueExpression::Variable(variable.clone())); + element = self.rewrite_exists_in_value(element, &mut value)?; + self.resolve_value(&mut value); + Ok(LogicalOperator::ListComprehension { + input: Box::new(input), + list: (**list).clone(), + variable: variable.clone(), + label, + element: Box::new(element), + keep, + map: value, + output: output.clone(), + }) + })(); + self.variables = before; + *expr = ValueExpression::Variable(output); + lowered + } _ => Ok(input), } } @@ -791,6 +945,9 @@ impl<'a> LogicalPlanner<'a> { // Not a node: the marker keeps it in WITH scope and matches no node mapping. self.variables .insert(path_variable.clone(), "Path".to_string()); + if let Some(label) = self.variables.get(&ends[0]).cloned() { + self.path_labels.insert(path_variable.clone(), label); + } let length = rel.length.as_ref(); Ok(LogicalOperator::ShortestPath { input: Box::new(plan.expect("both end nodes are planned")), @@ -1114,6 +1271,13 @@ impl<'a> LogicalPlanner<'a> { LogicalOperator::OptionalMatch { optional, .. } => { self.extract_variable_from_plan(optional) } + LogicalOperator::ListComprehension { input, .. } => { + self.extract_variable_from_plan(input) + } + LogicalOperator::ComprehensionElement => Err(GraphError::PlanError { + message: "a comprehension's element rows bind no pattern variable".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }), LogicalOperator::ShortestPath { target_variable, .. } => Ok(target_variable.clone()), diff --git a/crates/lance-graph/src/parameter_substitution.rs b/crates/lance-graph/src/parameter_substitution.rs index 0d2f55a..f66765a 100644 --- a/crates/lance-graph/src/parameter_substitution.rs +++ b/crates/lance-graph/src/parameter_substitution.rs @@ -295,6 +295,23 @@ fn substitute_in_value_expression( } } ValueExpression::Field { base, .. } => substitute_in_value_expression(base, parameters)?, + ValueExpression::PatternComprehension { filter, map, .. } => { + if let Some(b) = filter { + substitute_in_boolean_expression(b, parameters)?; + } + substitute_in_value_expression(map, parameters)?; + } + ValueExpression::ListComprehension { + list, filter, map, .. + } => { + substitute_in_value_expression(list, parameters)?; + if let Some(b) = filter { + substitute_in_boolean_expression(b, parameters)?; + } + if let Some(m) = map { + substitute_in_value_expression(m, parameters)?; + } + } _ => {} } Ok(()) diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index 8645c32..9766749 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -665,7 +665,9 @@ fn comparison_operator(input: &str) -> IResult<&str, ComparisonOperator> { fn basic_value_expression(input: &str) -> IResult<&str, ValueExpression> { alt(( parse_vector_literal, // Try vector literal first [0.1, 0.2] - list_literal, // Then a list of any expressions [a.x, 'y'] (or []) + pattern_comprehension, + list_comprehension, + list_literal, // Then a list of any expressions [a.x, 'y'] (or []) map_literal, parse_parameter, // Try $parameter function_call, // Regular function calls @@ -686,6 +688,51 @@ fn expression(input: &str) -> IResult<&str, ValueExpression> { ))(input) } +// `[(a)-[:R]->(b) WHERE predicate | value]` +fn pattern_comprehension(input: &str) -> IResult<&str, ValueExpression> { + let (input, _) = tuple((char('['), multispace0))(input)?; + let (input, pattern) = path_pattern(input)?; + let (input, filter) = opt(where_clause)(input)?; + let (input, map) = preceded(tuple((multispace0, char('|'), multispace0)), expression)(input)?; + let (input, _) = tuple((multispace0, char(']')))(input)?; + Ok(( + input, + ValueExpression::PatternComprehension { + pattern: Box::new(pattern), + filter: filter.map(|w| Box::new(w.expression)), + map: Box::new(map), + }, + )) +} + +// `[x IN list WHERE predicate | value]` +fn list_comprehension(input: &str) -> IResult<&str, ValueExpression> { + let (input, (_, _, variable, _, _, _, list)) = tuple(( + char('['), + multispace0, + identifier, + multispace1, + tag_no_case("IN"), + multispace1, + value_expression, + ))(input)?; + let (input, filter) = opt(where_clause)(input)?; + let (input, map) = opt(preceded( + tuple((multispace0, char('|'), multispace0)), + expression, + ))(input)?; + let (input, _) = tuple((multispace0, char(']')))(input)?; + Ok(( + input, + ValueExpression::ListComprehension { + variable: variable.to_string(), + list: Box::new(list), + filter: filter.map(|w| Box::new(w.expression)), + map: map.map(Box::new), + }, + )) +} + // `{key: value, ...}` fn map_literal(input: &str) -> IResult<&str, ValueExpression> { let entry = tuple(( diff --git a/crates/lance-graph/src/pattern_order.rs b/crates/lance-graph/src/pattern_order.rs index 6d496f2..6f8e004 100644 --- a/crates/lance-graph/src/pattern_order.rs +++ b/crates/lance-graph/src/pattern_order.rs @@ -234,6 +234,30 @@ fn value_variables<'a>(expr: &'a ValueExpression, out: &mut HashSet<&'a str>) { ValueExpression::Boolean(b) => boolean_variables(b, out), ValueExpression::Map(entries) => entries.iter().for_each(|(_, v)| value_variables(v, out)), ValueExpression::Field { base, .. } => value_variables(base, out), + ValueExpression::PatternComprehension { + pattern, + filter, + map, + } => { + let nodes = std::iter::once(&pattern.start_node) + .chain(pattern.segments.iter().map(|s| &s.end_node)); + out.extend(nodes.filter_map(|n| n.variable.as_deref())); + if let Some(b) = filter { + boolean_variables(b, out); + } + value_variables(map, out); + } + ValueExpression::ListComprehension { + list, filter, map, .. + } => { + value_variables(list, out); + if let Some(b) = filter { + boolean_variables(b, out); + } + if let Some(m) = map { + value_variables(m, out); + } + } ValueExpression::Literal(_) | ValueExpression::Parameter(_) | ValueExpression::VectorLiteral(_) => {} diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index 561f1ac..dbe4311 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -698,6 +698,45 @@ impl SemanticAnalyzer { } } ValueExpression::Field { base, .. } => self.analyze_value_expression(base)?, + ValueExpression::PatternComprehension { + pattern, + filter, + map, + } => { + // New variables of the pattern are local to the comprehension. + let before = self.variables.clone(); + let result = (|| -> Result<()> { + self.analyze_graph_pattern(&GraphPattern::Path((**pattern).clone()))?; + if let Some(b) = filter { + self.analyze_boolean_expression(b)?; + } + self.analyze_value_expression(map) + })(); + self.variables = before; + result?; + } + ValueExpression::ListComprehension { + variable, + list, + filter, + map, + } => { + self.analyze_value_expression(list)?; + // The variable is local to the comprehension. + let before = self.variables.clone(); + self.register_projection_alias(variable); + let result = (|| -> Result<()> { + if let Some(b) = filter { + self.analyze_boolean_expression(b)?; + } + if let Some(m) = map { + self.analyze_value_expression(m)?; + } + Ok(()) + })(); + self.variables = before; + result?; + } } Ok(()) } diff --git a/python/python/tests/test_engine_values.py b/python/python/tests/test_engine_values.py index 2a2d57d..fa35a5c 100644 --- a/python/python/tests/test_engine_values.py +++ b/python/python/tests/test_engine_values.py @@ -3,9 +3,10 @@ """Nodes as values, lists of nodes, UNWIND, parameters and inferred labels. -The constructs LDBC's official IC queries use: `a = b` on nodes, `x IN list`, `collect(node)` -then `UNWIND`, inline properties that read a variable, chained comparisons, unlabeled -nodes whose label follows from a relationship's endpoints. Answers are written out by hand. +The constructs LDBC's official IC queries use: `a = b` on nodes, `x IN list`, +`collect(node)` then `UNWIND`, inline properties that read a variable, chained +comparisons, unlabeled nodes whose label follows from a relationship's endpoints. +Answers are written out by hand. """ from datetime import datetime, timezone @@ -56,13 +57,15 @@ def utc(*args): CASES = [ # A node compared with a node compares ids. ( - "MATCH (a:Person {id: 1})-[:KNOWS*1..2]->(b) WHERE NOT a = b RETURN DISTINCT b.id", + "MATCH (a:Person {id: 1})-[:KNOWS*1..2]->(b) WHERE NOT a = b RETURN DISTINCT " + "b.id", [(2,), (3,), (4,)], ), # `p` is a Post: HAS_CREATOR goes from Post to Person. ("MATCH (a:Person {id: 2})<-[:HAS_CREATOR]-(p) RETURN p.id", [(10,)]), ( - "MATCH (x:City {name: 'X'}), (p:Person)-[:LIVES_IN]->(c) WHERE c IN [x] RETURN p.id", + "MATCH (x:City {name: 'X'}), (p:Person)-[:LIVES_IN]->(c) WHERE c IN [x] " + "RETURN p.id", [(1,), (2,)], ), # A list of nodes, and a pattern after WITH that starts at a new node. @@ -91,7 +94,8 @@ def utc(*args): # A scalar from WITH survives a variable-length path. ( "MATCH (c:City {name: 'Y'}) WITH c.id AS cid " - "MATCH (:Person {id: 1})-[:KNOWS*1..2]->(b)-[:LIVES_IN]->(c2) WHERE c2.id = cid " + "MATCH (:Person {id: 1})-[:KNOWS*1..2]->(b)-[:LIVES_IN]->(c2) WHERE c2.id = " + "cid " "RETURN DISTINCT b.id", [(3,), (4,)], ), @@ -135,7 +139,8 @@ def utc(*args): ( "MATCH (p:Person {id: 3})<-[:HAS_CREATOR]-(post:Post) WITH p, post " "ORDER BY post.created ASC " - "WITH p, head(collect({id: post.id, at: post.created})) AS first RETURN first.id", + "WITH p, head(collect({id: post.id, at: post.created})) AS first RETURN " + "first.id", [(11,)], ), ( @@ -144,16 +149,52 @@ def utc(*args): ), ( "MATCH (a:Post {id: 10}), (b:Post {id: 11}) " - "RETURN toInteger(duration.inSeconds(a.created, b.created).seconds / 86400.0) AS days", + "RETURN toInteger(duration.inSeconds(a.created, b.created).seconds / " + "86400.0) AS days", [(152,)], ), ( - "MATCH (a:Person {id: 4}), (b:Person {id: 1}), p = shortestPath((a)-[:KNOWS*]-(b)) " + "MATCH (a:Person {id: 4}), (b:Person {id: 1}), p = " + "shortestPath((a)-[:KNOWS*]-(b)) " "RETURN CASE p IS NULL WHEN true THEN -1 ELSE length(p) END AS len", [(2,)], ), ] +# List and pattern comprehensions. +CASES += [ + ( + "MATCH (p:Person) OPTIONAL MATCH (p)<-[:HAS_CREATOR]-(post:Post) " + "WITH p, collect(post) AS posts RETURN p.id, [x IN posts WHERE x.id > 10 | " + "x.id] AS l", + [(1, []), (2, []), (3, [11, 12]), (4, [])], + ), + # A pattern predicate inside the comprehension (LDBC IC10's shape). + ( + "MATCH (p:Person) OPTIONAL MATCH (p)-[:KNOWS]->(f:Person) WITH p, collect(f) " + "AS fs " + "RETURN p.id, size([f IN fs WHERE (f)-[:LIVES_IN]->(:City {name: 'Y'})]) AS " + "inY", + [(1, 1), (2, 1), (3, 1), (4, 0)], + ), + ( + "MATCH (p:Person) RETURN p.id, [(p)<-[:HAS_CREATOR]-(post) WHERE post.id > " + "10 | post.id] AS l", + [(1, []), (2, []), (3, [11, 12]), (4, [])], + ), + ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) " + "RETURN a.id, b.id, size([(a)-[:KNOWS]->(m)-[:KNOWS]->(b) | m]) AS via", + [(1, 2, 0), (1, 3, 1), (2, 3, 0), (3, 4, 0)], + ), + ( + "MATCH (a:Person {id: 1}), (b:Person {id: 4}), p = " + "shortestPath((a)-[:KNOWS*]->(b)) " + "RETURN [n IN nodes(p) | n.name] AS names", + [(["a", "c", "d"],)], + ), +] + def rows(table): return sorted(tuple(r.values()) for r in table.to_pylist()) From f8bdb94cf84cb4e93aa50fbfc888a8a88fa41753 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Sat, 24 Oct 2026 09:57:13 -0400 Subject: [PATCH 4/5] perf(planner): seed EXISTS marks and comprehensions from the original rows Each mark or comprehension planned its seed from the rows with the previous ones already joined on, so k of them in one projection copied the input 2^k times. Seeding all of them from the rows the projection or WHERE started from keeps the copies linear: LDBC IC14 drops from 641 to 225 physical operators and from 188 to 62 ms. --- crates/lance-graph/src/logical_plan.rs | 51 ++++++++++++++++++++++---- 1 file changed, 43 insertions(+), 8 deletions(-) diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index e12cae6..ef6520f 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -218,6 +218,10 @@ pub struct LogicalPlanner<'a> { comprehensions: usize, /// The node label of each path variable's nodes, for `nodes(path)`. path_labels: HashMap, + /// The rows a projection or WHERE started from, while its EXISTS marks and + /// comprehensions are being added: each is seeded from these, not from the rows with + /// the previous ones joined on, so planned copies grow linearly, not exponentially. + seed_source: Option, } impl<'a> LogicalPlanner<'a> { @@ -230,6 +234,7 @@ impl<'a> LogicalPlanner<'a> { node_lists: HashMap::new(), comprehensions: 0, path_labels: HashMap::new(), + seed_source: None, } } @@ -405,7 +410,7 @@ impl<'a> LogicalPlanner<'a> { ) -> Result { match clause { ReadingClause::Match(match_clause) if match_clause.optional => { - self.plan_optional_match(base, match_clause) + self.plan_optional_match(base, match_clause, None) } ReadingClause::Match(match_clause) => { let plan = self.plan_match_clause_with_base(base, match_clause, filter)?; @@ -644,7 +649,10 @@ impl<'a> LogicalPlanner<'a> { predicate: &BooleanExpression, ) -> Result { let mut predicate = predicate.clone(); - let input = self.rewrite_exists(input, &mut predicate)?; + let saved = self.seed_source.replace(input.clone()); + let input = self.rewrite_exists(input, &mut predicate); + self.seed_source = saved; + let input = input?; self.resolve_boolean(&mut predicate); Ok(LogicalOperator::Filter { input: Box::new(input), @@ -687,7 +695,8 @@ impl<'a> LogicalPlanner<'a> { }; // The pattern's own new variables are local to the subquery. let before = self.variables.clone(); - let planned = self.plan_optional_match(Some(input), &inner); + let seed = self.seed_source.clone(); + let planned = self.plan_optional_match(Some(input), &inner, seed); self.variables = before; let LogicalOperator::OptionalMatch { input, @@ -783,7 +792,8 @@ impl<'a> LogicalPlanner<'a> { }; let before = self.variables.clone(); let planned = (|| -> Result<(LogicalOperator, ValueExpression)> { - let planned = self.plan_optional_match(Some(input), &inner)?; + let seed = self.seed_source.clone(); + let planned = self.plan_optional_match(Some(input), &inner, seed)?; let mut value = (**map).clone(); let LogicalOperator::OptionalMatch { input, @@ -793,7 +803,10 @@ impl<'a> LogicalPlanner<'a> { else { unreachable!("plan_optional_match returns an OptionalMatch") }; - let optional = self.rewrite_exists_in_value(*optional, &mut value)?; + let saved = self.seed_source.replace((*optional).clone()); + let optional = self.rewrite_exists_in_value(*optional, &mut value); + self.seed_source = saved; + let optional = optional?; self.resolve_value(&mut value); let mut projections: Vec = keys .iter() @@ -863,10 +876,12 @@ impl<'a> LogicalPlanner<'a> { let output = format!("__list_{}", self.comprehensions); // The variable is local: an element of the list, a node when it has a label. let before = self.variables.clone(); + let saved_seed = self.seed_source.clone(); let kind = label.clone().unwrap_or_else(|| "Unwound".to_string()); self.variables.insert(variable.clone(), kind); let lowered = (|| -> Result { let mut element = LogicalOperator::ComprehensionElement; + self.seed_source = Some(LogicalOperator::ComprehensionElement); let mut keep = filter.as_deref().cloned(); if let Some(k) = &mut keep { element = self.rewrite_exists(element, k)?; @@ -890,6 +905,7 @@ impl<'a> LogicalPlanner<'a> { }) })(); self.variables = before; + self.seed_source = saved_seed; *expr = ValueExpression::Variable(output); lowered } @@ -971,6 +987,7 @@ impl<'a> LogicalPlanner<'a> { &mut self, base: Option, clause: &MatchClause, + seed_source: Option, ) -> Result { let unsupported = |what: &str| GraphError::UnsupportedFeature { feature: format!("OPTIONAL MATCH {what}"), @@ -1018,7 +1035,7 @@ impl<'a> LogicalPlanner<'a> { } let seed = LogicalOperator::Distinct { input: Box::new(LogicalOperator::Project { - input: Box::new(base.clone()), + input: Box::new(seed_source.unwrap_or_else(|| base.clone())), projections: carried .iter() .map(|var| ProjectionItem { @@ -1300,9 +1317,18 @@ impl<'a> LogicalPlanner<'a> { let mut projections: Vec = Vec::new(); let mut items = return_clause.items.clone(); let mut input = input; + let saved = self.seed_source.replace(input.clone()); for item in &mut items { - input = self.rewrite_exists_in_value(input, &mut item.expression)?; + let rewritten = self.rewrite_exists_in_value(input, &mut item.expression); + input = match rewritten { + Ok(input) => input, + Err(e) => { + self.seed_source = saved; + return Err(e); + } + }; } + self.seed_source = saved; for item in &items { let alias = &item.alias; @@ -1384,9 +1410,18 @@ impl<'a> LogicalPlanner<'a> { let mut projections = Vec::with_capacity(with_clause.items.len()); let mut items = with_clause.items.clone(); let mut input = input; + let saved = self.seed_source.replace(input.clone()); for item in &mut items { - input = self.rewrite_exists_in_value(input, &mut item.expression)?; + let rewritten = self.rewrite_exists_in_value(input, &mut item.expression); + input = match rewritten { + Ok(input) => input, + Err(e) => { + self.seed_source = saved; + return Err(e); + } + }; } + self.seed_source = saved; for item in &items { let alias = match (&item.alias, &item.expression) { (Some(alias), ValueExpression::Variable(var)) => { From 559b803271960080e30280020cf90ef4624e779f Mon Sep 17 00:00:00 2001 From: prrao87 Date: Sat, 24 Oct 2026 10:01:04 -0400 Subject: [PATCH 5/5] docs: document query parameters and the wider Cypher coverage --- README.md | 7 +++++-- crates/lance-graph-python/src/graph.rs | 9 +++++++++ python/README.md | 10 ++++++++++ 3 files changed, 24 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 0c9d964..46ce826 100644 --- a/README.md +++ b/README.md @@ -23,8 +23,11 @@ DataFusion as extension operators, falling back to a join for any hop the CSR ca and `refresh()`es to new dataset versions. - **Smarter planning**: patterns start at their most selective node, and DataFusion gets column statistics to size its joins. -- **Much wider Cypher**, enough for all 14 LDBC Interactive complex queries: `OPTIONAL MATCH`, `EXISTS`, chained - `WITH`, relationship alternation, `shortestPath` and more. +- **Much wider Cypher**, enough to run the official LDBC Interactive complex queries (IC14 without `reduce()`, + which openCypher leaves out): `OPTIONAL MATCH`, `EXISTS` and pattern predicates, chained `WITH`, `UNWIND`, + list and pattern comprehensions, maps, relationship alternation, `shortestPath`, `$parameters` and more, + following openCypher's + [standardisation scope](https://github.com/opencypher/openCypher/blob/master/docs/standardisation-scope.adoc). - **Correctness fixes** to behavior inherited from lance-graph (for example, undirected patterns dropped incoming edges), each with a regression test. - **Current dependencies** (Lance 13, DataFusion 54, Arrow 58, PyO3 0.28), with Delta Lake and Unity Catalog support diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 9c973c9..569b47f 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -1955,6 +1955,15 @@ impl CypherEngine { /// /// The whole result is held in memory; use ``execute_batches`` for large outputs. /// + /// Parameters + /// ---------- + /// query : str + /// The Cypher query + /// params : dict, optional + /// Values for ``$name`` parameters in the query: numbers, strings, booleans, + /// ``None``, and ``datetime``/``date`` values (compared as timestamps and dates). + /// ``execute_batches`` and ``explain`` take the same argument. + /// /// Raises /// ------ /// ValueError diff --git a/python/README.md b/python/README.md index 6537d3f..2ca0e71 100644 --- a/python/README.md +++ b/python/README.md @@ -85,8 +85,18 @@ result3 = engine.execute("MATCH (p:Person) RETURN count(*) as total") print(result1.to_pylist()) # [{'p.name': 'Alice'}] + +# $parameters, including datetime and date values +result4 = engine.execute( + "MATCH (p:Person) WHERE p.age >= $min_age RETURN p.name", {"min_age": 30} +) ``` +Giving each relationship its endpoint labels, as in +`.with_relationship("lives_in", "src", "dst", "Person", "City")`, lets queries leave +nodes unlabeled where the relationship decides the label (`(p:Person)-[:lives_in]->(c)`), +and is required for relationship-type alternation such as `[:likes_post|likes_comment]`. + ### 3. Lance datasets: snapshots, refresh and streaming `CypherEngine` reads Lance datasets lazily from storage (projections and filters are pushed into the