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 653dbc6..569b47f 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 { @@ -1939,15 +1955,30 @@ 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 /// 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 +1996,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 +2024,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 +2154,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..52d40c7 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, @@ -389,6 +398,30 @@ 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, + }, + /// `[(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), @@ -410,7 +443,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 @@ -562,6 +596,94 @@ 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::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(_) + | 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/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index 0d2a7b7..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(), } } @@ -206,10 +209,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)?; @@ -227,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/aggregate_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs index 9528106..0cf71fc 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs @@ -3,24 +3,78 @@ //! 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, 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(); 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(); + 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) { @@ -30,22 +84,31 @@ 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 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)); } 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 +140,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..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); + + 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)) + } - // Convert sort items to DataFusion sort expressions - let sort_exprs: Vec = sort_items + /// `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( @@ -256,6 +276,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)? @@ -267,7 +288,15 @@ 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 { + 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 +313,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/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/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..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 => {} } } @@ -153,16 +157,26 @@ 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) } 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 ddc93af..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; @@ -101,8 +102,33 @@ 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::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/cypher_udf.rs b/crates/lance-graph/src/datafusion_planner/cypher_udf.rs new file mode 100644 index 0000000..e4f47a2 --- /dev/null +++ b/crates/lance-graph/src/datafusion_planner/cypher_udf.rs @@ -0,0 +1,175 @@ +// 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::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::{ + lit, ColumnarValue, Expr, ReturnFieldArgs, 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) +} + +/// `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 ace45c4..8b2bd0a 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,13 +163,31 @@ 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, )), + "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())), @@ -194,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) @@ -291,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 { @@ -397,15 +416,109 @@ 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), + VE::ListComprehension { .. } | VE::PatternComprehension { .. } => { + unreachable!("the logical planner turns comprehensions into operators") + } } } +/// 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)) @@ -436,7 +549,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) @@ -462,6 +575,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/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..ef6520f 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) @@ -151,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 @@ -189,6 +212,16 @@ 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, + /// 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, + /// 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> { @@ -198,6 +231,10 @@ impl<'a> LogicalPlanner<'a> { config, fan_out: None, exists_marks: 0, + node_lists: HashMap::new(), + comprehensions: 0, + path_labels: HashMap::new(), + seed_source: None, } } @@ -316,14 +353,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), }; } @@ -380,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)?; @@ -401,17 +431,124 @@ 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); + } + } + 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), + _ => {} + } + } + + /// [`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 +566,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, @@ -503,7 +649,11 @@ 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), predicate, @@ -545,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, @@ -615,6 +766,149 @@ 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), + 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 seed = self.seed_source.clone(); + let planned = self.plan_optional_match(Some(input), &inner, seed)?; + let mut value = (**map).clone(); + let LogicalOperator::OptionalMatch { + input, + optional, + keys, + } = planned + else { + unreachable!("plan_optional_match returns an OptionalMatch") + }; + 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() + .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 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)?; + 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; + self.seed_source = saved_seed; + *expr = ValueExpression::Variable(output); + lowered + } _ => Ok(input), } } @@ -667,6 +961,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")), @@ -690,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}"), @@ -721,10 +1019,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 + input: Box::new(seed_source.unwrap_or_else(|| base.clone())), + projections: carried .iter() .map(|var| ProjectionItem { expression: ValueExpression::Variable(var.clone()), @@ -843,17 +1155,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 @@ -958,6 +1288,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()), @@ -980,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; @@ -1027,8 +1373,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(), }); } @@ -1062,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)) => { @@ -1083,11 +1440,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 +1502,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..f66765a 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, .. } @@ -284,6 +288,30 @@ 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)?, + 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 a3883fe..9766749 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 @@ -568,15 +665,100 @@ 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 []) - parse_parameter, // Try $parameter - function_call, // Regular function calls + 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 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) +} + +// `[(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(( + 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> { @@ -644,18 +826,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)) } @@ -664,16 +857,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() @@ -689,7 +882,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, @@ -803,7 +996,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)?; @@ -838,10 +1032,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)?; @@ -945,10 +1137,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(( @@ -1037,7 +1226,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_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..6f8e004 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, .. } @@ -226,12 +231,57 @@ 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::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(_) => {} } } +/// 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..dbe4311 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)?; } @@ -458,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 { @@ -663,6 +691,52 @@ 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)?, + 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(()) } @@ -711,10 +785,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 +871,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/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 diff --git a/python/python/tests/test_engine_values.py b/python/python/tests/test_engine_values.py new file mode 100644 index 0000000..fa35a5c --- /dev/null +++ b/python/python/tests/test_engine_values.py @@ -0,0 +1,217 @@ +# 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"],)]), +] + +# 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,)], + ), +] + +# 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()) + + +@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",)]