From f71fe3c481b5c8599c24e904d163b8867eae29af Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 09:58:52 -0700 Subject: [PATCH 01/66] fix(cypher): preserve integer list precision Keep integer lists as Int64 values and convert numeric lists only at vector boundaries. Add singleton-row planning so standalone RETURN examples execute on configured engines. Issue #3. Implemented with assistance from Codex; coordinated runtime and benchmark validation pending. --- .../src/datafusion_planner/analysis.rs | 2 +- .../src/datafusion_planner/builder/helpers.rs | 4 +- .../src/datafusion_planner/builder/mod.rs | 5 +- .../src/datafusion_planner/vector_ops.rs | 88 ++++++++--------- crates/lance-graph/src/logical_plan.rs | 11 ++- crates/lance-graph/src/parser.rs | 44 +++++++++ docs/cypher.md | 11 ++- docs/vector-search.md | 4 + python/python/tests/test_integer_lists.py | 95 +++++++++++++++++++ 9 files changed, 210 insertions(+), 54 deletions(-) create mode 100644 python/python/tests/test_integer_lists.py diff --git a/crates/lance-graph/src/datafusion_planner/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index 89abee0..3329517 100644 --- a/crates/lance-graph/src/datafusion_planner/analysis.rs +++ b/crates/lance-graph/src/datafusion_planner/analysis.rs @@ -256,7 +256,7 @@ fn analyze_operator( } analyze_operator(element, analysis, rel_counter)?; } - LogicalOperator::ComprehensionElement => {} + LogicalOperator::ComprehensionElement | LogicalOperator::SingleRow => {} } Ok(()) } diff --git a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs index f2b8333..34647d5 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs @@ -112,7 +112,7 @@ impl DataFusionPlanner { LogicalOperator::ListComprehension { input, .. } => { Self::collect_variables(input, vars) } - LogicalOperator::ComprehensionElement => {} + LogicalOperator::ComprehensionElement | LogicalOperator::SingleRow => {} } } @@ -176,7 +176,7 @@ impl DataFusionPlanner { LogicalOperator::ShortestPath { input, .. } => Self::collect_labels(input, labels), LogicalOperator::Union { .. } => {} LogicalOperator::ListComprehension { input, .. } => Self::collect_labels(input, labels), - LogicalOperator::ComprehensionElement => {} + LogicalOperator::ComprehensionElement | LogicalOperator::SingleRow => {} } } } diff --git a/crates/lance-graph/src/datafusion_planner/builder/mod.rs b/crates/lance-graph/src/datafusion_planner/builder/mod.rs index cb2d705..c2ca6a4 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/mod.rs @@ -22,7 +22,7 @@ mod join_builder; use super::DataFusionPlanner; use crate::error::Result; use crate::logical_plan::*; -use datafusion::logical_expr::LogicalPlan; +use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder}; use super::analysis::PlanningContext; @@ -34,6 +34,9 @@ impl DataFusionPlanner { op: &LogicalOperator, ) -> Result { match op { + LogicalOperator::SingleRow => LogicalPlanBuilder::empty(true) + .build() + .map_err(|e| self.plan_error("Failed to build singleton row", e)), LogicalOperator::ScanByLabel { variable, label, diff --git a/crates/lance-graph/src/datafusion_planner/vector_ops.rs b/crates/lance-graph/src/datafusion_planner/vector_ops.rs index 794eec9..a831952 100644 --- a/crates/lance-graph/src/datafusion_planner/vector_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/vector_ops.rs @@ -26,18 +26,7 @@ pub fn extract_vectors(array: &ArrayRef) -> Result>> { }); } let value_array = list_array.value(i); - let float_array = value_array - .as_any() - .downcast_ref::() - .ok_or_else(|| GraphError::ExecutionError { - message: "Expected Float32Array in vector".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - })?; - - let vec: Vec = (0..float_array.len()) - .map(|j| float_array.value(j)) - .collect(); - vectors.push(vec); + vectors.push(vector_values(value_array.as_ref())?); } return Ok(vectors); } @@ -53,18 +42,7 @@ pub fn extract_vectors(array: &ArrayRef) -> Result>> { }); } let value_array = list_array.value(i); - let float_array = value_array - .as_any() - .downcast_ref::() - .ok_or_else(|| GraphError::ExecutionError { - message: "Expected Float32Array in vector".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - })?; - - let vec: Vec = (0..float_array.len()) - .map(|j| float_array.value(j)) - .collect(); - vectors.push(vec); + vectors.push(vector_values(value_array.as_ref())?); } return Ok(vectors); } @@ -86,33 +64,32 @@ pub fn extract_single_vector_from_scalar( location: snafu::Location::new(file!(), line!(), column!()), })?; - let list_array = array - .as_any() - .downcast_ref::() + extract_vectors(&array)? + .into_iter() + .next() .ok_or_else(|| GraphError::ExecutionError { - message: "Expected FixedSizeListArray for vector scalar".to_string(), + message: "Empty vector array".to_string(), location: snafu::Location::new(file!(), line!(), column!()), - })?; + }) +} - if list_array.is_empty() { +/// Numeric lists become Float32 only when consumed as vectors. +fn vector_values(values: &dyn Array) -> Result> { + if !values.data_type().is_numeric() || values.null_count() != 0 { return Err(GraphError::ExecutionError { - message: "Empty vector array".to_string(), + message: "Vector elements must be non-null numbers".to_string(), location: snafu::Location::new(file!(), line!(), column!()), }); } - - let value_array = list_array.value(0); - let float_array = value_array - .as_any() - .downcast_ref::() - .ok_or_else(|| GraphError::ExecutionError { - message: "Expected Float32Array in vector".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), + let floats = + arrow::compute::cast(values, &arrow::datatypes::DataType::Float32).map_err(|e| { + GraphError::ExecutionError { + message: format!("Failed to convert vector elements to Float32: {e}"), + location: snafu::Location::new(file!(), line!(), column!()), + } })?; - - Ok((0..float_array.len()) - .map(|j| float_array.value(j)) - .collect()) + let floats = floats.as_any().downcast_ref::().unwrap(); + Ok(floats.values().to_vec()) } /// Compute L2 (Euclidean) distance between two vectors @@ -234,6 +211,31 @@ mod tests { use super::*; use std::sync::Arc; + #[test] + fn test_integer_list_vectors_and_scalar() { + use arrow::array::{Int64Builder, ListBuilder}; + use datafusion::scalar::ScalarValue; + + let mut builder = ListBuilder::new(Int64Builder::new()); + builder.values().append_slice(&[1, 0, 1]); + builder.append(true); + let array: ArrayRef = Arc::new(builder.finish()); + assert_eq!(extract_vectors(&array).unwrap(), vec![vec![1.0, 0.0, 1.0]]); + let scalar = ScalarValue::try_from_array(&array, 0).unwrap(); + assert_eq!( + extract_single_vector_from_scalar(&scalar).unwrap(), + vec![1.0, 0.0, 1.0] + ); + } + + #[test] + fn test_vector_elements_must_be_non_null_numbers() { + use arrow::array::{Int64Array, StringArray}; + + assert!(vector_values(&Int64Array::from(vec![Some(1), None])).is_err()); + assert!(vector_values(&StringArray::from(vec!["1", "2"])).is_err()); + } + #[test] fn test_l2_distance() { let a = vec![1.0, 0.0, 0.0]; diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index ed89e08..fd66b89 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -18,6 +18,8 @@ use std::collections::HashMap; /// A logical plan operator - describes what operation to perform #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub enum LogicalOperator { + /// One row with no bound variables, for queries without reading clauses. + SingleRow, /// Scan all nodes with a specific label ScanByLabel { variable: String, @@ -385,10 +387,7 @@ impl<'a> LogicalPlanner<'a> { let mut plan = base_plan; if reading_clauses.is_empty() && plan.is_none() { - return Err(GraphError::PlanError { - message: "Query must have at least one MATCH or UNWIND clause".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - }); + plan = Some(LogicalOperator::SingleRow); } for clause in reading_clauses { @@ -1266,6 +1265,10 @@ impl<'a> LogicalPlanner<'a> { #[allow(clippy::only_used_in_recursion)] fn extract_variable_from_plan(&self, plan: &LogicalOperator) -> Result { match plan { + LogicalOperator::SingleRow => Err(GraphError::PlanError { + message: "a singleton row binds no pattern variable".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }), LogicalOperator::ScanByLabel { variable, .. } => Ok(variable.clone()), LogicalOperator::Unwind { alias, .. } => Ok(alias.clone()), LogicalOperator::Expand { diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index 9766749..7f215d6 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -1119,6 +1119,7 @@ fn float32_literal(input: &str) -> IResult<&str, f32> { // Parse vector literal: [0.1, 0.2, 0.3] fn parse_vector_literal(input: &str) -> IResult<&str, ValueExpression> { + let literal = input; let (input, _) = char('[')(input)?; let (input, _) = multispace0(input)?; @@ -1130,6 +1131,17 @@ fn parse_vector_literal(input: &str) -> IResult<&str, ValueExpression> { let (input, _) = multispace0(input)?; let (input, _) = char(']')(input)?; + // Integer lists must retain Int64 precision outside vector operations. + if !literal[..literal.len() - input.len()] + .chars() + .any(|c| matches!(c, '.' | 'e' | 'E')) + { + return Err(nom::Err::Error(nom::error::Error::new( + literal, + nom::error::ErrorKind::Verify, + ))); + } + Ok((input, ValueExpression::VectorLiteral(values))) } @@ -2281,6 +2293,38 @@ mod tests { } } + #[test] + fn test_integer_lists_preserve_precision() { + for literal in ["[1, 2, 3]", "[-19791209300775, 16777217]", "[]"] { + assert!(parse_vector_literal(literal).is_err()); + let (remaining, expression) = value_expression(literal).unwrap(); + assert!(remaining.is_empty()); + let ValueExpression::ScalarFunction { name, args } = expression else { + panic!("expected an ordinary list"); + }; + assert_eq!(name, "__list"); + assert!(args + .iter() + .all(|v| matches!(v, ValueExpression::Literal(PropertyValue::Integer(_))))); + } + let (_, expression) = value_expression("[19791209300775]").unwrap(); + assert_eq!( + expression, + ValueExpression::ScalarFunction { + name: "__list".to_string(), + args: vec![ValueExpression::Literal(PropertyValue::Integer( + 19791209300775 + ))], + } + ); + for literal in ["[1, 2.0]", "[1e-3, 2]", "[-1.0, 2]"] { + assert!(matches!( + value_expression(literal).unwrap().1, + ValueExpression::VectorLiteral(_) + )); + } + } + #[test] fn test_parse_vector_literal() { let result = parse_vector_literal("[0.1, 0.2, 0.3]"); diff --git a/docs/cypher.md b/docs/cypher.md index b7e4674..1339b04 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -177,13 +177,18 @@ read-only subset of openCypher, enough to run all 14 (IC14 without `reduce()`). | `WHERE` | After `MATCH`, `OPTIONAL MATCH` (filters inside the optional part) or `WITH`. | | `WITH [DISTINCT]` | With `ORDER BY`, `LIMIT` and `WHERE`; chain as many as you need. No `SKIP` inside `WITH`. | | `UNWIND list AS x` | | -| `RETURN [DISTINCT]` | With `AS` aliases. | +| `RETURN [DISTINCT]` | With `AS` aliases. Without reading clauses, expressions run on one implicit row. | | `ORDER BY`, `SKIP`, `LIMIT` | `SKIP` and `LIMIT` take integer literals, not parameters. | !!! failure "Not supported" - Writes (`CREATE`, `MERGE`, `SET`, `DELETE`, `REMOVE`), `UNION`, `CALL`, `FOREACH`, and queries with - neither `MATCH` nor `UNWIND`. + Writes (`CREATE`, `MERGE`, `SET`, `DELETE`, `REMOVE`), `UNION`, `CALL`, and `FOREACH`. + +Integer list literals retain signed 64-bit values: `RETURN [1, 2, 3] AS ids` returns integers, and +`UNWIND [1131, 19791209300775] AS id RETURN id` preserves both IDs exactly. Numeric literals with a decimal +point or exponent, such as `[0.1, 1e-3]`, form float32 vectors. Whole-number lists can also be used as query +vectors; conversion to float32 happens inside the vector operation. A configured engine is still required +for standalone `RETURN` queries. ### Patterns diff --git a/docs/vector-search.md b/docs/vector-search.md index 8c5d5d7..21e8a9a 100644 --- a/docs/vector-search.md +++ b/docs/vector-search.md @@ -32,6 +32,10 @@ vector_similarity(a, b, metric) `metric` is a bare keyword (not a string): `cosine`, `l2` or `dot`. One argument must be a property; the other can be a property or a vector literal. +Whole-number query vectors such as `[1, 0]` also work. They remain integer lists until the vector function +converts their elements to float32. Outside vector functions, integer lists preserve signed 64-bit precision. +Vector elements must be non-null numbers. + ```python result = engine.execute(""" MATCH (d:Document) diff --git a/python/python/tests/test_integer_lists.py b/python/python/tests/test_integer_lists.py new file mode 100644 index 0000000..22f9a56 --- /dev/null +++ b/python/python/tests/test_integer_lists.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Integer list values retain precision; only vector operations convert to float32.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + +LARGE_ID = 19791209300775 + + +@pytest.fixture(params=[False, True], ids=["join", "csr"]) +def engine(request): + people = pa.table( + { + "id": [1131, LARGE_ID], + "ids": pa.array([[1131, LARGE_ID], [1, 2]], type=pa.list_(pa.int64())), + "embedding": pa.array( + [[1, 0, 1], [0, 1, 0]], type=pa.list_(pa.float32(), 3) + ), + } + ) + config = GraphConfig.builder().node_table("Person", "id").build() + return CypherEngine(config, {"Person": people}, use_csr=request.param) + + +def test_unwind_large_ids(engine): + result = engine.execute( + f"UNWIND [1131, {LARGE_ID}] AS id " + "MATCH (p:Person {id: id}) RETURN p.id AS id ORDER BY id" + ) + assert result.to_pylist() == [{"id": 1131}, {"id": LARGE_ID}] + assert result.schema.field("id").type == pa.int64() + + +@pytest.mark.parametrize( + ("query", "expected"), + [ + ("RETURN [1, 2, 3] AS l", [1, 2, 3]), + (f"WITH [1131, {LARGE_ID}] AS l RETURN l", [1131, LARGE_ID]), + (f"RETURN [x IN [1131, {LARGE_ID}] | x] AS l", [1131, LARGE_ID]), + (f"RETURN [-{LARGE_ID}, -1, 0] AS l", [-LARGE_ID, -1, 0]), + ], +) +def test_integer_list_values(engine, query, expected): + result = engine.execute(query) + assert result.column(0).to_pylist() == [expected] + assert pa.types.is_int64(result.schema.field(0).type.value_type) + + +def test_integer_list_column_comparison(engine): + result = engine.execute( + f"MATCH (p:Person) WHERE p.ids = [1131, {LARGE_ID}] RETURN p.id AS id" + ) + assert result.to_pylist() == [{"id": 1131}] + + +@pytest.mark.parametrize("metric", ["l2", "cosine", "dot"]) +@pytest.mark.parametrize("function", ["vector_distance", "vector_similarity"]) +def test_integer_query_vector(engine, metric, function): + query = ( + f"MATCH (p:Person) RETURN p.id AS id, " + f"{function}(p.embedding, {{vector}}, {metric}) AS score ORDER BY id" + ) + integer = engine.execute(query.format(vector="[1, 0, 1]")) + decimal = engine.execute(query.format(vector="[1.0, 0.0, 1.0]")) + assert integer.equals(decimal) + + +@pytest.mark.parametrize( + ("query", "expected"), + [ + ("RETURN 1 AS n", [{"n": 1}]), + ("RETURN NULL AS n", [{"n": None}]), + ("RETURN [] AS l", [{"l": []}]), + ("RETURN count(*) AS n", [{"n": 1}]), + ("RETURN 1 AS n LIMIT 0", []), + ("RETURN 1 AS n SKIP 1", []), + ("WITH 2 AS n RETURN n + 1 AS n", [{"n": 3}]), + ("UNWIND [] AS n RETURN n", []), + ], +) +def test_singleton_queries(engine, query, expected): + assert engine.execute(query).to_pylist() == expected + + +def test_singleton_parameters_and_unknown_variables(engine): + assert engine.execute("RETURN $n AS n", {"n": LARGE_ID}).to_pylist() == [ + {"n": LARGE_ID} + ] + with pytest.raises( + ValueError, match="(?i)(unknown|undefined|not defined|not found)" + ): + engine.execute("RETURN missing AS n") From 0d3681b1f0afce8e158bc97c4df74311e7111e2a Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:02:01 -0700 Subject: [PATCH 02/66] feat(cypher): support list and map parameters Preserve integer IDs in nested parameter expressions and round-trip containers through CypherQuery.parameters. Validate keys, integer bounds and nesting; document Arrow list coercion. Part of #20; based on #3. Implemented with assistance from Codex. Runtime and benchmark validation pending. --- crates/lance-graph-python/src/graph.rs | 50 ++++++++- .../lance-graph/src/parameter_substitution.rs | 94 +++++++++++----- docs/api-reference.md | 3 +- docs/cypher.md | 17 ++- docs/vector-search.md | 7 +- .../python/tests/test_container_parameters.py | 105 ++++++++++++++++++ 6 files changed, 232 insertions(+), 44 deletions(-) create mode 100644 python/python/tests/test_container_parameters.py diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 778d688..eadcabd 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -42,7 +42,7 @@ use lance_graph::{ use pyo3::{ exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, prelude::*, - types::PyDict, + types::{PyDict, PyInt, PyList}, IntoPyObject, }; use serde_json::Value as JsonValue; @@ -849,10 +849,25 @@ impl CypherQuery { // Helper functions to convert between Python and JSON values fn python_to_json(value: &Bound<'_, PyAny>) -> PyResult { + python_to_json_value(value, 0) +} + +/// Match serde_json's nesting bound and reject cycles before exhausting the Rust stack. +fn python_to_json_value(value: &Bound<'_, PyAny>, depth: usize) -> PyResult { + if depth >= 128 { + return Err(PyValueError::new_err( + "Parameter nesting exceeds 128 levels (possibly a cyclic container)", + )); + } if value.is_none() { Ok(JsonValue::Null) } else if let Ok(b) = value.extract::() { Ok(JsonValue::Bool(b)) + } else if value.is_instance_of::() { + let i = value.extract::().map_err(|_| { + PyValueError::new_err("Integer parameters must fit signed 64-bit values") + })?; + Ok(JsonValue::Number(i.into())) } else if let Ok(i) = value.extract::() { Ok(JsonValue::Number(i.into())) } else if let Ok(f) = value.extract::() { @@ -862,6 +877,21 @@ fn python_to_json(value: &Bound<'_, PyAny>) -> PyResult { )) } else if let Ok(s) = value.extract::() { Ok(JsonValue::String(s)) + } else if let Ok(values) = value.cast::() { + let values = values + .iter() + .map(|v| python_to_json_value(&v, depth + 1)) + .collect::>>()?; + Ok(JsonValue::Array(values)) + } else if let Ok(values) = value.cast::() { + let mut map = serde_json::Map::new(); + for (key, value) in values.iter() { + let key = key + .extract::() + .map_err(|_| PyValueError::new_err("Map parameter keys must be strings"))?; + map.insert(key, python_to_json_value(&value, depth + 1)?); + } + Ok(JsonValue::Object(map)) } else if value.hasattr("isoformat")? { // datetime and date: ISO 8601 text, which DataFusion casts when comparing with // timestamp and date columns. @@ -890,9 +920,19 @@ fn json_to_python(py: Python, value: &JsonValue) -> PyResult> { } } JsonValue::String(s) => Ok(s.as_str().into_pyobject(py)?.unbind().into()), - JsonValue::Array(_) | JsonValue::Object(_) => { - // For complex types, convert to string representation - Ok(value.to_string().into_pyobject(py)?.unbind().into()) + JsonValue::Array(values) => { + let values = values + .iter() + .map(|v| json_to_python(py, v)) + .collect::>>()?; + Ok(PyList::new(py, values)?.into_any().unbind()) + } + JsonValue::Object(values) => { + let map = PyDict::new(py); + for (key, value) in values { + map.set_item(key, json_to_python(py, value)?)?; + } + Ok(map.into_any().unbind()) } } } @@ -1999,7 +2039,7 @@ impl CypherEngine { /// 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). + /// ``None``, lists, string-keyed maps, and ``datetime``/``date`` values (compared as timestamps and dates). /// ``execute_batches`` and ``explain`` take the same argument. /// /// Raises diff --git a/crates/lance-graph/src/parameter_substitution.rs b/crates/lance-graph/src/parameter_substitution.rs index f66765a..e9bbfa9 100644 --- a/crates/lance-graph/src/parameter_substitution.rs +++ b/crates/lance-graph/src/parameter_substitution.rs @@ -240,29 +240,7 @@ fn substitute_in_value_expression( location: snafu::Location::new(file!(), line!(), column!()), })?; - // Check for array to VectorLiteral conversion - if let serde_json::Value::Array(arr) = param_value { - let mut floats = Vec::new(); - for v in arr { - if let Some(f) = v.as_f64() { - floats.push(f as f32); - } else { - return Err(GraphError::PlanError { - message: format!( - "Parameter ${} is a list but contains non-numeric values. Only float vectors are supported as list parameters currently.", - name - ), - location: snafu::Location::new(file!(), line!(), column!()), - }); - } - } - *expr = ValueExpression::VectorLiteral(floats); - return Ok(()); - } - - // Scalar conversion - let prop_val = json_to_property_value(param_value)?; - *expr = ValueExpression::Literal(prop_val); + *expr = json_to_expression(param_value)?; } ValueExpression::ScalarFunction { args, .. } | ValueExpression::AggregateFunction { args, .. } => { @@ -317,6 +295,26 @@ fn substitute_in_value_expression( Ok(()) } +/// Keep container parameters in the same representation as literal lists and maps. +fn json_to_expression(value: &serde_json::Value) -> Result { + match value { + serde_json::Value::Array(values) => Ok(ValueExpression::ScalarFunction { + name: "__list".to_string(), + args: values + .iter() + .map(json_to_expression) + .collect::>()?, + }), + serde_json::Value::Object(values) => Ok(ValueExpression::Map( + values + .iter() + .map(|(key, value)| Ok((key.clone(), json_to_expression(value)?))) + .collect::>()?, + )), + _ => Ok(ValueExpression::Literal(json_to_property_value(value)?)), + } +} + fn json_to_property_value(value: &serde_json::Value) -> Result { match value { serde_json::Value::Null => Ok(PropertyValue::Null), @@ -324,21 +322,57 @@ fn json_to_property_value(value: &serde_json::Value) -> Result { serde_json::Value::Number(n) => { if let Some(i) = n.as_i64() { Ok(PropertyValue::Integer(i)) + } else if n.is_u64() { + Err(GraphError::PlanError { + message: "Integer parameters must fit signed 64-bit values".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }) } else if let Some(f) = n.as_f64() { Ok(PropertyValue::Float(f)) } else { Err(GraphError::PlanError { - message: format!("Number parameter could not be converted to i64 or f64: {}", n), + message: format!( + "Number parameter could not be converted to i64 or f64: {}", + n + ), location: snafu::Location::new(file!(), line!(), column!()), }) } } serde_json::Value::String(s) => Ok(PropertyValue::String(s.clone())), - serde_json::Value::Array(_) | serde_json::Value::Object(_) => { - Err(GraphError::PlanError { - message: "Complex types (List, Map) are not fully supported as parameters yet (except float vectors).".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - }) - } + serde_json::Value::Array(_) | serde_json::Value::Object(_) => Err(GraphError::PlanError { + message: + "List and map parameters must be used as expressions, not inline property values." + .to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn container_expressions_preserve_ids_and_nesting() { + let list = json_to_expression(&json!([19791209300775_i64, null])).unwrap(); + assert_eq!( + list, + ValueExpression::ScalarFunction { + name: "__list".to_string(), + args: vec![ + ValueExpression::Literal(PropertyValue::Integer(19791209300775)), + ValueExpression::Literal(PropertyValue::Null) + ], + } + ); + let map = json_to_expression(&json!({"ids": [19791209300775_i64]})).unwrap(); + assert!( + matches!(map, ValueExpression::Map(entries) if entries.len() == 1 && entries[0].0 == "ids") + ); + assert!(json_to_expression(&json!([])).is_ok()); + assert!(json_to_expression(&json!({})).is_ok()); + assert!(json_to_expression(&json!(u64::MAX)).is_err()); } } diff --git a/docs/api-reference.md b/docs/api-reference.md index 81ffd03..d61f8b1 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -101,7 +101,8 @@ be opened. `#!python execute(query: str, params: dict | None = None) -> pyarrow.Table` : Runs `query` and returns the whole result. `params` binds `$name` parameters: `int`, `float`, `str`, - `bool`, `None`, `datetime.datetime` and `datetime.date`. + `bool`, `None`, lists, string-keyed nested maps, `datetime.datetime` and `datetime.date`. + Integers must fit signed 64-bit values; nesting beyond 128 levels, cycles and unsupported values raise `ValueError`. `#!python execute_batches(query: str, params: dict | None = None) -> pyarrow.RecordBatchReader` : Streams the result. Planning errors raise immediately; execution errors raise while reading. The reader diff --git a/docs/cypher.md b/docs/cypher.md index 1339b04..b8d8242 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -72,7 +72,7 @@ Output columns are named by their `AS` alias, or by the expression text (`p.name ### Parameters Write `$name` in the query and pass a dict as the second argument. Parameter values can be numbers, -strings, booleans, `None`, and `datetime.datetime` or `datetime.date` values: +strings, booleans, `None`, lists, string-keyed maps, and `datetime.datetime` or `datetime.date` values: ```python import datetime @@ -83,10 +83,19 @@ engine.execute( ).to_pylist() # [{'name': 'David'}] ``` -!!! warning "No list or map parameters" +Lists retain integer IDs, and maps can contain nested lists/maps and nulls: - `$ids` bound to `[1, 2, 3]` raises `ValueError: Unsupported parameter type` on `CypherEngine`. Write the - list literally in the query (`WHERE p.id IN [1, 2, 3]`) for now. +```python +engine.execute( + "MATCH (p:Person) WHERE p.id IN $ids RETURN p.name AS name ORDER BY name", + {"ids": [1, 3]}, +).to_pylist() # [{'name': 'Alice'}, {'name': 'Carol'}] +``` + +Map keys must be strings, integer parameters must fit signed 64-bit values, and nesting is limited to 128 +levels. Cyclic containers and unsupported Python objects raise `ValueError`. Lists follow literal-list +coercion: mixed string/numeric values become strings because Arrow lists have one element type. Containers +are expression values; inline node/relationship property filters still require scalar parameters. Datetimes compare against timestamp columns, and dates against date columns. Parameters are substituted before planning, so they are also allowed in inline property maps: `(p:Person {id: $id})`. diff --git a/docs/vector-search.md b/docs/vector-search.md index 21e8a9a..6a2e36b 100644 --- a/docs/vector-search.md +++ b/docs/vector-search.md @@ -56,10 +56,9 @@ print(result.to_pylist()) Smaller distances and larger similarities mean closer vectors. Comparing vectors of different dimensions gives the worst possible score rather than an error. -!!! warning "Write vector literals into the query" - - `CypherEngine` doesn't accept list parameters yet, so `$query_vector` bound to a list raises - `ValueError`. Format the vector into the query text. +Query vectors may also be passed as list parameters, for example +`engine.execute("MATCH (d:Document) RETURN vector_distance(d.embedding, $v, l2)", {"v": [1, 0]})`. +Elements are converted to float32 only by the vector operation. ## `VectorSearch` diff --git a/python/python/tests/test_container_parameters.py b/python/python/tests/test_container_parameters.py new file mode 100644 index 0000000..14295f8 --- /dev/null +++ b/python/python/tests/test_container_parameters.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Container parameters share literal semantics without losing integer IDs.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, CypherQuery, GraphConfig + +BIG = 19791209300775 +CONFIG = GraphConfig.builder().node_table("Person", "id").build() +TABLES = { + "Person": pa.table( + { + "id": [1, BIG], + "embedding": pa.array([[1, 0], [0, 1]], type=pa.list_(pa.float32(), 2)), + } + ) +} + + +@pytest.fixture(params=[False, True], ids=["join", "csr"]) +def engine(request): + return CypherEngine(CONFIG, TABLES, use_csr=request.param) + + +def test_list_ids_and_streams(engine): + query = "UNWIND $ids AS id MATCH (p:Person {id: id}) RETURN p.id AS id ORDER BY id" + params = {"ids": [1, BIG]} + expected = [{"id": 1}, {"id": BIG}] + assert engine.execute(query, params).to_pylist() == expected + assert engine.execute_batches(query, params).read_all().to_pylist() == expected + assert ( + engine.execute( + "MATCH (p:Person) WHERE p.id IN $ids RETURN p.id AS id ORDER BY id", params + ).to_pylist() + == expected + ) + assert "Projection" in engine.explain(query, params) + + +@pytest.mark.parametrize( + "value", + [ + [], + [None], + [1, BIG], + [[1, BIG], []], + {}, + {"name": "a", "ids": [1, BIG], "missing": None}, + ], +) +def test_container_return_and_parameter_roundtrip(engine, value): + assert engine.execute("RETURN $value AS value", {"value": value}).to_pylist() == [ + {"value": value} + ] + query = ( + CypherQuery("RETURN $value AS value") + .with_config(CONFIG) + .with_parameter("value", value) + ) + assert query.parameters() == {"value": value} + assert query.execute(TABLES).to_pylist() == [{"value": value}] + + +def test_map_field_and_comprehension(engine): + assert engine.execute( + "WITH $value AS m RETURN m.name AS name, m.ids AS ids", + {"value": {"name": "a", "ids": [BIG]}}, + ).to_pylist() == [{"name": "a", "ids": [BIG]}] + assert engine.execute( + "RETURN [x IN $ids WHERE x > 1 | x] AS ids", {"ids": [1, BIG]} + ).to_pylist() == [{"ids": [BIG]}] + + +@pytest.mark.parametrize("vector", [[1, 0], [1.0, 0.0]]) +def test_vector_parameter(engine, vector): + query = ( + "MATCH (p:Person) RETURN p.id AS id, " + "vector_distance(p.embedding, $v, l2) AS d ORDER BY id" + ) + assert engine.execute(query, {"v": vector}).equals( + engine.execute(query.replace("$v", "[1.0, 0.0]")) + ) + + +def test_mixed_list_follows_literal_coercion(engine): + assert engine.execute("RETURN $v AS v", {"v": [1, "a"]}).to_pylist() == [ + {"v": ["1", "a"]} + ] + + +@pytest.mark.parametrize( + "value", [{1: "bad"}, [object()], {"nested": {"bad"}}, [2**64], 2**64] +) +def test_invalid_container_values(engine, value): + with pytest.raises(ValueError): + engine.execute("RETURN $v AS v", {"v": value}) + + +def test_recursive_container_is_rejected(engine): + values = [] + values.append(values) + with pytest.raises(ValueError, match="nesting|cyclic"): + engine.execute("RETURN $v AS v", {"v": values}) From a61cea3775778ec05a28473fdd5a68352bc76030 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 09:54:33 -0700 Subject: [PATCH 03/66] fix(graph): prevent stale concurrent refresh publication Retry refresh when another replacement has been published since its starting snapshot. Add deterministic publication coverage and preserve failed-refresh and active-reader behavior. Resolves #12. Implemented with assistance from Codex; runtime and benchmark validation pending. --- crates/lance-graph-python/src/graph.rs | 123 ++++++++++++++++++++--- docs/api-reference.md | 3 +- docs/lance.md | 3 + python/python/tests/test_engine_lance.py | 23 ++++- 4 files changed, 134 insertions(+), 18 deletions(-) diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index eadcabd..4c21596 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -1961,6 +1961,16 @@ impl CypherEngine { .clone() } + /// Publish only if no other refresh replaced the snapshot used to build this state. + fn try_publish_refresh(&self, expected: &Arc, replacement: EngineState) -> bool { + let mut current = self.state.write().expect("engine state lock poisoned"); + if !Arc::ptr_eq(¤t, expected) { + return false; + } + *current = Arc::new(replacement); + true + } + fn parse( &self, query: &str, @@ -2125,25 +2135,31 @@ impl CypherEngine { /// The replacement is validated in full before it is swapped in; on error nothing /// changes and the previous snapshot stays usable. Queries already running finish /// on the snapshot they started with. PyArrow inputs are unchanged. + /// If another refresh publishes first, retry from its snapshot instead of replacing + /// it with a candidate built from an older state. fn refresh(&self, py: Python) -> PyResult<()> { - let current = self.snapshot(); - let mut sources = current.sources.clone(); - for source in sources.values_mut() { - if let EngineSource::Lance { - uri, - storage_options, - dataset, - } = source - { - let latest = open_lance_dataset(py, uri, None, storage_options)?; - if !same_lance_snapshot(&latest, dataset) { - *dataset = Arc::new(latest); + loop { + let current = self.snapshot(); + let mut sources = current.sources.clone(); + for source in sources.values_mut() { + if let EngineSource::Lance { + uri, + storage_options, + dataset, + } = source + { + let latest = open_lance_dataset(py, uri, None, storage_options)?; + if !same_lance_snapshot(&latest, dataset) { + *dataset = Arc::new(latest); + } } } + let state = + build_engine_state(py, &self.config, sources, self.use_csr, Some(¤t))?; + if self.try_publish_refresh(¤t, state) { + return Ok(()); + } } - let state = build_engine_state(py, &self.config, sources, self.use_csr, Some(¤t))?; - *self.state.write().expect("engine state lock poisoned") = Arc::new(state); - Ok(()) } /// Node id -> row index statistics per node label (experimental, ``use_csr=True``). @@ -2493,6 +2509,83 @@ pub fn register_graph_module(py: Python, parent_module: &Bound<'_, PyModule>) -> mod tests { use super::*; + #[test] + fn test_refresh_rejects_stale_snapshot() { + use arrow_array::Int64Array; + use std::sync::mpsc; + use std::time::Duration; + + fn state(id: i64) -> EngineState { + let batch = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from(vec![id])) as arrow_array::ArrayRef, + )]) + .unwrap(); + let provider = + Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch.clone()]]).unwrap()); + let context = SessionContext::new(); + context.register_table("person", provider.clone()).unwrap(); + EngineState { + sources: HashMap::from([("Person".to_string(), EngineSource::Arrow(batch))]), + catalog: Arc::new( + InMemoryCatalog::new() + .with_node_source("Person", Arc::new(DefaultTableSource::new(provider))), + ), + context, + csr: HashMap::new(), + node_index: HashMap::new(), + key_stats: HashMap::new(), + } + } + + let engine = CypherEngine { + config: RustGraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(), + state: RwLock::new(Arc::new(state(1))), + use_csr: false, + }; + let reader_snapshot = engine.snapshot(); + let (started_tx, started_rx) = mpsc::channel(); + let (published_tx, published_rx) = mpsc::channel(); + std::thread::scope(|threads| { + let engine = &engine; + let older = threads.spawn(move || { + let starting_snapshot = engine.snapshot(); + started_tx.send(()).unwrap(); + published_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + engine.try_publish_refresh(&starting_snapshot, state(2)) + }); + started_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + assert!(engine.try_publish_refresh(&engine.snapshot(), state(3))); + published_tx.send(()).unwrap(); + assert!(!older.join().unwrap()); + }); + + for (snapshot, expected) in [(reader_snapshot, 1), (engine.snapshot(), 3)] { + let query = RustCypherQuery::new("MATCH (p:Person) RETURN p.id") + .unwrap() + .with_config(engine.config.clone()); + let result = RT + .runtime + .block_on(query.execute_with_catalog_and_context( + snapshot.catalog.clone(), + snapshot.context.clone(), + )) + .unwrap(); + assert_eq!( + result + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + expected + ); + } + } + #[test] fn test_split_vector_column_with_alias() { let (alias, column) = split_vector_column("d.embedding"); diff --git a/docs/api-reference.md b/docs/api-reference.md index d61f8b1..95a1261 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -113,7 +113,8 @@ be opened. `#!python refresh() -> None` : Moves every Lance dataset to its latest version and updates the CSR, node indexes and statistics. - All-or-nothing; running queries keep their snapshot. See [Lance datasets](lance.md#refresh). + All-or-nothing; running queries keep their snapshot. If another refresh publishes first, retries from + that snapshot so an older candidate cannot replace it. See [Lance datasets](lance.md#refresh). `#!python dataset_versions() -> dict[str, int]` : The pinned version of each Lance dataset. PyArrow inputs are not listed. diff --git a/docs/lance.md b/docs/lance.md index d18a1cd..23037af 100644 --- a/docs/lance.md +++ b/docs/lance.md @@ -71,6 +71,9 @@ What `refresh()` guarantees: - **Schema changes are picked up.** New columns are queryable after the refresh. - **All or nothing.** The replacement is validated in full before it's swapped in. On error, nothing changes and the previous snapshot keeps working. +- **Concurrent refreshes cannot overwrite a newer published snapshot with an older candidate.** If another + refresh publishes first, the overlapping refresh retries from that snapshot and reopens the latest datasets. + The retry can add work when refreshes overlap; it does not block queries while the replacement is built. - **Running queries are unaffected.** Queries and `execute_batches()` readers that started before the refresh finish on their original snapshot. - **The CSR updates incrementally**, as described in [Keeping the CSR current](csr.md#keeping-the-csr-current). diff --git a/python/python/tests/test_engine_lance.py b/python/python/tests/test_engine_lance.py index 1b815d3..0c81356 100644 --- a/python/python/tests/test_engine_lance.py +++ b/python/python/tests/test_engine_lance.py @@ -153,9 +153,10 @@ def test_refresh_sees_table_replaced_at_same_location(graph): assert count_people(engine) == 3 -def test_failed_refresh_keeps_previous_snapshot(graph): +@pytest.mark.parametrize("use_csr", [False, True]) +def test_failed_refresh_keeps_previous_snapshot(graph, use_csr): path, person, knows = graph - engine = CypherEngine(CONFIG, {"Person": person, "KNOWS": knows}) + engine = CypherEngine(CONFIG, {"Person": person, "KNOWS": knows}, use_csr=use_csr) # Dropping the configured id column must not be accepted. lance.write_dataset( @@ -255,6 +256,24 @@ def test_execute_batches_partial_consumption(graph): assert count_people(engine) == N +@pytest.mark.parametrize("use_csr", [False, True]) +def test_execute_batches_keeps_snapshot_across_refresh(graph, use_csr): + path, person, knows = graph + engine = CypherEngine(CONFIG, {"Person": person, "KNOWS": knows}, use_csr=use_csr) + reader = engine.execute_batches("MATCH (p:Person) RETURN p.id") + first = reader.read_next_batch() + + lance.write_dataset(people(N, N + 5), path / "person.lance", mode="append") + engine.refresh() + + remaining = reader.read_all() + assert first.num_rows + remaining.num_rows == N + assert pc.sum(first.column(0)).as_py() + ( + pc.sum(remaining.column(0)).as_py() or 0 + ) == sum(range(N)) + assert count_people(engine) == N + 5 + + def test_execute_batches_reports_planning_errors(graph): _, person, knows = graph engine = CypherEngine(CONFIG, {"Person": person, "KNOWS": knows}) From 181875fc841745f32984f7d472be1515a488f11e Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 09:56:04 -0700 Subject: [PATCH 04/66] refactor(graph): add read-only config mapping iterators Migrate mapping readers while keeping public fields for source compatibility until a breaking release. Document migration and test normalized keys and original labels. Part of #23. Implemented with assistance from Codex; runtime and benchmark validation pending. --- crates/lance-graph-python/src/graph.rs | 26 ++++++---- crates/lance-graph/README.md | 7 +++ crates/lance-graph/src/config.rs | 51 ++++++++++++++----- .../builder/join_builder.rs | 6 +-- crates/lance-graph/src/query.rs | 19 ++++--- 5 files changed, 75 insertions(+), 34 deletions(-) diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 4c21596..1112d3d 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -348,13 +348,19 @@ impl GraphConfig { /// Get node labels fn node_labels(&self) -> Vec { // Extract from the config's node mappings - self.inner.node_mappings.keys().cloned().collect() + self.inner + .iter_node_mappings() + .map(|(name, _)| name.clone()) + .collect() } /// Get relationship types fn relationship_types(&self) -> Vec { // Extract from the config's relationship mappings - self.inner.relationship_mappings.keys().cloned().collect() + self.inner + .iter_relationship_mappings() + .map(|(name, _)| name.clone()) + .collect() } fn __repr__(&self) -> String { @@ -1477,10 +1483,10 @@ fn validate_sources( } Ok(()) }; - for mapping in config.node_mappings.values() { + for (_, mapping) in config.iter_node_mappings() { check(&mapping.label, &mapping.id_field)?; } - for mapping in config.relationship_mappings.values() { + for (_, mapping) in config.iter_relationship_mappings() { check(&mapping.relationship_type, &mapping.source_id_field)?; check(&mapping.relationship_type, &mapping.target_id_field)?; } @@ -1587,7 +1593,7 @@ fn build_node_index( previous: Option<&EngineState>, ) -> PyResult> { let mut entries = HashMap::new(); - for mapping in config.node_mappings.values() { + for (_, mapping) in config.iter_node_mappings() { let label = &mapping.label; let Some(source) = sources .iter() @@ -1697,7 +1703,7 @@ fn build_key_statistics( let mut to_count = Vec::new(); // Relationship type and endpoint columns, for the fan-out once counts are known. let mut endpoints = None; - for mapping in config.node_mappings.values() { + for (_, mapping) in config.iter_node_mappings() { if !mapping.label.eq_ignore_ascii_case(name) { continue; } @@ -1705,7 +1711,7 @@ fn build_key_statistics( distinct.insert(id.name().clone(), rows()?); } } - for mapping in config.relationship_mappings.values() { + for (_, mapping) in config.iter_relationship_mappings() { if !mapping.relationship_type.eq_ignore_ascii_case(name) { continue; } @@ -1804,7 +1810,7 @@ fn build_csr_entries( previous: Option<&EngineState>, ) -> PyResult> { let mut entries = HashMap::new(); - for mapping in config.relationship_mappings.values() { + for (_, mapping) in config.iter_relationship_mappings() { let rel_type = &mapping.relationship_type; let Some((name, source)) = sources .iter() @@ -2267,8 +2273,8 @@ impl CypherEngine { fn __repr__(&self) -> String { format!( "CypherEngine(nodes={}, relationships={})", - self.config.node_mappings.len(), - self.config.relationship_mappings.len() + self.config.iter_node_mappings().len(), + self.config.iter_relationship_mappings().len() ) } } diff --git a/crates/lance-graph/README.md b/crates/lance-graph/README.md index 5fecb3a..9750021 100644 --- a/crates/lance-graph/README.md +++ b/crates/lance-graph/README.md @@ -84,6 +84,13 @@ let config = GraphConfig::builder() .build()?; ``` +For Rust callers migrating away from direct mapping-field access, use `get_node_mapping(label)` and +`get_relationship_mapping(type_name)` for case-insensitive lookups. `iter_node_mappings()` and +`iter_relationship_mappings()` yield read-only `(normalized_key, mapping)` pairs; iteration order is +unspecified. Build mappings through `GraphConfigBuilder` so keys are normalized. Public fields remain +available for compatibility until a breaking release; replacing direct writes with builder calls avoids +that future migration. The serialized configuration and Python API are unchanged. + ## Executing Cypher Queries - `CypherQuery::new` parses Cypher text into the internal AST. diff --git a/crates/lance-graph/src/config.rs b/crates/lance-graph/src/config.rs index e35eb75..668791a 100644 --- a/crates/lance-graph/src/config.rs +++ b/crates/lance-graph/src/config.rs @@ -21,16 +21,8 @@ use std::collections::HashMap; /// # Future API Changes /// /// These fields may become private in a future major version to enforce invariants. -/// Code should migrate to using accessor methods (`get_node_mapping()`, `get_relationship_mapping()`) -/// rather than direct field access. -/// -/// # TODO: API Safety -/// -/// TODO: Make `node_mappings` and `relationship_mappings` private to prevent external code -/// from bypassing key normalization. This would require: -/// 1. Adding iterator methods for external access (e.g., `iter_node_mappings()`) -/// 2. Making this a breaking API change in next major version -/// 3. Ensuring all internal code uses accessor methods +/// Use the lookup and iterator accessors for reads, and [`GraphConfigBuilder`] for writes. +/// The mapping fields remain public for source compatibility until a breaking release. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GraphConfig { /// Mapping of node labels to their field configurations @@ -126,6 +118,18 @@ impl GraphConfig { self.relationship_mappings.get(&rel_type.to_lowercase()) } + /// Iterate over normalized node labels and their mappings without allowing mutation. + pub fn iter_node_mappings(&self) -> impl ExactSizeIterator { + self.node_mappings.iter() + } + + /// Iterate over normalized relationship types and their mappings without allowing mutation. + pub fn iter_relationship_mappings( + &self, + ) -> impl ExactSizeIterator { + self.relationship_mappings.iter() + } + /// Validate the configuration /// /// Checks for: @@ -430,6 +434,29 @@ impl RelationshipMapping { mod tests { use super::*; + #[test] + fn test_mapping_iterators_keep_normalized_keys_and_original_labels() { + let config = GraphConfig::builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + .unwrap(); + let (key, node) = config.iter_node_mappings().next().unwrap(); + assert_eq!(key, "person"); + assert_eq!(node.label, "Person"); + assert!(std::ptr::eq( + node, + config.get_node_mapping("PERSON").unwrap() + )); + let (key, relationship) = config.iter_relationship_mappings().next().unwrap(); + assert_eq!(key, "knows"); + assert_eq!(relationship.relationship_type, "KNOWS"); + assert!(std::ptr::eq( + relationship, + config.get_relationship_mapping("KnOwS").unwrap() + )); + } + #[test] fn test_graph_config_builder() { let config = GraphConfig::builder() @@ -439,8 +466,8 @@ mod tests { .build() .unwrap(); - assert_eq!(config.node_mappings.len(), 2); - assert_eq!(config.relationship_mappings.len(), 1); + assert_eq!(config.iter_node_mappings().len(), 2); + assert_eq!(config.iter_relationship_mappings().len(), 1); let person_mapping = config.get_node_mapping("Person").unwrap(); assert_eq!(person_mapping.id_field, "person_id"); diff --git a/crates/lance-graph/src/datafusion_planner/builder/join_builder.rs b/crates/lance-graph/src/datafusion_planner/builder/join_builder.rs index 7fc3ac3..97be618 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/join_builder.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/join_builder.rs @@ -242,10 +242,8 @@ impl DataFusionPlanner { .find(|r| r.alias == *var) { // Get the relationship mapping to find src/dst field names - if let Some(rel_map) = self - .config - .relationship_mappings - .get(&rel_instance.rel_type) + if let Some(rel_map) = + self.config.get_relationship_mapping(&rel_instance.rel_type) { // Generate composite join keys for both src_id and dst_id // This ensures we're matching the exact same relationship instance diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 3b43c20..d7f890e 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -480,7 +480,7 @@ impl CypherQuery { let mut catalog = InMemoryCatalog::new(); // Register node sources - for label in config.node_mappings.keys() { + for (label, _) in config.iter_node_mappings() { let table_provider = ctx.table_provider(label) .await @@ -497,7 +497,7 @@ impl CypherQuery { } // Register relationship sources - for rel_type in config.relationship_mappings.keys() { + for (rel_type, _) in config.iter_relationship_mappings() { let table_provider = ctx.table_provider(rel_type) .await @@ -745,12 +745,15 @@ impl CypherQuery { let mut required_tables: HashSet = HashSet::new(); // Use original label/type names (not lowercase keys) for namespace resolution // The namespace needs the original casing to find files on disk - required_tables.extend(config.node_mappings.values().map(|m| m.label.clone())); required_tables.extend( config - .relationship_mappings - .values() - .map(|m| m.relationship_type.clone()), + .iter_node_mappings() + .map(|(_, mapping)| mapping.label.clone()), + ); + required_tables.extend( + config + .iter_relationship_mappings() + .map(|(_, mapping)| mapping.relationship_type.clone()), ); if required_tables.is_empty() { @@ -816,7 +819,7 @@ impl CypherQuery { providers.insert(normalized_table_name.clone(), provider); } - for label in config.node_mappings.keys() { + for (label, _) in config.iter_node_mappings() { let provider = providers .get(label) .ok_or_else(|| GraphError::ConfigError { @@ -831,7 +834,7 @@ impl CypherQuery { catalog = catalog.with_node_source(label, table_source); } - for rel_type in config.relationship_mappings.keys() { + for (rel_type, _) in config.iter_relationship_mappings() { let provider = providers .get(rel_type) .ok_or_else(|| GraphError::ConfigError { From 356f20494dc41fb1d82b32312d568500fe0ae217 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:02:41 -0700 Subject: [PATCH 05/66] feat(cypher): validate known comparison and arithmetic types Use catalog schemas and scoped aliases to diagnose incompatible operands before planning. Preserve numeric and temporal coercions and defer unknown types to DataFusion. Add Rust semantic and Python execution coverage. Part of #23. Implemented with assistance from Codex; runtime tests queued and benchmark validation pending. --- crates/lance-graph/src/query.rs | 2 +- crates/lance-graph/src/semantic.rs | 108 +++++-- crates/lance-graph/src/semantic/types.rs | 377 +++++++++++++++++++++++ docs/api-reference.md | 2 +- docs/cypher.md | 7 +- python/python/tests/test_engine_types.py | 97 ++++++ 6 files changed, 560 insertions(+), 33 deletions(-) create mode 100644 crates/lance-graph/src/semantic/types.rs create mode 100644 python/python/tests/test_engine_types.py diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index d7f890e..5fb7b37 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -883,7 +883,7 @@ impl CypherQuery { let config = self.require_config()?; // Phase 1: Semantic Analysis - let mut analyzer = SemanticAnalyzer::new(config.clone()); + let mut analyzer = SemanticAnalyzer::new(config.clone()).with_catalog(catalog.clone()); let semantic = analyzer.analyze(&self.ast, &self.parameters)?; if !semantic.errors.is_empty() { return Err(GraphError::PlanError { diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index b8cf649..61679ea 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -13,6 +13,10 @@ use crate::case_insensitive::CaseInsensitiveLookup; use crate::config::GraphConfig; use crate::error::{GraphError, Result}; use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +mod types; +use types::{arithmetic_operator, comparison_operator, ExpressionType}; /// Semantic analyzer - validates and enriches the AST pub struct SemanticAnalyzer { @@ -21,6 +25,8 @@ pub struct SemanticAnalyzer { /// Variables whose scope a WITH ended; still type-checked at the end. retired: Vec, current_scope: ScopeType, + catalog: Option>, + expression_types: HashMap, } /// Information about a variable in the query @@ -70,9 +76,20 @@ impl SemanticAnalyzer { variables: HashMap::new(), retired: Vec::new(), current_scope: ScopeType::Match, + catalog: None, + expression_types: HashMap::new(), } } + /// Use the execution snapshot's schemas for checks that configuration alone cannot make. + pub(crate) fn with_catalog( + mut self, + catalog: Arc, + ) -> Self { + self.catalog = Some(catalog); + self + } + /// Analyze a Cypher query AST pub fn analyze( &mut self, @@ -214,6 +231,7 @@ impl SemanticAnalyzer { /// Analyze UNWIND clause and register variables fn analyze_unwind_clause(&mut self, unwind_clause: &UnwindClause) -> Result<()> { self.analyze_value_expression(&unwind_clause.expression)?; + let element_type = self.expression_type(&unwind_clause.expression).element(); // Register the aliased variable (normalize to lowercase for case-insensitive behavior) let var_name = &unwind_clause.alias; @@ -236,8 +254,9 @@ impl SemanticAnalyzer { properties: HashSet::new(), defined_in: self.current_scope.clone(), }; - self.variables.insert(var_name_lower, var_info); + self.variables.insert(var_name_lower.clone(), var_info); } + self.expression_types.insert(var_name_lower, element_type); Ok(()) } @@ -308,6 +327,19 @@ impl SemanticAnalyzer { }; self.variables.insert(var_name_lower, var_info); } + // MATCH binds an entity even when the name previously came from UNWIND. + let name = var_name.to_lowercase(); + let info = &self.variables[&name]; + let mut labels = info.labels.clone(); + if let Some(ExpressionType::Entity(_, inferred)) = self.expression_types.get(&name) { + for label in inferred { + if !labels.contains(label) { + labels.push(label.clone()); + } + } + } + self.expression_types + .insert(name, ExpressionType::Entity(VariableType::Node, labels)); } Ok(()) } @@ -356,9 +388,14 @@ impl SemanticAnalyzer { /// Analyze boolean expression and check variable references fn analyze_boolean_expression(&mut self, expr: &BooleanExpression) -> Result<()> { match expr { - BooleanExpression::Comparison { left, right, .. } => { + BooleanExpression::Comparison { + left, + operator, + right, + } => { self.analyze_value_expression(left)?; self.analyze_value_expression(right)?; + self.validate_binary(left, comparison_operator(operator), right)?; } BooleanExpression::And(left, right) | BooleanExpression::Or(left, right) => { self.analyze_boolean_expression(left)?; @@ -374,12 +411,20 @@ impl SemanticAnalyzer { patterns, where_clause, } => { - for pattern in patterns { - self.analyze_graph_pattern(pattern)?; - } - if let Some(w) = where_clause { - self.analyze_boolean_expression(w)?; - } + let before = self.variables.clone(); + let before_types = self.expression_types.clone(); + let result = (|| -> Result<()> { + for pattern in patterns { + self.analyze_graph_pattern(pattern)?; + } + if let Some(w) = where_clause { + self.analyze_boolean_expression(w)?; + } + Ok(()) + })(); + self.variables = before; + self.expression_types = before_types; + result?; } BooleanExpression::In { expression, list } => { self.analyze_value_expression(expression)?; @@ -623,26 +668,15 @@ impl SemanticAnalyzer { self.analyze_value_expression(otherwise)?; } } - ValueExpression::Arithmetic { left, right, .. } => { + ValueExpression::Arithmetic { + left, + operator, + right, + } => { // Validate arithmetic operands recursively self.analyze_value_expression(left)?; self.analyze_value_expression(right)?; - - // If both sides are literals, ensure they are numeric - let is_numeric_literal = |pv: &PropertyValue| { - matches!(pv, PropertyValue::Integer(_) | PropertyValue::Float(_)) - }; - - if let (ValueExpression::Literal(l1), ValueExpression::Literal(l2)) = - (&**left, &**right) - { - if !(is_numeric_literal(l1) && is_numeric_literal(l2)) { - return Err(GraphError::PlanError { - message: "Arithmetic requires numeric literal operands".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - }); - } - } + self.validate_binary(left, arithmetic_operator(operator), right)?; } ValueExpression::VectorDistance { left, right, .. } => { // Validate vector distance function arguments @@ -705,6 +739,7 @@ impl SemanticAnalyzer { } => { // New variables of the pattern are local to the comprehension. let before = self.variables.clone(); + let before_types = self.expression_types.clone(); let result = (|| -> Result<()> { self.analyze_graph_pattern(&GraphPattern::Path((**pattern).clone()))?; if let Some(b) = filter { @@ -713,6 +748,7 @@ impl SemanticAnalyzer { self.analyze_value_expression(map) })(); self.variables = before; + self.expression_types = before_types; result?; } ValueExpression::ListComprehension { @@ -724,7 +760,11 @@ impl SemanticAnalyzer { self.analyze_value_expression(list)?; // The variable is local to the comprehension. let before = self.variables.clone(); + let before_types = self.expression_types.clone(); + let element_type = self.expression_type(list).element(); self.register_projection_alias(variable); + self.expression_types + .insert(variable.to_lowercase(), element_type); let result = (|| -> Result<()> { if let Some(b) = filter { self.analyze_boolean_expression(b)?; @@ -735,6 +775,7 @@ impl SemanticAnalyzer { Ok(()) })(); self.variables = before; + self.expression_types = before_types; result?; } } @@ -774,7 +815,10 @@ impl SemanticAnalyzer { for item in &return_clause.items { self.analyze_value_expression(&item.expression)?; if let Some(alias) = &item.alias { + let value_type = self.expression_type(&item.expression); self.register_projection_alias(alias); + self.expression_types + .insert(alias.to_lowercase(), value_type); } } Ok(()) @@ -788,6 +832,7 @@ impl SemanticAnalyzer { } // WITH ends the scope of everything it does not carry; `WITH b AS x` renames b. let mut kept = HashMap::new(); + let mut kept_types = HashMap::new(); for item in &with_clause.items { let carried = match &item.expression { ValueExpression::Variable(v) => self.variables.get(&v.to_lowercase()).cloned(), @@ -801,10 +846,15 @@ impl SemanticAnalyzer { if let Some(info) = self.variables.get(&p.variable.to_lowercase()) { kept.insert(p.variable.to_lowercase(), info.clone()); } + if let Some(value_type) = self.expression_types.get(&p.variable.to_lowercase()) + { + kept_types.insert(p.variable.to_lowercase(), value_type.clone()); + } continue; } (None, _) => continue, }; + kept_types.insert(name.to_lowercase(), self.expression_type(&item.expression)); let info = carried.unwrap_or_else(|| VariableInfo { name: name.clone(), variable_type: VariableType::Property, @@ -817,6 +867,7 @@ impl SemanticAnalyzer { self.retired .extend(std::mem::take(&mut self.variables).into_values()); self.variables = kept; + self.expression_types = kept_types; // Validate ORDER BY within WITH if present if let Some(order_by) = &with_clause.order_by { for item in &order_by.items { @@ -865,10 +916,7 @@ impl SemanticAnalyzer { /// Validate types and operations fn validate_types(&self, errors: &mut Vec) { - // TODO: Implement type checking - // - Check that properties exist on nodes/relationships - // - Check that comparison operations are valid for data types - // - Check that arithmetic operations are valid + // Expression operations are checked during traversal, while their scope is available. // Check that properties referenced in patterns exist in schema when property fields are defined for var_info in self.variables.values().chain(&self.retired) { @@ -1815,7 +1863,7 @@ mod tests { assert!(result .errors .iter() - .any(|e| e.contains("Arithmetic requires numeric literal operands"))); + .any(|e| e.contains("Type mismatch") && e.contains("Utf8") && e.contains("Int64"))); } #[test] diff --git a/crates/lance-graph/src/semantic/types.rs b/crates/lance-graph/src/semantic/types.rs new file mode 100644 index 0000000..7331815 --- /dev/null +++ b/crates/lance-graph/src/semantic/types.rs @@ -0,0 +1,377 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Conservative expression checks against the current catalog and lexical scope. + +use super::*; +use arrow_schema::{DataType, TimeUnit}; +use datafusion_expr::binary::BinaryTypeCoercer; +use datafusion_expr::Operator; + +#[derive(Clone, Debug, PartialEq)] +pub(super) enum ExpressionType { + Unknown, + Scalar(DataType), + Entity(VariableType, Vec), + List(Box), + Map(Vec<(String, Self)>), +} + +impl ExpressionType { + pub(super) fn element(self) -> Self { + match self { + Self::List(element) => *element, + Self::Scalar( + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _), + ) => Self::Scalar(field.data_type().clone()), + _ => Self::Unknown, + } + } + + fn known(&self) -> bool { + !matches!(self, Self::Unknown | Self::Scalar(DataType::Null)) + } +} + +impl SemanticAnalyzer { + fn variable_expression_type(&self, name: &str) -> ExpressionType { + if let Some(value_type) = self.expression_types.get(&name.to_lowercase()) { + return value_type.clone(); + } + match self.variables.get_ci(name) { + Some(info) if info.variable_type != VariableType::Property => { + ExpressionType::Entity(info.variable_type.clone(), info.labels.clone()) + } + _ => ExpressionType::Unknown, + } + } + + fn property_expression_type(&self, base: ExpressionType, property: &str) -> ExpressionType { + match base { + ExpressionType::Entity(kind, labels) => { + let Some(catalog) = &self.catalog else { + return ExpressionType::Unknown; + }; + let types = labels + .iter() + .map(|label| { + let source = match kind { + VariableType::Node => catalog.node_source(label), + VariableType::Relationship => catalog.relationship_source(label), + _ => None, + }?; + source + .schema() + .fields() + .iter() + .find(|field| field.name().eq_ignore_ascii_case(property)) + .map(|field| ExpressionType::Scalar(field.data_type().clone())) + }) + .collect::>>(); + types.map(common_type).unwrap_or(ExpressionType::Unknown) + } + ExpressionType::Map(fields) => fields + .into_iter() + .find(|(name, _)| name.eq_ignore_ascii_case(property)) + .map(|(_, value_type)| value_type) + .unwrap_or(ExpressionType::Unknown), + ExpressionType::Scalar(DataType::Struct(fields)) => fields + .iter() + .find(|field| field.name().eq_ignore_ascii_case(property)) + .map(|field| ExpressionType::Scalar(field.data_type().clone())) + .unwrap_or(ExpressionType::Unknown), + // Date component access and function-specific fields remain with DataFusion. + _ => ExpressionType::Unknown, + } + } + + pub(super) fn expression_type(&self, expression: &ValueExpression) -> ExpressionType { + use ExpressionType::{List, Map, Scalar, Unknown}; + match expression { + ValueExpression::Variable(name) => self.variable_expression_type(name), + ValueExpression::Property(property) + | ValueExpression::Literal(PropertyValue::Property(property)) => self + .property_expression_type( + self.variable_expression_type(&property.variable), + &property.property, + ), + ValueExpression::Literal(value) => match value { + PropertyValue::Integer(_) => Scalar(DataType::Int64), + PropertyValue::Float(_) => Scalar(DataType::Float64), + PropertyValue::String(_) => Scalar(DataType::Utf8), + PropertyValue::Boolean(_) => Scalar(DataType::Boolean), + PropertyValue::Null => Scalar(DataType::Null), + PropertyValue::Variable(name) => self.variable_expression_type(name), + _ => Unknown, + }, + ValueExpression::Arithmetic { + left, + operator, + right, + } => match (self.expression_type(left), self.expression_type(right)) { + (Scalar(left), Scalar(right)) => { + BinaryTypeCoercer::new(&left, &arithmetic_operator(operator), &right) + .get_result_type() + .map(Scalar) + .unwrap_or(Unknown) + } + _ => Unknown, + }, + ValueExpression::ScalarFunction { name, args } + | ValueExpression::AggregateFunction { name, args, .. } => { + let first = || { + args.first() + .map(|value| self.expression_type(value)) + .unwrap_or(Unknown) + }; + match name.to_lowercase().as_str() { + "count" | "size" | "length" | "tointeger" | "to_epoch_ms" => { + Scalar(DataType::Int64) + } + "tofloat" => Scalar(DataType::Float64), + "tolower" | "lower" | "toupper" | "upper" | "tostring" => { + Scalar(DataType::Utf8) + } + "date" => Scalar(DataType::Date32), + "datetime" => Scalar(DataType::Timestamp( + TimeUnit::Millisecond, + Some("UTC".into()), + )), + "collect" => List(Box::new(first())), + "head" | "__index" => first().element(), + "range" => List(Box::new(Scalar(DataType::Int64))), + "nodes" => List(Box::new(Unknown)), + "__list" => List(Box::new(common_type( + args.iter() + .map(|value| self.expression_type(value)) + .collect(), + ))), + "coalesce" => common_type( + args.iter() + .map(|value| self.expression_type(value)) + .collect(), + ), + // SUM/AVG output widths depend on DataFusion's accumulator type. + "min" | "max" => first(), + _ => Unknown, + } + } + ValueExpression::Case { whens, otherwise } => common_type( + whens + .iter() + .map(|(_, value)| self.expression_type(value)) + .chain(otherwise.iter().map(|value| self.expression_type(value))) + .collect(), + ), + ValueExpression::Boolean(_) => Scalar(DataType::Boolean), + ValueExpression::Map(fields) => Map(fields + .iter() + .map(|(name, value)| (name.clone(), self.expression_type(value))) + .collect()), + ValueExpression::Field { base, field } => { + self.property_expression_type(self.expression_type(base), field) + } + ValueExpression::VectorLiteral(_) => List(Box::new(Scalar(DataType::Float32))), + ValueExpression::VectorDistance { .. } | ValueExpression::VectorSimilarity { .. } => { + Scalar(DataType::Float64) + } + // Their local scopes have ended; do not infer from similarly named outer variables. + ValueExpression::ListComprehension { .. } + | ValueExpression::PatternComprehension { .. } + | ValueExpression::Parameter(_) => Unknown, + } + } + + pub(super) fn validate_binary( + &self, + left: &ValueExpression, + operator: Operator, + right: &ValueExpression, + ) -> Result<()> { + use ExpressionType::{Entity, Scalar}; + let left_type = self.expression_type(left); + let right_type = self.expression_type(right); + if !left_type.known() || !right_type.known() { + return Ok(()); + } + let comparison = matches!( + operator, + Operator::Eq + | Operator::NotEq + | Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq + ); + let compatible = match (&left_type, &right_type) { + (Scalar(left), Scalar(right)) => { + let string_numeric = (left.is_numeric() + && matches!( + right, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + )) + || (right.is_numeric() + && matches!( + left, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + )); + !string_numeric + && BinaryTypeCoercer::new(left, &operator, right) + .get_result_type() + .is_ok() + } + (Entity(left, _), Entity(right, _)) => comparison && left == right, + // Collection comparison rules remain with the planner; arithmetic on entities does not. + (Entity(..), _) | (_, Entity(..)) => false, + _ => comparison, + }; + if compatible { + Ok(()) + } else { + Err(GraphError::PlanError { + message: format!( + "Type mismatch in '{} {} {}': {} and {}", + display_value(left), + operator, + display_value(right), + display_type(&left_type), + display_type(&right_type) + ), + location: snafu::Location::new(file!(), line!(), column!()), + }) + } + } +} + +fn common_type(types: Vec) -> ExpressionType { + let mut known = types + .into_iter() + .filter(|value| !matches!(value, ExpressionType::Scalar(DataType::Null))); + let Some(first) = known.next() else { + return ExpressionType::Unknown; + }; + if known.all(|value| value == first) { + first + } else { + ExpressionType::Unknown + } +} + +fn display_type(value: &ExpressionType) -> String { + match value { + ExpressionType::Scalar(data_type) => data_type.to_string(), + ExpressionType::Entity(kind, _) => format!("{kind:?}"), + _ => format!("{value:?}"), + } +} + +fn display_value(value: &ValueExpression) -> String { + match value { + ValueExpression::Property(property) => { + format!("{}.{}", property.variable, property.property) + } + ValueExpression::Variable(name) => name.clone(), + ValueExpression::Literal(PropertyValue::Integer(value)) => value.to_string(), + ValueExpression::Literal(PropertyValue::Float(value)) => value.to_string(), + ValueExpression::Literal(PropertyValue::String(value)) => format!("{value:?}"), + ValueExpression::Arithmetic { + left, + operator, + right, + } => format!( + "({} {} {})", + display_value(left), + arithmetic_operator(operator), + display_value(right) + ), + _ => format!("{value:?}"), + } +} + +pub(super) fn arithmetic_operator(operator: &ArithmeticOperator) -> Operator { + match operator { + ArithmeticOperator::Add => Operator::Plus, + ArithmeticOperator::Subtract => Operator::Minus, + ArithmeticOperator::Multiply => Operator::Multiply, + ArithmeticOperator::Divide => Operator::Divide, + ArithmeticOperator::Modulo => Operator::Modulo, + } +} + +pub(super) fn comparison_operator(operator: &ComparisonOperator) -> Operator { + match operator { + ComparisonOperator::Equal => Operator::Eq, + ComparisonOperator::NotEqual => Operator::NotEq, + ComparisonOperator::LessThan => Operator::Lt, + ComparisonOperator::LessThanOrEqual => Operator::LtEq, + ComparisonOperator::GreaterThan => Operator::Gt, + ComparisonOperator::GreaterThanOrEqual => Operator::GtEq, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_schema::{Field, Schema}; + use lance_graph_catalog::{InMemoryCatalog, SimpleTableSource}; + + fn analyze(query: &str) -> SemanticResult { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + Field::new("age", DataType::Int32, true), + Field::new( + "created", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + true, + ), + ])); + let catalog = InMemoryCatalog::new() + .with_node_source("Person", Arc::new(SimpleTableSource::new(schema))); + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + SemanticAnalyzer::new(config) + .with_catalog(Arc::new(catalog)) + .analyze( + &crate::parser::parse_cypher_query(query).unwrap(), + &HashMap::new(), + ) + .unwrap() + } + + #[test] + fn rejects_known_operand_types_in_their_scope() { + for (query, expression, types) in [ + ("MATCH (p:Person) WHERE p.name = 1 RETURN p.id", "p.name = 1", ["Utf8", "Int64"]), + ("MATCH (p:Person) RETURN p + 1", "p + 1", ["Node", "Int64"]), + ("MATCH (p:Person) WITH p.name AS value RETURN value * 2", "value * 2", ["Utf8", "Int64"]), + ("MATCH (p:Person) WITH collect(p.name) AS names UNWIND names AS name RETURN name / 2", "name / 2", ["Utf8", "Int64"]), + ("MATCH (p:Person) RETURN [p IN [p.name] | p + 1]", "p + 1", ["Utf8", "Int64"]), + ] { + let errors = analyze(query).errors.join("\n"); + assert!(errors.contains("Type mismatch") && errors.contains(expression) + && types.iter().all(|name| errors.contains(name)), "{query}: {errors}"); + } + } + + #[test] + fn preserves_coercions_nulls_and_unknown_types() { + for query in [ + "MATCH (p:Person) RETURN p.age + 0.5", + "MATCH (p:Person) WHERE p.created > '2020-01-01T00:00:00Z' RETURN p.id", + "MATCH (p:Person) RETURN p.age + null", + "MATCH (p:Person) RETURN p.unknown + 1", + "MATCH (p:Person) WITH p AS q RETURN q.age + 1", + "MATCH (p:Person) RETURN [p IN [p.age] | p + 1], p.name", + "MATCH (p:Person) RETURN [p IN [p.name, p.age] | p + 1]", + "MATCH (p:Person) WITH collect(p) AS people UNWIND people AS person RETURN person.age + 1", + ] { + let errors = analyze(query).errors; + assert!(errors.is_empty(), "{query}: {errors:?}"); + } + } +} diff --git a/docs/api-reference.md b/docs/api-reference.md index 95a1261..741bcf9 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -32,7 +32,7 @@ Keys match labels case-insensitively. | Exception | When | | --- | --- | -| `ValueError` | Invalid query, unknown label, missing dataset or column, unsupported parameter type, empty `datasets` | +| `ValueError` | Invalid query, incompatible known operand types (with expression and type details), unknown label, missing dataset or column, unsupported parameter type, empty `datasets` | | `NotImplementedError` | A Cypher feature that parses but isn't supported | | `RuntimeError` | Execution or I/O failures | | `TypeError` | A dataset value of an unsupported type, such as a LanceDB `Table` | diff --git a/docs/cypher.md b/docs/cypher.md index b8d8242..b871d80 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -297,6 +297,11 @@ case-insensitive. `match (P:person) return p.NAME` works, and the output column | Exception | Raised when | | --- | --- | -| `ValueError` | The query doesn't parse, refers to an unknown label, or a dataset is missing a configured column. | +| `ValueError` | The query doesn't parse, refers to an unknown label, a dataset is missing a configured column, or a comparison/arithmetic expression has incompatible known types. Type errors name the expression and its operand types. | | `NotImplementedError` | The query parses but uses an unsupported feature (for example a function that isn't implemented, or more than 20 hops). | | `RuntimeError` | Execution fails, for example when a Lance dataset can't be read. | + +Known comparison and arithmetic types are checked against the current dataset schemas before planning. +For example, comparing a string property with a number reports the expression and both types in a +`ValueError`. Mixed numeric types and supported date/timestamp coercions remain valid. Types that cannot +be determined in the current scope are left to the planner, so these checks are not a complete type system. diff --git a/python/python/tests/test_engine_types.py b/python/python/tests/test_engine_types.py new file mode 100644 index 0000000..a629b00 --- /dev/null +++ b/python/python/tests/test_engine_types.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Schema-aware errors retain the expression, operand types and lexical scope.""" + +from datetime import datetime, timezone + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture(params=[False, True]) +def engine(request): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + return CypherEngine( + config, + { + "Person": pa.table( + { + "id": [1, 2], + "name": ["Ada", "Bob"], + "age": pa.array([30, 40], pa.int32()), + "created": [datetime(2020, 1, 1, tzinfo=timezone.utc)] * 2, + } + ), + "KNOWS": pa.table({"src": [1], "dst": [2], "label": ["friend"]}), + }, + use_csr=request.param, + ) + + +@pytest.mark.parametrize( + ("query", "expression", "types"), + [ + ( + "MATCH (p:Person) WHERE p.name = 1 RETURN p.id", + "p.name = 1", + ("Utf8", "Int64"), + ), + ("MATCH (p:Person) RETURN p + 1", "p + 1", ("Node", "Int64")), + ( + "MATCH (p:Person) WITH p.name AS value RETURN value * 2", + "value * 2", + ("Utf8", "Int64"), + ), + ( + "MATCH (p:Person) WITH collect(p.name) AS names UNWIND names AS name " + "RETURN name / 2", + "name / 2", + ("Utf8", "Int64"), + ), + ( + "MATCH (p:Person) RETURN [p IN [p.name] | p + 1]", + "p + 1", + ("Utf8", "Int64"), + ), + ( + "MATCH (p:Person)-[r:KNOWS]->(q:Person) WHERE r.label = 1 RETURN q.id", + "r.label = 1", + ("Utf8", "Int64"), + ), + ], +) +def test_type_errors_name_expression_and_operands(engine, query, expression, types): + with pytest.raises(ValueError, match="Type mismatch") as error: + engine.execute(query) + assert expression in str(error.value) + assert all(name in str(error.value) for name in types) + + +def test_numeric_temporal_and_null_coercions(engine): + result = engine.execute( + "MATCH (p:Person) WHERE p.created >= $since " + "RETURN p.age + 0.5 AS adjusted, p.age + null AS missing ORDER BY adjusted", + {"since": datetime(2020, 1, 1, tzinfo=timezone.utc)}, + ) + assert result.to_pylist() == [ + {"adjusted": 30.5, "missing": None}, + {"adjusted": 40.5, "missing": None}, + ] + + +def test_comprehension_shadow_does_not_change_outer_type(engine): + result = engine.execute( + "MATCH (p:Person) RETURN [p IN [p.age] | p + 1] AS ages, p.name AS name " + "ORDER BY name" + ) + assert result.to_pylist() == [ + {"ages": [31], "name": "Ada"}, + {"ages": [41], "name": "Bob"}, + ] From 5bc730e65c4c0451c3f78dea764c8a801e3f018e Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:16:29 -0700 Subject: [PATCH 06/66] perf(cypher): cache snapshot-local scalar query templates Reuse unoptimized plans in the engine, binding fresh scalar values before optimization. Bound the cache to 128 templates and retain ordinary planning for containers and value-dependent output names. Include result, schema, refresh and concurrency regressions and scoped timing harnesses. Part of #14. Isolated core Rust suite: 344 passed, 1 optional timing diagnostic ignored. Formatting, focused Ruff and strict docs passed. Python runtime, bridge test linking, benchmark and Clippy evidence remain pending. Submitted by David Hughes (@stugorf), with AI Collaboration. --- Cargo.lock | 1 + crates/lance-graph-benches/Cargo.toml | 1 + .../benches/graph_execution.rs | 72 +++- crates/lance-graph-python/src/graph.rs | 224 ++++++++-- .../builder/aggregate_ops.rs | 8 +- .../datafusion_planner/builder/basic_ops.rs | 8 +- .../builder/comprehension_ops.rs | 8 +- .../src/datafusion_planner/expression.rs | 47 ++- .../src/datafusion_planner/join_ops.rs | 4 +- .../lance-graph/src/datafusion_planner/mod.rs | 3 + .../src/datafusion_planner/scan_ops.rs | 9 +- crates/lance-graph/src/query.rs | 387 +++++++++++++++++- docs/api-reference.md | 6 +- docs/lance.md | 15 + python/python/tests/test_plan_cache.py | 83 ++++ 15 files changed, 805 insertions(+), 71 deletions(-) create mode 100644 python/python/tests/test_plan_cache.py diff --git a/Cargo.lock b/Cargo.lock index 0f372ba..7c61376 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4110,6 +4110,7 @@ dependencies = [ "arrow-array", "arrow-schema", "criterion", + "datafusion", "futures", "lance", "lance-graph", diff --git a/crates/lance-graph-benches/Cargo.toml b/crates/lance-graph-benches/Cargo.toml index 6d0b613..3881098 100644 --- a/crates/lance-graph-benches/Cargo.toml +++ b/crates/lance-graph-benches/Cargo.toml @@ -10,6 +10,7 @@ lance-graph = { path = "../lance-graph", version = "0.2.0" } lance = "13.0.0" arrow-array = "58" arrow-schema = "58" +datafusion = { version = "54", default-features = false } criterion = { version = "0.5", features = ["async", "async_tokio", "html_reports"] } futures = "0.3" tempfile = "3" diff --git a/crates/lance-graph-benches/benches/graph_execution.rs b/crates/lance-graph-benches/benches/graph_execution.rs index 2b6831b..333ec6e 100644 --- a/crates/lance-graph-benches/benches/graph_execution.rs +++ b/crates/lance-graph-benches/benches/graph_execution.rs @@ -336,5 +336,75 @@ fn bench_cypher_execution(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_cypher_execution); +// A fixed snapshot isolates template reuse from dataset/catalog construction. +// The result count and checksum are checked outside the timed loops. +fn bench_query_templates(c: &mut Criterion) { + use datafusion::datasource::{DefaultTableSource, MemTable}; + use datafusion::execution::context::SessionContext; + use lance_graph::{GraphSourceCatalog, InMemoryCatalog}; + let rt = tokio::runtime::Runtime::new().unwrap(); + let context = SessionContext::new(); + let batch = make_people_batch(10_000); + let provider = Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch]]).unwrap()); + let catalog: Arc = Arc::new( + InMemoryCatalog::new() + .with_node_source("Person", Arc::new(DefaultTableSource::new(provider))), + ); + let config = GraphConfig::builder() + .node_table("Person", "person_id") + .build() + .unwrap(); + let text = "MATCH (p:Person) WHERE p.age >= $minimum RETURN p.person_id AS id ORDER BY id"; + let fresh = || { + CypherQuery::new(text) + .unwrap() + .with_config(config.clone()) + .with_parameter("minimum", 50) + }; + let mut template = fresh(); + assert!(template.prepare_cached_plan(catalog.clone()).unwrap()); + let expected = rt + .block_on(fresh().execute_with_catalog_and_context(catalog.clone(), context.clone())) + .unwrap(); + let cached = rt + .block_on(template.execute_with_catalog_and_context(catalog.clone(), context.clone())) + .unwrap(); + assert_eq!(cached, expected); + let expected_ids: Vec = (0..10_000).filter(|id| 20 + id % 60 >= 50).collect(); + assert_eq!(cached.num_rows(), expected_ids.len()); + let ids = cached + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + ids.values().iter().map(|id| i64::from(*id)).sum::(), + expected_ids.iter().map(|id| i64::from(*id)).sum::() + ); + let mut group = c.benchmark_group("query_templates"); + group.bench_function("cold_parse_plan_execute", |b| { + b.iter(|| { + black_box( + rt.block_on( + fresh().execute_with_catalog_and_context(catalog.clone(), context.clone()), + ) + .unwrap(), + ); + }) + }); + group.bench_function("cached_bind_plan_execute", |b| { + b.iter(|| { + let query = template.clone().with_parameter("minimum", black_box(50)); + black_box( + rt.block_on( + query.execute_with_catalog_and_context(catalog.clone(), context.clone()), + ) + .unwrap(), + ); + }) + }); + group.finish(); +} + +criterion_group!(benches, bench_cypher_execution, bench_query_templates); criterion_main!(benches); diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 1112d3d..c535661 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -14,8 +14,8 @@ //! Graph query functionality for Lance datasets -use std::collections::HashMap; -use std::sync::{Arc, RwLock}; +use std::collections::{HashMap, VecDeque}; +use std::sync::{Arc, Mutex, RwLock}; use arrow::compute::concat_batches; use arrow::ffi_stream::ArrowArrayStreamReader; @@ -1234,6 +1234,7 @@ enum EngineSource { /// keep the `Arc` they started with, so they finish on their original snapshot. struct EngineState { sources: HashMap, + plans: Mutex>, catalog: Arc, context: SessionContext, /// CSR adjacencies by relationship type; empty unless the engine uses CSR. @@ -1244,6 +1245,33 @@ struct EngineState { key_stats: HashMap, } +/// A small LRU of unoptimized plans. Values are bound after lookup, before optimization. +struct CachedQuery { + text: String, + parameter_types: Vec<(String, u8)>, + query: RustCypherQuery, +} + +// Containers change expression shape; keep them on the uncached path. +fn parameter_types(parameters: &HashMap) -> Option> { + let mut types = parameters + .iter() + .map(|(name, value)| { + let kind = match value { + JsonValue::Null => 0, + JsonValue::Bool(_) => 1, + JsonValue::String(_) => 2, + JsonValue::Number(value) if value.is_i64() => 3, + JsonValue::Number(value) if value.is_f64() => 4, + _ => return None, + }; + Some((name.clone(), kind)) + }) + .collect::>>()?; + types.sort_unstable(); + Some(types) +} + /// Distinct counts for one dataset's node-id and edge-endpoint columns. #[derive(Clone)] struct KeyStatsEntry { @@ -1573,6 +1601,7 @@ fn build_engine_state( } Ok(EngineState { sources, + plans: Mutex::default(), catalog: Arc::new(catalog), context: ctx, csr, @@ -1982,21 +2011,69 @@ impl CypherEngine { query: &str, params: Option<&Bound<'_, PyDict>>, state: &EngineState, + use_cache: bool, ) -> 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)?); + parameters.insert( + key.extract::()?.to_lowercase(), + python_to_json(&value)?, + ); + } + self.plan_query(query, parameters, state, use_cache) + .map_err(graph_error_to_pyerr) + } + + fn plan_query( + &self, + text: &str, + parameters: HashMap, + state: &EngineState, + use_cache: bool, + ) -> Result { + let key = use_cache.then(|| parameter_types(¶meters)).flatten(); + if let Some(key) = &key { + let mut plans = state.plans.lock().unwrap(); + if let Some(index) = plans + .iter() + .position(|entry| entry.text == text && entry.parameter_types == *key) + { + let entry = plans.remove(index).unwrap(); + let query = entry.query.clone().with_parameters(parameters); + plans.push_back(entry); + return Ok(query); + } } - let query = RustCypherQuery::new(query) - .map_err(graph_error_to_pyerr)? + let query = RustCypherQuery::new(text)? .with_parameters(parameters) .with_config(self.config.clone()) .with_fan_out(state.fan_out()); - Ok(if self.use_csr { + let mut query = if self.use_csr { query.with_adjacency(state.adjacency_catalog()) } else { query - }) + }; + if let Some(parameter_types) = key { + // Some operators require literal values during planning. If a template + // cannot be built, ordinary planning preserves their existing behavior. + if query + .prepare_cached_plan(state.catalog.clone()) + .unwrap_or(false) + { + let mut plans = state.plans.lock().unwrap(); + plans + .retain(|entry| entry.text != text || entry.parameter_types != parameter_types); + if plans.len() == 128 { + plans.pop_front(); + } + plans.push_back(CachedQuery { + text: text.to_string(), + parameter_types, + query: query.clone(), + }); + } + } + Ok(query) } } @@ -2072,7 +2149,7 @@ impl CypherEngine { params: Option<&Bound<'_, PyDict>>, ) -> PyResult> { let state = self.snapshot(); - let cypher_query = self.parse(query, params, &state)?; + let cypher_query = self.parse(query, params, &state, true)?; let result_batch = RT .block_on( Some(py), @@ -2100,7 +2177,7 @@ impl CypherEngine { use arrow::pyarrow::IntoPyArrow; let state = self.snapshot(); - let cypher_query = self.parse(query, params, &state)?; + let cypher_query = self.parse(query, params, &state, true)?; let stream = RT .block_on( Some(py), @@ -2126,7 +2203,7 @@ impl CypherEngine { params: Option<&Bound<'_, PyDict>>, ) -> PyResult { let state = self.snapshot(); - let cypher_query = self.parse(query, params, &state)?; + let cypher_query = self.parse(query, params, &state, false)?; RT.block_on( Some(py), cypher_query @@ -2254,7 +2331,7 @@ impl CypherEngine { vector_search: &VectorSearch, ) -> PyResult> { let state = self.snapshot(); - let cypher_query = self.parse(query, None, &state)?; + let cypher_query = self.parse(query, None, &state, false)?; let vs = vector_search.inner.clone(); // Execute query to get candidates, then apply vector reranking @@ -2515,41 +2592,112 @@ pub fn register_graph_module(py: Python, parent_module: &Bound<'_, PyModule>) -> mod tests { use super::*; + fn test_state(id: i64) -> EngineState { + let batch = RecordBatch::try_from_iter([( + "id", + Arc::new(arrow_array::Int64Array::from(vec![id])) as arrow_array::ArrayRef, + )]) + .unwrap(); + let provider = + Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch.clone()]]).unwrap()); + let context = SessionContext::new(); + context.register_table("person", provider.clone()).unwrap(); + EngineState { + sources: HashMap::from([("Person".to_string(), EngineSource::Arrow(batch))]), + plans: Mutex::default(), + catalog: Arc::new( + InMemoryCatalog::new() + .with_node_source("Person", Arc::new(DefaultTableSource::new(provider))), + ), + context, + csr: HashMap::new(), + node_index: HashMap::new(), + key_stats: HashMap::new(), + } + } + + #[test] + fn test_plan_cache_reuses_bounds_and_refreshes() { + let engine = CypherEngine { + config: RustGraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(), + state: RwLock::new(Arc::new(test_state(1))), + use_csr: false, + }; + let snapshot = engine.snapshot(); + let text = "MATCH (p:Person) WHERE p.id = $value RETURN p.id"; + for value in [1, 2, 3] { + engine + .plan_query( + text, + HashMap::from([("value".into(), value.into())]), + &snapshot, + true, + ) + .unwrap(); + assert_eq!(snapshot.plans.lock().unwrap().len(), 1); + // The retained template still has its first parameters: later calls + // really used lookup and rebinding rather than replacing the entry. + assert_eq!( + snapshot.plans.lock().unwrap()[0].query.parameters()["value"], + JsonValue::from(1) + ); + } + for value in [JsonValue::Null, JsonValue::from(1.5)] { + engine + .plan_query( + text, + HashMap::from([("value".into(), value)]), + &snapshot, + true, + ) + .unwrap(); + } + assert_eq!(snapshot.plans.lock().unwrap().len(), 3); + engine + .plan_query( + "MATCH (p:Person) WHERE p.id IN $ids RETURN p.id", + HashMap::from([("ids".into(), serde_json::json!([1, 2]))]), + &snapshot, + true, + ) + .unwrap(); + assert_eq!(snapshot.plans.lock().unwrap().len(), 3); + for value in 0..130 { + engine + .plan_query( + &format!("MATCH (p:Person) RETURN p.id AS id LIMIT {value}"), + HashMap::new(), + &snapshot, + true, + ) + .unwrap(); + } + assert_eq!(snapshot.plans.lock().unwrap().len(), 128); + assert!(snapshot + .plans + .lock() + .unwrap() + .iter() + .all(|entry| entry.text != text)); + assert!(engine.try_publish_refresh(&snapshot, test_state(2))); + assert!(engine.snapshot().plans.lock().unwrap().is_empty()); + assert_eq!(snapshot.plans.lock().unwrap().len(), 128); + } + #[test] fn test_refresh_rejects_stale_snapshot() { - use arrow_array::Int64Array; use std::sync::mpsc; use std::time::Duration; - fn state(id: i64) -> EngineState { - let batch = RecordBatch::try_from_iter([( - "id", - Arc::new(Int64Array::from(vec![id])) as arrow_array::ArrayRef, - )]) - .unwrap(); - let provider = - Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch.clone()]]).unwrap()); - let context = SessionContext::new(); - context.register_table("person", provider.clone()).unwrap(); - EngineState { - sources: HashMap::from([("Person".to_string(), EngineSource::Arrow(batch))]), - catalog: Arc::new( - InMemoryCatalog::new() - .with_node_source("Person", Arc::new(DefaultTableSource::new(provider))), - ), - context, - csr: HashMap::new(), - node_index: HashMap::new(), - key_stats: HashMap::new(), - } - } - let engine = CypherEngine { config: RustGraphConfig::builder() .node_table("Person", "id") .build() .unwrap(), - state: RwLock::new(Arc::new(state(1))), + state: RwLock::new(Arc::new(test_state(1))), use_csr: false, }; let reader_snapshot = engine.snapshot(); @@ -2561,10 +2709,10 @@ mod tests { let starting_snapshot = engine.snapshot(); started_tx.send(()).unwrap(); published_rx.recv_timeout(Duration::from_secs(5)).unwrap(); - engine.try_publish_refresh(&starting_snapshot, state(2)) + engine.try_publish_refresh(&starting_snapshot, test_state(2)) }); started_rx.recv_timeout(Duration::from_secs(5)).unwrap(); - assert!(engine.try_publish_refresh(&engine.snapshot(), state(3))); + assert!(engine.try_publish_refresh(&engine.snapshot(), test_state(3))); published_tx.send(()).unwrap(); assert!(!older.join().unwrap()); }); 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 0cf71fc..35e8d00 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs @@ -96,7 +96,7 @@ impl DataFusionPlanner { 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); + let mut aggregate = self.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 { @@ -105,10 +105,10 @@ impl DataFusionPlanner { } agg_exprs.push(aggregate.alias(column)); } - outer.push(super::super::expression::to_df_value_expr(&rewritten).alias(alias)); + outer.push(self.value_expr(&rewritten).alias(alias)); } else { // Group expressions: use raw expression for grouping, no alias - group_exprs.push(super::super::expression::to_df_value_expr(&p.expression)); + group_exprs.push(self.value_expr(&p.expression)); } } @@ -130,7 +130,7 @@ impl DataFusionPlanner { })); } else if !super::super::expression::contains_aggregate(&p.expression) { // Re-create the expression and apply alias - let expr = super::super::expression::to_df_value_expr(&p.expression); + let expr = self.value_expr(&p.expression); let aliased = if let Some(alias) = &p.alias { expr.alias(alias.to_lowercase()) } else { 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 c10b116..9a6c2d3 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs @@ -87,7 +87,7 @@ impl DataFusionPlanner { ) -> Result { let input_plan = self.build_operator(ctx, input)?; let predicate = super::super::expression::bind_boolean(predicate, input_plan.schema()); - let expr = super::super::expression::to_df_boolean_expr(&predicate); + let expr = self.boolean_expr(&predicate); LogicalPlanBuilder::from(input_plan) .filter(expr) .map_err(|e| self.plan_error("Failed to build filter", e))? @@ -175,7 +175,7 @@ impl DataFusionPlanner { continue; } let bound = super::super::expression::bind_value(&p.expression, &schema); - let expr = super::super::expression::to_df_value_expr(&bound); + let expr = self.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 { @@ -234,7 +234,7 @@ impl DataFusionPlanner { let bound = super::super::expression::bind_value(&item.expression, input_plan.schema()); SortExpr { - expr: super::super::expression::to_df_value_expr(&bound), + expr: self.value_expr(&bound), asc: matches!(item.direction, crate::ast::SortDirection::Ascending), nulls_first: true, } @@ -289,7 +289,7 @@ impl DataFusionPlanner { // Convert expression to DataFusion Expr let bound = super::super::expression::bind_value(expression, input_plan.schema()); - let df_expr = super::super::expression::to_df_value_expr(&bound); + let df_expr = self.value_expr(&bound); // A list of node ids: unnest under a temporary name, then join each id back to the // node table so the alias carries the node's columns like any matched node. let unnested = match label { diff --git a/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs index 43b77d2..8858abc 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs @@ -16,7 +16,7 @@ use datafusion::logical_expr::{ }; use super::super::analysis::PlanningContext; -use super::super::expression::{bind_boolean, bind_value, to_df_boolean_expr, to_df_value_expr}; +use super::super::expression::{bind_boolean, bind_value}; use crate::ast::{BooleanExpression, ValueExpression}; use crate::case_insensitive::qualify_column; use crate::datafusion_planner::DataFusionPlanner; @@ -54,7 +54,7 @@ impl DataFusionPlanner { .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 list = self.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)); @@ -130,9 +130,9 @@ impl DataFusionPlanner { // 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()))); + kept = kept.and(self.boolean_expr(&bind_boolean(keep, elements.schema()))); } - let value = to_df_value_expr(&bind_value(map, elements.schema())); + let value = self.value_expr(&bind_value(map, elements.schema())); let collected = array_agg(value) .filter(kept) .order_by(vec![col(&pos).sort(true, false)]) diff --git a/crates/lance-graph/src/datafusion_planner/expression.rs b/crates/lance-graph/src/datafusion_planner/expression.rs index 8b2bd0a..66fe7c7 100644 --- a/crates/lance-graph/src/datafusion_planner/expression.rs +++ b/crates/lance-graph/src/datafusion_planner/expression.rs @@ -19,6 +19,43 @@ use datafusion_functions_aggregate::min_max::max; use datafusion_functions_aggregate::min_max::min; use datafusion_functions_aggregate::sum::sum; +impl super::DataFusionPlanner { + pub(crate) fn value_expr(&self, value: &ValueExpression) -> Expr { + self.type_parameters(to_df_value_expr(value)) + } + + pub(crate) fn boolean_expr(&self, value: &BooleanExpression) -> Expr { + self.type_parameters(to_df_boolean_expr(value)) + } + + fn type_parameters(&self, expression: Expr) -> Expr { + use datafusion::common::tree_node::{Transformed, TreeNode}; + if self.parameter_types.is_empty() { + return expression; + } + expression + .transform_up(|expression| { + if let Expr::Placeholder(mut parameter) = expression { + parameter.field = self + .parameter_types + .get(parameter.id.trim_start_matches('$')) + .map(|data_type| { + std::sync::Arc::new(arrow_schema::Field::new( + "", + data_type.clone(), + true, + )) + }); + Ok(Transformed::yes(Expr::Placeholder(parameter))) + } else { + Ok(Transformed::no(expression)) + } + }) + .expect("typing placeholders does not fail") + .data + } +} + /// Helper function to create LIKE expressions with consistent settings fn create_like_expr(expression: &ValueExpression, pattern: &str, case_insensitive: bool) -> Expr { Expr::Like(datafusion::logical_expr::Like { @@ -131,10 +168,7 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { } VE::Literal(PV::Variable(v)) => col(v.to_lowercase()), VE::Literal(PV::Parameter(name)) => { - panic!( - "Parameter ${} should have been substituted during semantic analysis", - name - ); + datafusion::logical_expr::placeholder(format!("${}", name.to_lowercase())) } VE::Literal(PV::Property(prop)) => { // Create qualified column name: variable__property (lowercase for case-insensitivity) @@ -411,10 +445,7 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { lit(scalar) } VE::Parameter(name) => { - panic!( - "Parameter ${} should have been substituted during semantic analysis", - name - ); + datafusion::logical_expr::placeholder(format!("${}", name.to_lowercase())) } VE::Boolean(b) => to_df_boolean_expr(b), VE::Map(entries) => datafusion::functions::core::expr_fn::named_struct( diff --git a/crates/lance-graph/src/datafusion_planner/join_ops.rs b/crates/lance-graph/src/datafusion_planner/join_ops.rs index da9ed1c..61282b2 100644 --- a/crates/lance-graph/src/datafusion_planner/join_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/join_ops.rs @@ -59,9 +59,7 @@ impl DataFusionPlanner { // Apply target property filters (e.g., (b {age: 30})) for (k, v) in target_properties.iter() { - let lit_expr = super::expression::to_df_value_expr( - &crate::ast::ValueExpression::Literal(v.clone()), - ); + let lit_expr = self.value_expr(&crate::ast::ValueExpression::Literal(v.clone())); let filter_expr = Expr::BinaryExpr(BinaryExpr { left: Box::new(col(k.to_lowercase())), op: Operator::Eq, diff --git a/crates/lance-graph/src/datafusion_planner/mod.rs b/crates/lance-graph/src/datafusion_planner/mod.rs index 6c8d8ab..8c58f29 100644 --- a/crates/lance-graph/src/datafusion_planner/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/mod.rs @@ -47,6 +47,7 @@ pub struct DataFusionPlanner { pub(crate) config: GraphConfig, pub(crate) catalog: Option>, pub(crate) adjacency: Option, + pub(crate) parameter_types: std::collections::HashMap, } impl DataFusionPlanner { @@ -55,6 +56,7 @@ impl DataFusionPlanner { config, catalog: None, adjacency: None, + parameter_types: Default::default(), } } @@ -63,6 +65,7 @@ impl DataFusionPlanner { config, catalog: Some(catalog), adjacency: None, + parameter_types: Default::default(), } } diff --git a/crates/lance-graph/src/datafusion_planner/scan_ops.rs b/crates/lance-graph/src/datafusion_planner/scan_ops.rs index 33434f8..5a5466f 100644 --- a/crates/lance-graph/src/datafusion_planner/scan_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/scan_ops.rs @@ -44,9 +44,8 @@ impl DataFusionPlanner { let filter_exprs: Vec = properties .iter() .map(|(k, v)| { - let lit_expr = super::expression::to_df_value_expr( - &crate::ast::ValueExpression::Literal(v.clone()), - ); + let lit_expr = + self.value_expr(&crate::ast::ValueExpression::Literal(v.clone())); Expr::BinaryExpr(BinaryExpr { left: Box::new(col(k)), op: Operator::Eq, @@ -140,9 +139,7 @@ impl DataFusionPlanner { // Apply relationship property filters (e.g., -[r {since: 2020}]->) for (k, v) in relationship_properties.iter() { - let lit_expr = super::expression::to_df_value_expr( - &crate::ast::ValueExpression::Literal(v.clone()), - ); + let lit_expr = self.value_expr(&crate::ast::ValueExpression::Literal(v.clone())); let filter_expr = Expr::BinaryExpr(BinaryExpr { left: Box::new(col(k)), op: Operator::Eq, diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 5fb7b37..7a4373b 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -129,7 +129,24 @@ pub struct CypherQuery { adjacency: Option, /// Optional mean fan-out per relationship type and direction; see [`FanOut`] fan_out: Option, + /// An unoptimized template belonging to the caller's catalog snapshot. + cached_plan: Option>, } +struct CachedPlan { + graph: crate::logical_plan::LogicalOperator, + datafusion: datafusion::logical_expr::LogicalPlan, + catalog: Arc, + parameter_types: HashMap, +} + +impl std::fmt::Debug for CachedPlan { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CachedPlan") + .field("parameter_types", &self.parameter_types) + .finish_non_exhaustive() + } +} + impl CypherQuery { /// Create a new Cypher query from a query string pub fn new(query: &str) -> Result { @@ -142,11 +159,13 @@ impl CypherQuery { parameters: HashMap::new(), adjacency: None, fan_out: None, + cached_plan: None, }) } /// Set the graph configuration for this query pub fn with_config(mut self, config: GraphConfig) -> Self { + self.cached_plan = None; self.config = Some(config); self } @@ -157,6 +176,7 @@ impl CypherQuery { /// [`crate::csr_expand::CsrQueryPlanner`]. Each adjacency must be built from the /// same relationship data the context serves, or results will be wrong. pub fn with_adjacency(mut self, adjacency: crate::csr_expand::AdjacencyCatalog) -> Self { + self.cached_plan = None; self.adjacency = Some(adjacency); self } @@ -164,6 +184,7 @@ impl CypherQuery { /// Mean fan-out per relationship type and direction, so a pattern that continues /// earlier rows can start from its cheaper bound end (see [`crate::logical_plan::FanOut`]). pub fn with_fan_out(mut self, fan_out: crate::logical_plan::FanOut) -> Self { + self.cached_plan = None; self.fan_out = Some(fan_out); self } @@ -207,6 +228,85 @@ impl CypherQuery { &self.parameters } + /// Build a snapshot-local scalar-parameter template for the Python engine's cache. + /// Binding must precede optimization so pushed CSR filters never retain old values. + #[doc(hidden)] + pub fn prepare_cached_plan( + &mut self, + catalog: Arc, + ) -> Result { + if self.scalar_parameters().is_none() || self.has_parameter_named_columns() { + return Ok(false); + } + let (graph, datafusion) = self.build_logical_plans(catalog.clone(), true)?; + self.cached_plan = Some(Arc::new(CachedPlan { + graph, + datafusion, + catalog, + parameter_types: self + .scalar_parameters() + .unwrap() + .into_iter() + .map(|(name, value)| (name, value.data_type())) + .collect(), + })); + Ok(true) + } + + // Unaliased expressions use their bound values as column names. Reusing a + // placeholder-named projection would change the observable Arrow schema. + fn has_parameter_named_columns(&self) -> bool { + self.ast + .return_clause + .items + .iter() + .chain(self.ast.with_clause.iter().flat_map(|clause| &clause.items)) + .chain( + self.ast + .further_with_stages + .iter() + .flat_map(|stage| &stage.with_clause.items), + ) + .any(|item| { + if item.alias.is_some() { + return false; + } + let mut found = false; + crate::ast::transform_value(&mut item.expression.clone(), &mut |value| { + found |= matches!( + value, + crate::ast::ValueExpression::Parameter(_) + | crate::ast::ValueExpression::Literal( + crate::ast::PropertyValue::Parameter(_) + ) + ); + }); + found + }) + } + + fn scalar_parameters(&self) -> Option> { + use datafusion::common::ScalarValue; + self.parameters + .iter() + .map(|(name, value)| { + let value = match value { + serde_json::Value::Null => ScalarValue::Null, + serde_json::Value::Bool(value) => ScalarValue::Boolean(Some(*value)), + serde_json::Value::String(value) => ScalarValue::Utf8(Some(value.clone())), + serde_json::Value::Number(value) if value.is_i64() => { + ScalarValue::Int64(value.as_i64()) + } + serde_json::Value::Number(value) if value.is_f64() => { + ScalarValue::Float64(value.as_f64()) + } + _ => return None, + }; + Some((name.clone(), value)) + }) + .collect() + } + /// Get the required config, returning an error if not set fn require_config(&self) -> Result<&GraphConfig> { self.config.as_ref().ok_or_else(|| GraphError::ConfigError { @@ -876,6 +976,37 @@ impl CypherQuery { ) -> Result<( crate::logical_plan::LogicalOperator, datafusion::logical_expr::LogicalPlan, + )> { + if let Some(template) = &self.cached_plan { + if let Some(parameters) = self.scalar_parameters() { + let same_types = parameters.len() == template.parameter_types.len() + && parameters.iter().all(|(name, value)| { + template.parameter_types.get(name) == Some(&value.data_type()) + }); + if !Arc::ptr_eq(&catalog, &template.catalog) || !same_types { + return self.build_logical_plans(catalog, false); + } + let plan = template + .datafusion + .clone() + .with_param_values(parameters) + .map_err(|error| GraphError::PlanError { + message: format!("Failed to bind query parameters: {error}"), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + return Ok((template.graph.clone(), plan)); + } + } + self.build_logical_plans(catalog, false) + } + + fn build_logical_plans( + &self, + catalog: Arc, + parameterized: bool, + ) -> Result<( + crate::logical_plan::LogicalOperator, + datafusion::logical_expr::LogicalPlan, )> { use crate::datafusion_planner::{DataFusionPlanner, GraphPhysicalPlanner}; use crate::semantic::SemanticAnalyzer; @@ -894,12 +1025,24 @@ impl CypherQuery { // Phase 2: Graph Logical Plan let mut logical_planner = LogicalPlanner::new(config).with_fan_out(self.fan_out.as_ref()); - // Use the transformed AST (with parameters substituted) - let logical_plan = logical_planner.plan(&semantic.ast)?; + let ast = if parameterized { + &self.ast + } else { + &semantic.ast + }; + let logical_plan = logical_planner.plan(ast)?; // Phase 3: DataFusion Logical Plan let mut df_planner = DataFusionPlanner::with_catalog(config.clone(), catalog); df_planner.adjacency = self.adjacency.clone(); + if parameterized { + df_planner.parameter_types = self + .scalar_parameters() + .unwrap_or_default() + .into_iter() + .map(|(name, value)| (name, value.data_type())) + .collect(); + } let df_logical_plan = df_planner.plan(&logical_plan)?; Ok((logical_plan, df_logical_plan)) @@ -1325,6 +1468,7 @@ impl CypherQueryBuilder { parameters: self.parameters, adjacency: None, fan_out: None, + cached_plan: None, }; Ok(query) @@ -1360,6 +1504,245 @@ mod tests { assert!(query.parameters().contains_key("maxage")); } + #[tokio::test] + async fn cached_templates_rebind_scalars_and_preserve_schema() { + use arrow_array::{ArrayRef, Int64Array, StringArray}; + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + let datasets = HashMap::from([( + "Person".into(), + RecordBatch::try_from_iter([ + ("id", Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef), + ( + "name", + Arc::new(StringArray::from(vec!["Ada", "Bob", "Cam"])) as ArrayRef, + ), + ]) + .unwrap(), + )]); + let cases = [ + ( + "MATCH (p:Person) WHERE p.id >= $value RETURN p.id ORDER BY p.id", + vec![serde_json::json!(1), serde_json::json!(3)], + ), + ( + "MATCH (p:Person) RETURN p.id + $value AS adjusted ORDER BY adjusted", + vec![serde_json::json!(1.5), serde_json::json!(8.5)], + ), + ( + "MATCH (p:Person) WHERE p.name = $value RETURN p.id", + vec![serde_json::json!("Ada"), serde_json::json!("Cam")], + ), + ( + "MATCH (p:Person) RETURN $value AS value", + vec![serde_json::json!(true), serde_json::json!(false)], + ), + ( + "MATCH (p:Person) RETURN $value AS value", + vec![serde_json::Value::Null], + ), + ( + "MATCH (p:Person) RETURN $value", + vec![serde_json::json!(1), serde_json::json!(2)], + ), + ( + "MATCH (p:Person) RETURN p.id + $value", + vec![serde_json::json!(1), serde_json::json!(2)], + ), + ( + "MATCH (p:Person) RETURN p.id AS id, [x IN range(1, $value) | x + p.id] AS values ORDER BY id", + vec![serde_json::json!(1), serde_json::json!(3)], + ), + ]; + for (text, values) in cases { + let mut template = CypherQuery::new(text) + .unwrap() + .with_config(config.clone()) + .with_parameter("value", values[0].clone()); + let (catalog, context) = template + .build_catalog_and_context_from_datasets(datasets.clone()) + .await + .unwrap(); + let catalog: Arc = Arc::new(catalog); + let cached = template.prepare_cached_plan(catalog.clone()).unwrap(); + assert_eq!(cached, !template.has_parameter_named_columns(), "{text}"); + for value in values { + let cached = template.clone().with_parameter("value", value.clone()); + let fresh = CypherQuery::new(text) + .unwrap() + .with_config(config.clone()) + .with_parameter("value", value); + let actual = cached + .execute_with_catalog_and_context(catalog.clone(), context.clone()) + .await + .unwrap(); + let expected = fresh + .execute_with_catalog_and_context(catalog.clone(), context.clone()) + .await + .unwrap(); + assert_eq!(actual, expected, "{text}"); + } + } + } + + #[tokio::test] + async fn cached_templates_respect_catalog_and_parameter_types() { + use arrow_array::{ArrayRef, Int64Array}; + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + let data = |id| { + HashMap::from([( + "Person".into(), + RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from(vec![id])) as ArrayRef, + )]) + .unwrap(), + )]) + }; + let mut query = CypherQuery::new("MATCH (p:Person) WHERE p.id = $value RETURN p.id") + .unwrap() + .with_config(config.clone()) + .with_parameter("value", 1); + let (catalog, _) = query + .build_catalog_and_context_from_datasets(data(1)) + .await + .unwrap(); + assert!(query.prepare_cached_plan(Arc::new(catalog)).unwrap()); + // A different catalog must never serve rows from the prepared snapshot. + assert_eq!(query.execute(data(2), None).await.unwrap().num_rows(), 0); + assert!(query + .clone() + .with_parameter("value", "wrong") + .execute(data(1), None) + .await + .unwrap_err() + .to_string() + .contains("Type mismatch")); + assert!(query.with_config(config).cached_plan.is_none()); + let mut containers = CypherQuery::new("MATCH (p:Person) WHERE p.id IN $ids RETURN p.id") + .unwrap() + .with_parameter("ids", serde_json::json!([1, 2])); + assert!(!containers + .prepare_cached_plan(Arc::new(crate::InMemoryCatalog::new())) + .unwrap()); + } + + // Opt-in diagnostic: reports stage means for a fixed local scan, with no timing + // assertions. Physical planning calls the query planner directly, avoiding the + // extra optimize() performed by SessionState::create_physical_plan(). + #[tokio::test] + #[ignore = "local stage timing; run with --ignored --nocapture"] + async fn profile_query_planning_stages() { + use crate::datafusion_planner::{DataFusionPlanner, GraphPhysicalPlanner}; + use crate::semantic::SemanticAnalyzer; + use arrow_array::{ArrayRef, Int64Array}; + use std::time::Instant; + let text = "MATCH (p:Person) WHERE p.id >= $minimum RETURN p.id AS id ORDER BY id"; + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + let parameters = HashMap::from([("minimum".into(), serde_json::json!(10))]); + let mut template = CypherQuery::new(text) + .unwrap() + .with_config(config.clone()) + .with_parameters(parameters.clone()); + let data = HashMap::from([( + "Person".into(), + RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from_iter_values(0..1000)) as ArrayRef, + )]) + .unwrap(), + )]); + let (catalog, context) = template + .build_catalog_and_context_from_datasets(data) + .await + .unwrap(); + let catalog: Arc = Arc::new(catalog); + template.prepare_cached_plan(catalog.clone()).unwrap(); + let state = context.state(); + let mut totals = [std::time::Duration::ZERO; 8]; + for _ in 0..100 { + let start = Instant::now(); + let ast = parse_cypher_query(text).unwrap(); + totals[0] += start.elapsed(); + let start = Instant::now(); + let semantic = SemanticAnalyzer::new(config.clone()) + .with_catalog(catalog.clone()) + .analyze(&ast, ¶meters) + .unwrap(); + assert!(semantic.errors.is_empty()); + totals[1] += start.elapsed(); + let start = Instant::now(); + let graph = LogicalPlanner::new(&config).plan(&semantic.ast).unwrap(); + totals[2] += start.elapsed(); + let start = Instant::now(); + let logical = DataFusionPlanner::with_catalog(config.clone(), catalog.clone()) + .plan(&graph) + .unwrap(); + totals[3] += start.elapsed(); + let start = Instant::now(); + let optimized = state.optimize(&logical).unwrap(); + totals[4] += start.elapsed(); + let start = Instant::now(); + let physical = state + .query_planner() + .create_physical_plan(&optimized, &state) + .await + .unwrap(); + totals[5] += start.elapsed(); + let start = Instant::now(); + let batches = datafusion::physical_plan::collect(physical, context.task_ctx()) + .await + .unwrap(); + totals[6] += start.elapsed(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 990 + ); + let checksum: i64 = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + }) + .sum(); + assert_eq!(checksum, (10 + 999) * 990 / 2); + let start = Instant::now(); + let _ = template.create_logical_plans(catalog.clone()).unwrap(); + totals[7] += start.elapsed(); + } + for (name, duration) in [ + "parse", + "semantic", + "graph_logical", + "datafusion_logical", + "optimize", + "physical", + "execute", + "cached_bind", + ] + .into_iter() + .zip(totals) + { + println!( + "{name}: {:.3} us mean (100 iterations, 1000 input rows, 990 result rows)", + duration.as_secs_f64() * 1_000_000.0 / 100.0 + ); + } + } + #[test] fn test_query_builder() { let config = GraphConfig::builder() diff --git a/docs/api-reference.md b/docs/api-reference.md index 741bcf9..e67ad00 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -87,7 +87,11 @@ CypherEngine(config: GraphConfig, datasets: dict, use_csr: bool = False) ``` Registers `datasets` once and reuses the catalog, DataFusion session and (with `use_csr=True`) CSR index -across queries. Safe to share between threads. +across queries. Safe to share between threads. `execute()` and `execute_batches()` reuse up to 128 +unoptimized query templates per snapshot, keyed by exact text and scalar parameter names/types. Values are +bound before optimization on every call. Containers, value-dependent column names and unsupported template +shapes use ordinary planning. `explain()` always plans directly; successful `refresh()` clears the cache. +See [Repeated queries](lance.md#repeated-queries) for the scope and limits. | Parameter | Description | | --- | --- | diff --git a/docs/lance.md b/docs/lance.md index 23037af..9ff45cc 100644 --- a/docs/lance.md +++ b/docs/lance.md @@ -82,6 +82,21 @@ What `refresh()` guarantees: A refresh reopens each dataset by its URI with the same `storage_options` it was opened with, so object-store credentials carry over. +## Repeated queries + +`execute()` and `execute_batches()` keep up to 128 query templates per engine snapshot. Reuse the exact +query text with `$name` parameters to reuse parsing, semantic analysis, graph logical planning and +DataFusion logical planning. Every execution binds its own values before optimization and physical planning; +rows and results are never cached. Query text and parameter names/types identify each template, with nulls +separate from non-null values. Least recently used templates are evicted at the limit. + +Lists and maps use ordinary planning because their shape can change the plan. Unaliased parameter expressions +also use ordinary planning to preserve their value-dependent column names; use `RETURN $value AS value` +when you want a stable name. Queries whose operators cannot produce a reusable template fall back to ordinary +planning. `explain()` plans the current values directly. A successful `refresh()` starts an empty cache for the +new snapshot, including a refresh whose datasets did not change; active readers retain their existing plans. +These rules apply with and without CSR. They do not promise a speedup for a particular workload. + ## Streaming large results `execute()` returns one table, held in memory. For big outputs, stream: diff --git a/python/python/tests/test_plan_cache.py b/python/python/tests/test_plan_cache.py new file mode 100644 index 0000000..276bb57 --- /dev/null +++ b/python/python/tests/test_plan_cache.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Repeated execution must bind current values before filters reach CSR scans.""" + +from concurrent.futures import ThreadPoolExecutor + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture(params=[False, True]) +def engine(request): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + return CypherEngine( + config, + { + "Person": pa.table({"id": [1, 2, 3], "age": [20, 30, 40]}), + "KNOWS": pa.table({"src": [1, 1, 2], "dst": [2, 3, 3]}), + }, + use_csr=request.param, + ) + + +def test_cached_target_filters_bind_each_value(engine): + query = ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) " + "WHERE a.id = $start AND b.age >= $minimum RETURN b.id AS id ORDER BY id" + ) + for minimum, ids in [(20, [2, 3]), (35, [3]), (45, []), (20, [2, 3])]: + params = {"start": 1, "minimum": minimum} + expected = [{"id": value} for value in ids] + assert engine.execute(query, params).to_pylist() == expected + assert engine.execute_batches(query, params).read_all().to_pylist() == expected + + +def test_cached_parameter_types_and_container_fallback(engine): + query = "MATCH (p:Person) WHERE p.id = 1 RETURN $value AS value" + for value in [1, 2, None, 1.5, "Ada", "Bob", True, False, [1, 2], [3], {"x": 1}]: + assert engine.execute(query, {"value": value}).to_pylist() == [{"value": value}] + membership = "MATCH (p:Person) WHERE p.id IN $ids RETURN p.id AS id ORDER BY id" + for ids in [[1, 2], [3], []]: + assert engine.execute(membership, {"ids": ids}).to_pylist() == [ + {"id": value} for value in ids + ] + + +def test_unaliased_parameters_keep_bound_column_names(engine): + query = "MATCH (p:Person) WHERE p.id = 1 RETURN $value" + first = engine.execute(query, {"value": 1}) + second = engine.execute(query, {"value": 2}) + assert first.column_names == ["1"] + assert second.column_names == ["2"] + assert first.column(0).to_pylist() == [1] + assert second.column(0).to_pylist() == [2] + + +def test_cached_queries_keep_concurrent_parameters_separate(engine): + query = "MATCH (p:Person) WHERE p.id = $id RETURN p.age AS age" + engine.execute(query, {"id": 1}) + + def execute(person_id): + return engine.execute(query, {"id": person_id}).to_pylist() + + ids = [1, 2, 3] * 8 + with ThreadPoolExecutor(max_workers=4) as pool: + results = list(pool.map(execute, ids)) + assert results == [[{"age": 10 * (person_id + 1)}] for person_id in ids] + + +def test_cached_query_still_rejects_missing_or_wrong_types(engine): + query = "MATCH (p:Person) WHERE p.id = $id RETURN p.age AS age" + engine.execute(query, {"id": 1}) + with pytest.raises(ValueError, match="[Pp]arameter"): + engine.execute(query) + with pytest.raises(ValueError, match="Type mismatch"): + engine.execute(query, {"id": "wrong"}) From 7adc028157ad120de5c05fbdd76a3ebab4785539 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:05:04 -0700 Subject: [PATCH 07/66] feat(cypher): bind parameters in skip and limit Resolve literal-or-parameter row counts before logical planning, including chained WITH limits. Reject invalid counts without truncation or wraparound. Part of #20. Implemented with assistance from Codex; coordinated runtime and benchmark checks pending. --- crates/lance-graph/src/ast.rs | 27 ++++++- crates/lance-graph/src/logical_plan.rs | 12 +-- .../lance-graph/src/parameter_substitution.rs | 41 +++++++++++ crates/lance-graph/src/parser.rs | 53 ++++++++------ crates/lance-graph/src/query.rs | 4 +- docs/cypher.md | 6 +- .../tests/test_pagination_parameters.py | 73 +++++++++++++++++++ 7 files changed, 183 insertions(+), 33 deletions(-) create mode 100644 python/python/tests/test_pagination_parameters.py diff --git a/crates/lance-graph/src/ast.rs b/crates/lance-graph/src/ast.rs index 52d40c7..7d4e183 100644 --- a/crates/lance-graph/src/ast.rs +++ b/crates/lance-graph/src/ast.rs @@ -10,6 +10,27 @@ use serde::{Deserialize, Serialize}; use std::collections::HashMap; +/// A pagination count, resolved to a nonnegative literal before logical planning. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum RowCount { + Literal(u64), + Parameter(String), +} + +impl RowCount { + /// Return a resolved count; unresolved parameters are a planning error. + pub fn value(&self) -> crate::error::Result { + match self { + Self::Literal(value) => Ok(*value), + Self::Parameter(name) => Err(crate::error::GraphError::PlanError { + message: format!("Unresolved pagination parameter: ${name}"), + location: snafu::Location::new(file!(), line!(), column!()), + }), + } + } +} + /// A complete Cypher query #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct CypherQuery { @@ -30,11 +51,11 @@ pub struct CypherQuery { /// RETURN clause pub return_clause: ReturnClause, /// LIMIT clause (optional) - pub limit: Option, + pub limit: Option, /// ORDER BY clause (optional) pub order_by: Option, /// SKIP/OFFSET clause (optional) - pub skip: Option, + pub skip: Option, } impl CypherQuery { @@ -476,7 +497,7 @@ pub struct WithClause { /// Optional ORDER BY within WITH pub order_by: Option, /// Optional LIMIT within WITH - pub limit: Option, + pub limit: Option, /// `WITH ... WHERE`: filters the WITH output, before any following MATCH #[serde(default)] pub where_clause: Option, diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index fd66b89..4359fc0 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -360,18 +360,18 @@ impl<'a> LogicalPlanner<'a> { } // Apply SKIP/OFFSET if present - if let Some(skip) = query.skip { + if let Some(skip) = &query.skip { plan = LogicalOperator::Offset { input: Box::new(plan), - offset: skip, + offset: skip.value()?, }; } // Apply LIMIT if present - if let Some(limit) = query.limit { + if let Some(limit) = &query.limit { plan = LogicalOperator::Limit { input: Box::new(plan), - count: limit, + count: limit.value()?, }; } @@ -1510,10 +1510,10 @@ impl<'a> LogicalPlanner<'a> { } // Apply LIMIT within WITH if present - if let Some(limit) = with_clause.limit { + if let Some(limit) = &with_clause.limit { plan = LogicalOperator::Limit { input: Box::new(plan), - count: limit, + count: limit.value()?, }; } diff --git a/crates/lance-graph/src/parameter_substitution.rs b/crates/lance-graph/src/parameter_substitution.rs index e9bbfa9..d0a0df4 100644 --- a/crates/lance-graph/src/parameter_substitution.rs +++ b/crates/lance-graph/src/parameter_substitution.rs @@ -7,6 +7,9 @@ pub fn substitute_parameters( query: &mut CypherQuery, parameters: &HashMap, ) -> Result<()> { + substitute_row_count(&mut query.skip, parameters)?; + substitute_row_count(&mut query.limit, parameters)?; + // Substitute in READING clauses for reading_clause in &mut query.reading_clauses { substitute_in_reading_clause(reading_clause, parameters)?; @@ -53,6 +56,22 @@ pub fn substitute_parameters( Ok(()) } +fn substitute_row_count( + count: &mut Option, + parameters: &HashMap, +) -> Result<()> { + if let Some(RowCount::Parameter(name)) = count { + let value = parameters.get(&name.to_lowercase()).and_then(serde_json::Value::as_u64) + .filter(|v| usize::try_from(*v).is_ok()) + .ok_or_else(|| GraphError::PlanError { + message: format!("Pagination parameter ${name} must be a supplied nonnegative integer fitting the platform row count"), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + *count = Some(RowCount::Literal(value)); + } + Ok(()) +} + fn substitute_in_reading_clause( clause: &mut ReadingClause, parameters: &HashMap, @@ -143,6 +162,7 @@ fn substitute_in_with_clause( with_clause: &mut WithClause, parameters: &HashMap, ) -> Result<()> { + substitute_row_count(&mut with_clause.limit, parameters)?; for item in &mut with_clause.items { substitute_in_value_expression(&mut item.expression, parameters)?; } @@ -354,6 +374,27 @@ mod tests { use super::*; use serde_json::json; + #[test] + fn pagination_requires_nonnegative_integral_parameters() { + let mut query = crate::parser::parse_cypher_query("RETURN 1 LIMIT $n").unwrap(); + substitute_parameters(&mut query, &HashMap::from([("n".to_string(), json!(0))])).unwrap(); + assert_eq!(query.limit, Some(RowCount::Literal(0))); + for value in [ + json!(null), + json!(-1), + json!(1.5), + json!(true), + json!("1"), + json!([]), + ] { + let mut query = crate::parser::parse_cypher_query("RETURN 1 LIMIT $n").unwrap(); + assert!( + substitute_parameters(&mut query, &HashMap::from([("n".to_string(), value)])) + .is_err() + ); + } + } + #[test] fn container_expressions_preserve_ids_and_nesting() { let list = json_to_expression(&json!([19791209300775_i64, null])).unwrap(); diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index 7f215d6..2f6035b 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -1309,35 +1309,37 @@ fn order_by_item(input: &str) -> IResult<&str, OrderByItem> { } // Parse a LIMIT clause -fn limit_clause(input: &str) -> IResult<&str, u64> { +fn limit_clause(input: &str) -> IResult<&str, RowCount> { let (input, _) = multispace0(input)?; let (input, _) = tag_no_case("LIMIT")(input)?; let (input, _) = multispace1(input)?; - let (input, limit) = integer_literal(input)?; - // Negative counts are an error, not a wrap-around to a huge u64. - let limit = u64::try_from(limit).map_err(|_| { - nom::Err::Failure(nom::error::Error::new(input, nom::error::ErrorKind::Verify)) - })?; - Ok((input, limit)) + row_count(input) } // Parse a SKIP clause -fn skip_clause(input: &str) -> IResult<&str, u64> { +fn skip_clause(input: &str) -> IResult<&str, RowCount> { let (input, _) = multispace0(input)?; let (input, _) = tag_no_case("SKIP")(input)?; let (input, _) = multispace1(input)?; - let (input, skip) = integer_literal(input)?; - let skip = u64::try_from(skip).map_err(|_| { + row_count(input) +} + +fn row_count(input: &str) -> IResult<&str, RowCount> { + if let Ok((remaining, name)) = parameter(input) { + return Ok((remaining, RowCount::Parameter(name))); + } + let (remaining, value) = integer_literal(input)?; + let value = u64::try_from(value).map_err(|_| { nom::Err::Failure(nom::error::Error::new(input, nom::error::ErrorKind::Verify)) })?; - Ok((input, skip)) + Ok((remaining, RowCount::Literal(value))) } // Parse pagination clauses (SKIP and LIMIT) -fn pagination_clauses(input: &str) -> IResult<&str, (Option, Option)> { +fn pagination_clauses(input: &str) -> IResult<&str, (Option, Option)> { let (mut remaining, _) = multispace0(input)?; - let mut skip: Option = None; - let mut limit: Option = None; + let mut skip: Option = None; + let mut limit: Option = None; loop { let before = remaining; @@ -1702,7 +1704,7 @@ mod tests { let query = "MATCH (n:Person) RETURN n.name LIMIT 10"; let result = parse_cypher_query(query).unwrap(); - assert_eq!(result.limit, Some(10)); + assert_eq!(result.limit, Some(RowCount::Literal(10))); } #[test] @@ -1710,7 +1712,7 @@ mod tests { let query = "MATCH (n:Person) RETURN n.name SKIP 5"; let result = parse_cypher_query(query).unwrap(); - assert_eq!(result.skip, Some(5)); + assert_eq!(result.skip, Some(RowCount::Literal(5))); assert_eq!(result.limit, None); } @@ -1719,8 +1721,8 @@ mod tests { let query = "MATCH (n:Person) RETURN n.name SKIP 5 LIMIT 10"; let result = parse_cypher_query(query).unwrap(); - assert_eq!(result.skip, Some(5)); - assert_eq!(result.limit, Some(10)); + assert_eq!(result.skip, Some(RowCount::Literal(5))); + assert_eq!(result.limit, Some(RowCount::Literal(10))); } #[test] @@ -1728,7 +1730,7 @@ mod tests { let query = "MATCH (n:Person) RETURN n.name ORDER BY n.age SKIP 5"; let result = parse_cypher_query(query).unwrap(); - assert_eq!(result.skip, Some(5)); + assert_eq!(result.skip, Some(RowCount::Literal(5))); assert!(result.order_by.is_some()); } @@ -1737,8 +1739,8 @@ mod tests { let query = "MATCH (n:Person) RETURN n.name ORDER BY n.age SKIP 5 LIMIT 10"; let result = parse_cypher_query(query).unwrap(); - assert_eq!(result.skip, Some(5)); - assert_eq!(result.limit, Some(10)); + assert_eq!(result.skip, Some(RowCount::Literal(5))); + assert_eq!(result.limit, Some(RowCount::Literal(10))); assert!(result.order_by.is_some()); } @@ -2293,6 +2295,15 @@ mod tests { } } + #[test] + fn pagination_accepts_parameters_without_losing_names() { + let query = parse_cypher_query("RETURN 1 AS n SKIP $offset LIMIT $count").unwrap(); + assert_eq!(query.skip, Some(RowCount::Parameter("offset".to_string()))); + assert_eq!(query.limit, Some(RowCount::Parameter("count".to_string()))); + assert!(parse_cypher_query("RETURN 1 LIMIT -1").is_err()); + assert!(parse_cypher_query("RETURN 1 LIMIT 1.5").is_err()); + } + #[test] fn test_integer_lists_preserve_precision() { for literal in ["[1, 2, 3]", "[-19791209300775, 16777217]", "[]"] { diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 7a4373b..c91c406 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -1454,8 +1454,8 @@ impl CypherQueryBuilder { items: self.order_by_items, }) }, - limit: self.limit, - skip: self.skip, + limit: self.limit.map(crate::ast::RowCount::Literal), + skip: self.skip.map(crate::ast::RowCount::Literal), }; // Generate query text from AST (simplified) diff --git a/docs/cypher.md b/docs/cypher.md index b871d80..a299aaa 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -92,6 +92,10 @@ engine.execute( ).to_pylist() # [{'name': 'Alice'}, {'name': 'Carol'}] ``` +Pagination also accepts parameters, for example `RETURN 1 AS n LIMIT $take` with `{"take": 1}`. +`WITH ... LIMIT $take` uses the same rules. Missing, null, boolean, fractional, negative and overflowing +counts raise `ValueError`; arbitrary expressions and `WITH SKIP` remain unsupported. + Map keys must be strings, integer parameters must fit signed 64-bit values, and nesting is limited to 128 levels. Cyclic containers and unsupported Python objects raise `ValueError`. Lists follow literal-list coercion: mixed string/numeric values become strings because Arrow lists have one element type. Containers @@ -187,7 +191,7 @@ read-only subset of openCypher, enough to run all 14 (IC14 without `reduce()`). | `WITH [DISTINCT]` | With `ORDER BY`, `LIMIT` and `WHERE`; chain as many as you need. No `SKIP` inside `WITH`. | | `UNWIND list AS x` | | | `RETURN [DISTINCT]` | With `AS` aliases. Without reading clauses, expressions run on one implicit row. | -| `ORDER BY`, `SKIP`, `LIMIT` | `SKIP` and `LIMIT` take integer literals, not parameters. | +| `ORDER BY`, `SKIP`, `LIMIT` | `SKIP` and `LIMIT` take nonnegative integer literals or parameters. | !!! failure "Not supported" diff --git a/python/python/tests/test_pagination_parameters.py b/python/python/tests/test_pagination_parameters.py new file mode 100644 index 0000000..4465e72 --- /dev/null +++ b/python/python/tests/test_pagination_parameters.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Pagination parameters resolve before planning and never wrap or truncate.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, CypherQuery, GraphConfig + +CONFIG = GraphConfig.builder().node_table("Person", "id").build() +TABLES = {"Person": pa.table({"id": [1, 2, 3, 4]})} + + +@pytest.fixture(params=[False, True]) +def engine(request): + return CypherEngine(CONFIG, TABLES, use_csr=request.param) + + +@pytest.mark.parametrize( + ("query", "params", "expected"), + [ + ( + "MATCH (p:Person) RETURN p.id AS id ORDER BY id SKIP $skip LIMIT $take", + {"skip": 1, "take": 2}, + [{"id": 2}, {"id": 3}], + ), + ( + "MATCH (p:Person) RETURN p.id AS id ORDER BY id LIMIT $take SKIP $skip", + {"skip": 1, "take": 2}, + [{"id": 2}, {"id": 3}], + ), + ( + "MATCH (p:Person) WITH p ORDER BY p.id LIMIT $take RETURN p.id AS id", + {"take": 2}, + [{"id": 1}, {"id": 2}], + ), + ( + "MATCH (p:Person) WITH p ORDER BY p.id LIMIT $first " + "WITH p LIMIT $second RETURN p.id AS id", + {"first": 3, "second": 1}, + [{"id": 1}], + ), + ("RETURN 1 AS n LIMIT $take", {"take": 0}, []), + ("RETURN 1 AS n SKIP $skip", {"skip": 1}, []), + ], +) +def test_parameterized_pagination(engine, query, params, expected): + assert engine.execute(query, params).to_pylist() == expected + assert engine.execute_batches(query, params).read_all().to_pylist() == expected + one_shot = CypherQuery(query).with_config(CONFIG) + for name, value in params.items(): + one_shot = one_shot.with_parameter(name, value) + assert one_shot.execute(TABLES).to_pylist() == expected + assert engine.explain(query, params) + + +@pytest.mark.parametrize("clause", ["LIMIT", "SKIP"]) +@pytest.mark.parametrize( + "params", + [ + {}, + {"n": None}, + {"n": -1}, + {"n": 1.5}, + {"n": True}, + {"n": "2"}, + {"n": []}, + {"n": 2**64}, + ], +) +def test_invalid_pagination(engine, clause, params): + with pytest.raises(ValueError): + engine.execute(f"RETURN 1 AS n {clause} $n", params) From 3c3e2afe8e2b17e6171eee8b1efcdd4d74a27c6c Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:16:57 -0700 Subject: [PATCH 08/66] fix(cypher): bypass templates for parameterized row counts SKIP and LIMIT parameters determine graph plan shape, including chained WITH limits. Keep these queries on ordinary planning and document the boundary. Add a regression for every row-count scope. Part of #14, dependent on #20. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/src/query.rs | 42 ++++++++++++++++++++++++++++++++- docs/api-reference.md | 2 +- docs/lance.md | 3 ++- 3 files changed, 44 insertions(+), 3 deletions(-) diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index c91c406..62c97dc 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -235,7 +235,28 @@ impl CypherQuery { &mut self, catalog: Arc, ) -> Result { - if self.scalar_parameters().is_none() || self.has_parameter_named_columns() { + let parameter_counts = self + .ast + .limit + .iter() + .chain(self.ast.skip.iter()) + .chain( + self.ast + .with_clause + .iter() + .filter_map(|clause| clause.limit.as_ref()), + ) + .chain( + self.ast + .further_with_stages + .iter() + .filter_map(|stage| stage.with_clause.limit.as_ref()), + ) + .any(|count| matches!(count, crate::ast::RowCount::Parameter(_))); + if parameter_counts + || self.scalar_parameters().is_none() + || self.has_parameter_named_columns() + { return Ok(false); } let (graph, datafusion) = self.build_logical_plans(catalog.clone(), true)?; @@ -1632,6 +1653,25 @@ mod tests { .unwrap()); } + #[test] + fn parameterized_row_counts_bypass_templates() { + for text in [ + "MATCH (p:Person) RETURN p.id LIMIT $count", + "MATCH (p:Person) RETURN p.id SKIP $count", + "MATCH (p:Person) WITH p LIMIT $count RETURN p.id", + "MATCH (p:Person) WITH p WITH p LIMIT $count RETURN p.id", + ] { + let mut query = CypherQuery::new(text).unwrap().with_parameter("count", 1); + assert!( + !query + .prepare_cached_plan(Arc::new(crate::InMemoryCatalog::new())) + .unwrap(), + "{text}" + ); + assert!(query.cached_plan.is_none()); + } + } + // Opt-in diagnostic: reports stage means for a fixed local scan, with no timing // assertions. Physical planning calls the query planner directly, avoiding the // extra optimize() performed by SessionState::create_physical_plan(). diff --git a/docs/api-reference.md b/docs/api-reference.md index e67ad00..d4150be 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -89,7 +89,7 @@ CypherEngine(config: GraphConfig, datasets: dict, use_csr: bool = False) Registers `datasets` once and reuses the catalog, DataFusion session and (with `use_csr=True`) CSR index across queries. Safe to share between threads. `execute()` and `execute_batches()` reuse up to 128 unoptimized query templates per snapshot, keyed by exact text and scalar parameter names/types. Values are -bound before optimization on every call. Containers, value-dependent column names and unsupported template +bound before optimization on every call. Containers, parameterized `SKIP`/`LIMIT`, value-dependent column names and unsupported template shapes use ordinary planning. `explain()` always plans directly; successful `refresh()` clears the cache. See [Repeated queries](lance.md#repeated-queries) for the scope and limits. diff --git a/docs/lance.md b/docs/lance.md index 9ff45cc..9aaec50 100644 --- a/docs/lance.md +++ b/docs/lance.md @@ -90,7 +90,8 @@ DataFusion logical planning. Every execution binds its own values before optimiz rows and results are never cached. Query text and parameter names/types identify each template, with nulls separate from non-null values. Least recently used templates are evicted at the limit. -Lists and maps use ordinary planning because their shape can change the plan. Unaliased parameter expressions +Lists, maps and parameterized `SKIP`/`LIMIT` counts use ordinary planning because their shape or value can +change the plan. Unaliased parameter expressions also use ordinary planning to preserve their value-dependent column names; use `RETURN $value AS value` when you want a stable name. Queries whose operators cannot produce a reusable template fall back to ordinary planning. `explain()` plans the current values directly. A successful `refresh()` starts an empty cache for the From dd15d82e6944416e32b44c415e3b7c00f663a1e8 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:18:54 -0700 Subject: [PATCH 09/66] test(cypher): document and exercise template reuse boundaries Release the GIL during template lookup/build, cover repeated parameterized limits, and document exact local stage measurements and benchmark proof. Coordinated Python, benchmark and SF1 validation remain pending. Part of #14. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph-benches/README.md | 24 ++++++++++++++++++++++++ crates/lance-graph-python/src/graph.rs | 11 ++++++----- python/python/tests/test_plan_cache.py | 18 ++++++++++++++++++ 3 files changed, 48 insertions(+), 5 deletions(-) diff --git a/crates/lance-graph-benches/README.md b/crates/lance-graph-benches/README.md index ce63a57..fc00241 100644 --- a/crates/lance-graph-benches/README.md +++ b/crates/lance-graph-benches/README.md @@ -26,3 +26,27 @@ cargo bench -p lance-graph-benches --bench graph_execution This crate is not published to crates.io and is excluded from releases. It exists solely for performance testing during development. + + +## Query template reuse + +Measure fresh parsing/planning/execution against reuse of an unoptimized scalar-parameter template on the +same 10,000-row Arrow snapshot. Both paths optimize and build a physical plan for each execution. The harness +checks identical Arrow results, row count and ID checksum before timing; catalog construction is excluded. + +```bash +cargo bench -p lance-graph-benches --bench graph_execution -- query_templates --warm-up-time 1 +``` + +For a stage breakdown on a separate fixed 1,000-row scan, run the optional core diagnostic: + +```bash +cargo test --release -p lance-graph profile_query_planning_stages -- --ignored --nocapture +``` + +It reports elapsed means for 100 iterations of parsing, semantic analysis, graph logical planning, +DataFusion logical planning, optimization, physical planning, execution and cached binding. Physical timing +calls DataFusion's query planner directly so optimization is counted only once. It verifies 990 result rows +and their ID checksum. These local measurements do not replace the two SF1 suites required for performance +submissions; record runtime, hardware, revision and build profile with any reported numbers. The engine's +Python cache also includes dictionary conversion and lookup, which this Rust microbenchmark does not time. diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index c535661..296f6b2 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -2008,6 +2008,7 @@ impl CypherEngine { fn parse( &self, + py: Python, query: &str, params: Option<&Bound<'_, PyDict>>, state: &EngineState, @@ -2020,7 +2021,7 @@ impl CypherEngine { python_to_json(&value)?, ); } - self.plan_query(query, parameters, state, use_cache) + py.detach(|| self.plan_query(query, parameters, state, use_cache)) .map_err(graph_error_to_pyerr) } @@ -2149,7 +2150,7 @@ impl CypherEngine { params: Option<&Bound<'_, PyDict>>, ) -> PyResult> { let state = self.snapshot(); - let cypher_query = self.parse(query, params, &state, true)?; + let cypher_query = self.parse(py, query, params, &state, true)?; let result_batch = RT .block_on( Some(py), @@ -2177,7 +2178,7 @@ impl CypherEngine { use arrow::pyarrow::IntoPyArrow; let state = self.snapshot(); - let cypher_query = self.parse(query, params, &state, true)?; + let cypher_query = self.parse(py, query, params, &state, true)?; let stream = RT .block_on( Some(py), @@ -2203,7 +2204,7 @@ impl CypherEngine { params: Option<&Bound<'_, PyDict>>, ) -> PyResult { let state = self.snapshot(); - let cypher_query = self.parse(query, params, &state, false)?; + let cypher_query = self.parse(py, query, params, &state, false)?; RT.block_on( Some(py), cypher_query @@ -2331,7 +2332,7 @@ impl CypherEngine { vector_search: &VectorSearch, ) -> PyResult> { let state = self.snapshot(); - let cypher_query = self.parse(query, None, &state, false)?; + let cypher_query = self.parse(py, query, None, &state, false)?; let vs = vector_search.inner.clone(); // Execute query to get candidates, then apply vector reranking diff --git a/python/python/tests/test_plan_cache.py b/python/python/tests/test_plan_cache.py index 276bb57..a0b2ea5 100644 --- a/python/python/tests/test_plan_cache.py +++ b/python/python/tests/test_plan_cache.py @@ -81,3 +81,21 @@ def test_cached_query_still_rejects_missing_or_wrong_types(engine): engine.execute(query) with pytest.raises(ValueError, match="Type mismatch"): engine.execute(query, {"id": "wrong"}) + + +@pytest.mark.parametrize( + "query", + [ + "MATCH (p:Person) RETURN p.id AS id ORDER BY id LIMIT $count", + "MATCH (p:Person) WITH p ORDER BY p.id LIMIT $count RETURN p.id AS id", + "MATCH (p:Person) WITH p WITH p ORDER BY p.id LIMIT $count RETURN p.id AS id", + ], +) +def test_parameterized_limits_replan_each_count(engine, query): + for count in [1, 3, 0, 2]: + expected = [{"id": value} for value in range(1, count + 1)] + assert engine.execute(query, {"count": count}).to_pylist() == expected + assert ( + engine.execute_batches(query, {"count": count}).read_all().to_pylist() + == expected + ) From 9f82994e18c2f1d1a7023213e74ec5e3c026cfa2 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:19:06 -0700 Subject: [PATCH 10/66] test(graph): restore refresh assertion array import Keep the existing refresh test array downcast in scope after sharing its real snapshot fixture with cache tests. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph-python/src/graph.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 296f6b2..4690d76 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -2690,6 +2690,7 @@ mod tests { #[test] fn test_refresh_rejects_stale_snapshot() { + use arrow_array::Int64Array; use std::sync::mpsc; use std::time::Duration; From aa368ebf091b459ef2ab40b0f248470355ece461 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:02:18 -0700 Subject: [PATCH 11/66] fix(query): resolve shadowing output aliases in sort Resolve all RETURN expressions before applying scalar aliases to the ORDER BY binding scope. Preserve whole-node projection behavior. Add focused planner and CSR/join Python regressions for issue #21. Document alias scope in the Cypher guide. Validation: Rust formatting, Python lint, strict docs build, spelling, and whitespace checks pass. Published 0.2.0 reproduction: 14 failures, 4 existing passes. Current-source runtime tests remain pending the coordinated build. Implemented by David Hughes with assistance from Codex. --- crates/lance-graph/src/logical_plan.rs | 49 ++++++++++++++ docs/cypher.md | 8 +-- python/python/tests/test_alias_shadowing.py | 74 +++++++++++++++++++++ 3 files changed, 127 insertions(+), 4 deletions(-) create mode 100644 python/python/tests/test_alias_shadowing.py diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index 4359fc0..a85a642 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -1386,6 +1386,17 @@ impl<'a> LogicalPlanner<'a> { } } + // RETURN aliases shadow input entities in ORDER BY, after every item is bound. + // Whole-node aliases expand to `alias.property`, so keep their entity binding. + self.variables.retain(|variable, _| { + !projections.iter().any(|projection| { + projection + .alias + .as_ref() + .is_some_and(|alias| alias.eq_ignore_ascii_case(variable)) + }) + }); + let mut plan = LogicalOperator::Project { input: Box::new(input), projections, @@ -1530,6 +1541,44 @@ mod tests { use super::*; use crate::{parser::parse_cypher_query, NodeMapping}; + #[test] + fn return_alias_shadows_node_only_after_projection() { + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + let query = + parse_cypher_query("MATCH (p:Person) RETURN p.name AS P, p.id AS id ORDER BY p") + .unwrap(); + let LogicalOperator::Sort { input, sort_items } = + LogicalPlanner::new(&config).plan(&query).unwrap() + else { + panic!("expected sort"); + }; + assert_eq!( + sort_items[0].expression, + ValueExpression::Variable("p".into()) + ); + let LogicalOperator::Project { projections, .. } = *input else { + panic!("expected projection"); + }; + assert_eq!( + projections[1].expression, + ValueExpression::Property(PropertyRef { + variable: "p".into(), + property: "id".into(), + }) + ); + + let query = parse_cypher_query("MATCH (p:Person) RETURN p AS p ORDER BY p").unwrap(); + let LogicalOperator::Sort { sort_items, .. } = + LogicalPlanner::new(&config).plan(&query).unwrap() + else { + panic!("expected sort"); + }; + assert_eq!(sort_items[0].expression, projections[1].expression); + } + #[test] fn test_relationship_query_logical_plan_structure() { let query_text = r#"MATCH (p:Person {name: "Alice"})-[:KNOWS]->(f:Person) WHERE f.age > 30 RETURN f.name"#; diff --git a/docs/cypher.md b/docs/cypher.md index a299aaa..7266eda 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -292,10 +292,10 @@ print(result.to_pylist()) Keywords, function names, variables, labels, relationship types and property names are all case-insensitive. `match (P:person) return p.NAME` works, and the output column keeps the spelling you wrote. -!!! bug "Known issue: an alias that shadows a variable" - - `RETURN p.name AS p` reuses the variable name `p` as an output alias, which can fail to plan once the - query also aggregates or uses `OPTIONAL MATCH`. Pick a different alias, such as `AS name`. +An output alias may reuse an input variable's name. Expressions in the same `RETURN` still use the input +variable; the following `ORDER BY` uses the output alias. For example, +`RETURN p.name AS p, p.id AS id ORDER BY p` returns names and IDs, sorted by name. This also works with +aggregation and `OPTIONAL MATCH`. ## Errors diff --git a/python/python/tests/test_alias_shadowing.py b/python/python/tests/test_alias_shadowing.py new file mode 100644 index 0000000..39f2b99 --- /dev/null +++ b/python/python/tests/test_alias_shadowing.py @@ -0,0 +1,74 @@ +"""Output aliases shadow input nodes only after their expressions are evaluated.""" + +import pytest +from lance_csr import CypherEngine +from test_engine_csr import CONFIG, KNOWS, PEOPLE + + +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize( + ("query", "columns", "expected"), + [ + ( + "MATCH (p:Person) RETURN p.name AS p ORDER BY p", + ["p"], + [(name,) for name in "abcdef"], + ), + ( + "MATCH (p:Person) RETURN p.name AS p, count(*) AS n ORDER BY p", + ["p", "n"], + [(name, 1) for name in "abcdef"], + ), + ( + "MATCH (p:Person) OPTIONAL MATCH (p)-[:KNOWS]->(f:Person) " + "RETURN p.name AS p, f.name AS f ORDER BY p, f", + ["p", "f"], + [ + ("a", "b"), + ("a", "b"), + ("a", "c"), + ("b", "c"), + ("c", "c"), + ("d", None), + ("e", "a"), + ("f", None), + ], + ), + ( + "MATCH (p:Person) OPTIONAL MATCH (p)-[:KNOWS]->(f:Person) " + "RETURN p.name AS p, count(f) AS n ORDER BY p", + ["p", "n"], + list(zip("abcdef", [3, 1, 1, 0, 1, 0], strict=True)), + ), + ( + "MATCH (p:Person) RETURN p.name AS p, p.id AS id ORDER BY p", + ["p", "id"], + list(zip("abcdef", range(1, 7), strict=True)), + ), + ( + "MATCH (p:Person) RETURN p.name AS P ORDER BY p DESC", + ["p"], + [(name,) for name in "fedcba"], + ), + ( + "MATCH (p:Person) RETURN p.name AS p ORDER BY toUpper(p) DESC", + ["p"], + [(name,) for name in "fedcba"], + ), + ( + "MATCH (p:Person) WITH p.name AS p ORDER BY p RETURN p", + ["p"], + [(name,) for name in "abcdef"], + ), + ( + "MATCH (p:Person) RETURN p AS p ORDER BY p", + ["p.id"], + [(key,) for key in range(1, 7)], + ), + ], +) +def test_output_alias_shadows_node(use_csr, query, columns, expected): + engine = CypherEngine(CONFIG, {"Person": PEOPLE, "KNOWS": KNOWS}, use_csr=use_csr) + result = engine.execute(query) + assert result.column_names == columns + assert [tuple(row.values()) for row in result.to_pylist()] == expected From 2c5e1ac5284a8fec908fe462da0c8cc86cf9f0ca Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:19:11 -0700 Subject: [PATCH 12/66] fix(cypher): lower empty maps as empty struct values DataFusion named_struct requires arguments. Use an empty Struct scalar for empty literal and parameter maps. Covered by the container parameter round-trip regression. Part of #20. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/src/datafusion_planner/expression.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/crates/lance-graph/src/datafusion_planner/expression.rs b/crates/lance-graph/src/datafusion_planner/expression.rs index 66fe7c7..1ec7894 100644 --- a/crates/lance-graph/src/datafusion_planner/expression.rs +++ b/crates/lance-graph/src/datafusion_planner/expression.rs @@ -448,6 +448,9 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { datafusion::logical_expr::placeholder(format!("${}", name.to_lowercase())) } VE::Boolean(b) => to_df_boolean_expr(b), + VE::Map(entries) if entries.is_empty() => lit(datafusion::common::ScalarValue::Struct( + std::sync::Arc::new(arrow::array::StructArray::new_empty_fields(1, None)), + )), VE::Map(entries) => datafusion::functions::core::expr_fn::named_struct( entries .iter() From 0dbe2e1d493331868fab5825d0e86b9fb85a96d2 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:26:37 -0700 Subject: [PATCH 13/66] perf(cypher): count unmapped comprehensions without list aggregation Specialize size() of unmapped list comprehensions to filtered COUNT(1), preserving null elements, duplicates, padding and input-row multiplicity. Keep mapped and pattern comprehensions on their existing path. Document the non-nullable Arrow count field and add focused plan/result tests plus a reference benchmark. Part of #18. Isolated Rust suite: 349 passed, one optional timing diagnostic ignored. Formatting, focused Ruff and strict docs passed. Python, benchmark and SF1 measurements pending; no performance claim. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph-benches/README.md | 14 +++ .../benches/graph_execution.rs | 46 ++++++- .../builder/comprehension_ops.rs | 113 ++++++++++++++++-- .../src/datafusion_planner/builder/mod.rs | 2 + crates/lance-graph/src/logical_plan.rs | 22 ++++ docs/api-reference.md | 2 + docs/cypher.md | 7 ++ .../python/tests/test_comprehension_count.py | 105 ++++++++++++++++ 8 files changed, 303 insertions(+), 8 deletions(-) create mode 100644 python/python/tests/test_comprehension_count.py diff --git a/crates/lance-graph-benches/README.md b/crates/lance-graph-benches/README.md index fc00241..2368aef 100644 --- a/crates/lance-graph-benches/README.md +++ b/crates/lance-graph-benches/README.md @@ -50,3 +50,17 @@ calls DataFusion's query planner directly so optimization is counted only once. and their ID checksum. These local measurements do not replace the two SF1 suites required for performance submissions; record runtime, hardware, revision and build profile with any reported numbers. The engine's Python cache also includes dictionary conversion and lookup, which this Rust microbenchmark does not time. + + +## Comprehension counts + +Compare `size([x IN list WHERE predicate])` with the equivalent identity-mapped comprehension, which keeps the +ordinary list-building path. Both queries run against 1,000 input rows with list lengths 16 and 256; results, +row count and summed counts must match before timing. This measures planning plus execution after parsing. + +```bash +cargo bench -p lance-graph-benches --bench graph_execution -- comprehension_count --warm-up-time 1 +``` + +The count path removes ordered list aggregation. It still numbers and unnests rows, evaluates predicates and +groups by input row, so these measurements do not establish a general bound on planning or memory use. diff --git a/crates/lance-graph-benches/benches/graph_execution.rs b/crates/lance-graph-benches/benches/graph_execution.rs index 333ec6e..bc1f959 100644 --- a/crates/lance-graph-benches/benches/graph_execution.rs +++ b/crates/lance-graph-benches/benches/graph_execution.rs @@ -406,5 +406,49 @@ fn bench_query_templates(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_cypher_execution, bench_query_templates); +fn bench_comprehension_counts(c: &mut Criterion) { + use arrow_array::Int64Array; + let rt = tokio::runtime::Runtime::new().unwrap(); + let config = GraphConfig::builder() + .node_table("Person", "person_id") + .build() + .unwrap(); + let data = HashMap::from([("Person".to_string(), make_people_batch(1000))]); + let mut group = c.benchmark_group("comprehension_count"); + for length in [16, 256] { + let query = |map| { + CypherQuery::new(&format!( + "MATCH (p:Person) RETURN p.person_id AS id, size([x IN range(0, {}) WHERE x % 2 = 0 {map}]) AS n ORDER BY id", length - 1 + )).unwrap().with_config(config.clone()) + }; + let count = query(""); + let reference = query("| x"); + let actual = execute_cypher_query(&rt, &count, data.clone()); + let expected = execute_cypher_query(&rt, &reference, data.clone()); + assert_eq!(actual.columns(), expected.columns()); + assert_eq!(actual.num_rows(), 1000); + let counts = actual + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(counts.values().iter().sum::(), 1000 * length / 2); + for (name, query) in [ + ("filtered_count", &count), + ("materialized_reference", &reference), + ] { + group.bench_with_input(BenchmarkId::new(name, length), &length, |b, _| { + b.iter(|| black_box(execute_cypher_query(&rt, query, data.clone()))); + }); + } + } + group.finish(); +} + +criterion_group!( + benches, + bench_cypher_execution, + bench_query_templates, + bench_comprehension_counts +); criterion_main!(benches); diff --git a/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs index 8858abc..45f6f4e 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs @@ -9,6 +9,7 @@ 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_aggregate::count::count; 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::{ @@ -42,6 +43,7 @@ impl DataFusionPlanner { keep: Option<&BooleanExpression>, map: &ValueExpression, output: &str, + count_only: bool, ) -> Result { let error = |what: &str| { let what = what.to_string(); @@ -132,12 +134,18 @@ impl DataFusionPlanner { if let Some(keep) = keep { kept = kept.and(self.boolean_expr(&bind_boolean(keep, elements.schema()))); } - let value = self.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"))? + let aggregate = if count_only { + // Count rows, not the mapped value: null elements still occupy a slot. + count(lit(1_i64)).filter(kept).build() + } else { + let value = self.value_expr(&bind_value(map, elements.schema())); + array_agg(value) + .filter(kept) + .order_by(vec![col(&pos).sort(true, false)]) + .build() + }; + let collected = aggregate + .map_err(error("Failed to aggregate comprehension values"))? .alias(output); // One group per input row: its number plus every input column, all equal within it. let names: HashSet = numbered @@ -160,7 +168,12 @@ impl DataFusionPlanner { .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)); + let value = if count_only { + col(output) + } else { + coalesce(vec![col(output), make_array(vec![])]) + }; + out.push(value.alias(output)); LogicalPlanBuilder::from(elements) .aggregate(group, vec![collected]) .and_then(|b| b.project(out)) @@ -168,3 +181,89 @@ impl DataFusionPlanner { .map_err(error("Failed to build list comprehension")) } } + +#[cfg(test)] +mod tests { + use arrow_array::{types::Int64Type, ArrayRef, Int64Array, ListArray, RecordBatch}; + use std::collections::HashMap; + use std::sync::Arc; + + fn data() -> HashMap { + HashMap::from([( + "Person".into(), + RecordBatch::try_from_iter([ + ( + "id", + Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef, + ), + ( + "values", + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), None, Some(1), Some(2)]), + Some(vec![]), + None, + Some(vec![None]), + ])) as ArrayRef, + ), + ]) + .unwrap(), + )]) + } + + #[tokio::test] + async fn comprehension_counts_keep_nulls_duplicates_and_empty_rows() { + let config = crate::GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + for (predicate, expected) in [ + ("", vec![4, 0, 0, 1]), + ("WHERE x > 1", vec![1, 0, 0, 0]), + ("WHERE x IS NULL", vec![1, 0, 0, 1]), + ("WHERE x IS NULL OR x > 0", vec![4, 0, 0, 1]), + ] { + // Duplicated input rows must remain distinct; an identity map is the + // unoptimized reference and also retains null elements in ARRAY_AGG. + let make_query = |map| { + crate::CypherQuery::new(&format!( + "MATCH (p:Person) UNWIND [1, 1] AS duplicate RETURN p.id AS id, \ + size([x IN p.values {predicate} {map}]) AS n ORDER BY id" + )) + .unwrap() + .with_config(config.clone()) + }; + let query = make_query(""); + let actual = query.execute(data(), None).await.unwrap(); + let reference = make_query("| x").execute(data(), None).await.unwrap(); + assert_eq!(actual.columns(), reference.columns(), "{predicate}"); + assert_eq!( + actual.schema().field(1).data_type(), + reference.schema().field(1).data_type() + ); + assert_eq!(actual.schema().field(0).name(), "id"); + assert_eq!(actual.schema().field(1).name(), "n"); + assert!(!actual.schema().field(1).is_nullable()); + let counts = actual + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + counts.values().to_vec(), + expected + .into_iter() + .flat_map(|n| [n, n]) + .collect::>() + ); + let plan = query.explain(data()).await.unwrap().to_lowercase(); + assert!(plan.contains("count("), "{plan}"); + assert!(!plan.contains("array_agg("), "{plan}"); + let mapped = make_query("| x") + .explain(data()) + .await + .unwrap() + .to_lowercase(); + assert!(mapped.contains("array_agg("), "{mapped}"); + } + } +} diff --git a/crates/lance-graph/src/datafusion_planner/builder/mod.rs b/crates/lance-graph/src/datafusion_planner/builder/mod.rs index c2ca6a4..5f07235 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/mod.rs @@ -117,6 +117,7 @@ impl DataFusionPlanner { keep, map, output, + count_only, } => self.build_list_comprehension( ctx, input, @@ -126,6 +127,7 @@ impl DataFusionPlanner { keep.as_ref(), map, output, + *count_only, ), LogicalOperator::ComprehensionElement => Ok(ctx .comprehension_elements diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index a85a642..e8ccc63 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -172,6 +172,9 @@ pub enum LogicalOperator { keep: Option, map: ValueExpression, output: String, + /// Count retained elements without building an ordered list. + #[serde(default)] + count_only: bool, }, /// The element rows of the enclosing [`LogicalOperator::ListComprehension`]. @@ -754,6 +757,24 @@ impl<'a> LogicalPlanner<'a> { None => Ok(input), } } + // The unmapped comprehension retains every passing element, including + // nulls and duplicates. Its size needs a filtered count, not ARRAY_AGG. + // Keep mapped comprehensions unchanged so their evaluation/errors remain. + ValueExpression::ScalarFunction { name, args } + if name.eq_ignore_ascii_case("size") + && matches!( + args.as_slice(), + [ValueExpression::ListComprehension { map: None, .. }] + ) => + { + let mut lowered = self.rewrite_exists_in_value(input, &mut args[0])?; + let LogicalOperator::ListComprehension { count_only, .. } = &mut lowered else { + unreachable!("a list comprehension lowers to its own operator") + }; + *count_only = true; + *expr = args.remove(0); + Ok(lowered) + } ValueExpression::ScalarFunction { args, .. } | ValueExpression::AggregateFunction { args, .. } => { for arg in args.iter_mut() { @@ -901,6 +922,7 @@ impl<'a> LogicalPlanner<'a> { keep, map: value, output: output.clone(), + count_only: false, }) })(); self.variables = before; diff --git a/docs/api-reference.md b/docs/api-reference.md index d4150be..7764763 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -107,6 +107,8 @@ be opened. : Runs `query` and returns the whole result. `params` binds `$name` parameters: `int`, `float`, `str`, `bool`, `None`, lists, string-keyed nested maps, `datetime.datetime` and `datetime.date`. Integers must fit signed 64-bit values; nesting beyond 128 levels, cycles and unsupported values raise `ValueError`. + Direct counts from `size()` of an unmapped list comprehension have a non-nullable Arrow field; + see [Cypher functions](cypher.md#functions). `#!python execute_batches(query: str, params: dict | None = None) -> pyarrow.RecordBatchReader` : Streams the result. Planning errors raise immediately; execution errors raise while reading. The reader diff --git a/docs/cypher.md b/docs/cypher.md index 7266eda..d8fda29 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -269,6 +269,13 @@ Any other function raises `NotImplementedError: Cypher function 'name' is not im missing: math functions (`abs`, `round`, `sqrt`, ...), most string and list functions, `reduce`, the list predicates `any` / `all` / `none` / `single`, and `relationships(p)`. +When `size()` consumes an unmapped list comprehension, such as `size([x IN values WHERE x > 0])`, the planner +counts retained elements directly. Each duplicate and each retained null element still contributes one; +an empty or null input list produces zero, as for the existing comprehension. A mapped comprehension +(`| expression`) or a pattern comprehension keeps the ordinary list-building path. The count path still +unnests and groups rows; `explain()` shows a filtered count in place of ordered list aggregation. A direct +count output is marked non-nullable in its Arrow schema; the previous list-size expression was marked nullable. + ### Example: a little of everything ```python diff --git a/python/python/tests/test_comprehension_count.py b/python/python/tests/test_comprehension_count.py new file mode 100644 index 0000000..ad04322 --- /dev/null +++ b/python/python/tests/test_comprehension_count.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Count-only comprehensions preserve multiplicity and predicate semantics.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture(params=[False, True]) +def engine(request): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + return CypherEngine( + config, + { + "Person": pa.table( + { + "id": [1, 2, 3, 4], + "values": pa.array( + [[1, None, 1, 2], [], None, [None]], type=pa.list_(pa.int64()) + ), + } + ), + # A duplicate edge affects the list's multiplicity; repeated matches + # inside EXISTS must still contribute a single truth value per element. + "KNOWS": pa.table({"src": [1, 1, 1, 2, 2], "dst": [2, 2, 3, 3, 3]}), + }, + use_csr=request.param, + ) + + +@pytest.mark.parametrize( + ("predicate", "expected"), + [ + ("", [4, 0, 0, 1]), + ("WHERE x > 1", [1, 0, 0, 0]), + ("WHERE x IS NULL", [1, 0, 0, 1]), + ("WHERE x IS NULL OR x > 0", [4, 0, 0, 1]), + ], +) +def test_count_matches_list_reference(engine, predicate, expected): + query = ( + "MATCH (p:Person) UNWIND [1, 1] AS duplicate RETURN p.id AS id, " + f"size([x IN p.values {predicate}]) AS n ORDER BY id" + ) + reference = query.replace(f"{predicate}])", f"{predicate} | x])") + rows = [ + {"id": index, "n": n} for index, n in enumerate(expected, 1) for _ in range(2) + ] + assert engine.execute(query).to_pylist() == rows + assert engine.execute_batches(query).read_all().to_pylist() == rows + assert engine.execute(reference).to_pylist() == rows + assert "array_agg(" not in engine.explain(query).lower() + assert "array_agg(" in engine.explain(reference).lower() + + +def test_count_with_node_lookup_and_exists_keeps_duplicate_elements(engine): + query = ( + "MATCH (p:Person) OPTIONAL MATCH (p)-[:KNOWS]->(f:Person) " + "WITH p, collect(f) AS friends " + "RETURN p.id AS id, size([f IN friends WHERE (f)-[:KNOWS]->(:Person)]) AS n " + "ORDER BY id" + ) + expected = [ + {"id": 1, "n": 2}, + {"id": 2, "n": 0}, + {"id": 3, "n": 0}, + {"id": 4, "n": 0}, + ] + assert engine.execute(query).to_pylist() == expected + assert ( + engine.execute(query.replace("(:Person)])", "(:Person) | f])")).to_pylist() + == expected + ) + + +def test_mapped_and_pattern_comprehensions_keep_normal_list_semantics(engine): + query = ( + "MATCH (p:Person) RETURN p.id AS id, " + "size([x IN p.values | null]) AS n ORDER BY id" + ) + assert engine.execute(query).to_pylist() == [ + {"id": 1, "n": 4}, + {"id": 2, "n": 0}, + {"id": 3, "n": 0}, + {"id": 4, "n": 1}, + ] + assert "array_agg(" in engine.explain(query).lower() + pattern = ( + "MATCH (p:Person) RETURN p.id AS id, " + "size([(p)-[:KNOWS]->(f) | f.id]) AS n ORDER BY id" + ) + assert engine.execute(pattern).to_pylist() == [ + {"id": 1, "n": 3}, + {"id": 2, "n": 2}, + {"id": 3, "n": 0}, + {"id": 4, "n": 0}, + ] + assert "array_agg(" in engine.explain(pattern).lower() From 1ac21f9c7797be5165113c265417fa80b5f87fdf Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:19:39 -0700 Subject: [PATCH 14/66] feat(cypher): add math, string, and list scalar functions Implement the bounded issue #20 scalar-function group with strict argument kinds, Unicode string operations, list operations, and stepped ranges. Preserve Cypher null, invalid-length, and rounding behavior. Accept empty quoted string literals and document the supported signatures and limits. Validation: 343 Rust library tests, 176 focused Python tests, workspace all-target clippy, Python lint, strict docs, spelling, and executable docs pass. Full Python run: 581 pass, two prerequisite empty-map parameter failures; the prerequisite correction is being integrated separately. Submitted by David Hughes (@stugorf), with AI Collaboration. --- .../src/datafusion_planner/expression.rs | 24 +- .../lance-graph/src/datafusion_planner/mod.rs | 1 + .../datafusion_planner/scalar_functions.rs | 489 ++++++++++++++++++ crates/lance-graph/src/parser.rs | 6 +- crates/lance-graph/src/semantic.rs | 27 +- docs/cypher.md | 41 +- python/python/tests/test_scalar_functions.py | 203 +++++++- 7 files changed, 753 insertions(+), 38 deletions(-) create mode 100644 crates/lance-graph/src/datafusion_planner/scalar_functions.rs diff --git a/crates/lance-graph/src/datafusion_planner/expression.rs b/crates/lance-graph/src/datafusion_planner/expression.rs index 1ec7894..0000f36 100644 --- a/crates/lance-graph/src/datafusion_planner/expression.rs +++ b/crates/lance-graph/src/datafusion_planner/expression.rs @@ -227,18 +227,18 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { VE::Variable(path) => col(qualify_column(path, &name.to_lowercase())), _ => Expr::Literal(datafusion::scalar::ScalarValue::Null, None), }, - "size" => Expr::Cast(datafusion::logical_expr::Cast::new( - Box::new(datafusion::functions_nested::expr_fn::cardinality( - to_df_value_expr(&args[0]), - )), - arrow_schema::DataType::Int64, - )), - // Cypher's range includes its upper bound, like DataFusion's gen_series. - "range" => datafusion::functions_nested::expr_fn::gen_series( - to_df_value_expr(&args[0]), - to_df_value_expr(&args[1]), - lit(1_i64), - ), + "abs" | "round" | "floor" | "ceil" | "sqrt" | "left" | "right" | "substring" + | "replace" | "split" | "trim" | "ltrim" | "rtrim" | "size" | "reverse" + | "last" | "tail" => { + super::scalar_functions::call(name, args.iter().map(to_df_value_expr).collect()) + } + "range" => { + let mut args: Vec<_> = args.iter().map(to_df_value_expr).collect(); + if args.len() == 2 { + args.push(lit(1_i64)); + } + super::scalar_functions::call(name, args) + } // list[i] (0-based in Cypher, 1-based in DataFusion's array_element). "__index" => datafusion::functions_nested::expr_fn::array_element( to_df_value_expr(&args[0]), diff --git a/crates/lance-graph/src/datafusion_planner/mod.rs b/crates/lance-graph/src/datafusion_planner/mod.rs index 8c58f29..f1f5d11 100644 --- a/crates/lance-graph/src/datafusion_planner/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/mod.rs @@ -20,6 +20,7 @@ mod config_helpers; mod cypher_udf; mod expression; mod join_ops; +mod scalar_functions; mod scan_ops; mod udf; pub mod vector_ops; diff --git a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs new file mode 100644 index 0000000..94cd8d7 --- /dev/null +++ b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs @@ -0,0 +1,489 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Cypher scalar semantics: strict argument kinds, Unicode characters, and nulls. +//! List operations reuse DataFusion's Arrow kernels; string slicing differs from SQL. + +use std::sync::Arc; + +use arrow::array::{ + new_null_array, Array, ArrayRef, AsArray, Float64Array, Int64Array, ListBuilder, StringArray, + StringBuilder, +}; +use arrow::compute::cast; +use arrow_schema::{DataType, Field}; +use datafusion::common::{exec_err, plan_err, ScalarValue}; +use datafusion::error::Result; +use datafusion::functions_nested::{extract, length, range, reverse}; +use datafusion::logical_expr::{ + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; + +#[derive(Debug, PartialEq, Eq, Hash)] +struct CypherScalar { + name: String, + signature: Signature, +} + +fn is_string(t: &DataType) -> bool { + matches!(t, DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View) +} + +fn list_field(t: &DataType) -> Option<&Arc> { + match t { + DataType::List(f) | DataType::LargeList(f) | DataType::FixedSizeList(f, _) => Some(f), + _ => None, + } +} + +impl ScalarUDFImpl for CypherScalar { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, types: &[DataType]) -> Result { + for (i, t) in types.iter().enumerate() { + let valid = match self.name.as_str() { + "abs" | "round" | "floor" | "ceil" | "sqrt" => t.is_integer() || t.is_floating(), + "range" => t.is_integer(), + "size" | "reverse" => is_string(t) || list_field(t).is_some(), + "last" | "tail" => list_field(t).is_some(), + "left" | "right" | "substring" if i > 0 => t.is_integer(), + _ => is_string(t), + }; + if !valid && !t.is_null() { + return plan_err!("{} does not accept {t} for argument {}", self.name, i + 1); + } + } + Ok(match self.name.as_str() { + "size" => DataType::Int64, + "abs" if types[0].is_integer() => DataType::Int64, + "abs" | "round" | "floor" | "ceil" | "sqrt" => DataType::Float64, + "range" => DataType::List(Arc::new(Field::new("item", DataType::Int64, true))), + "split" => DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + "last" => list_field(&types[0]) + .map(|f| f.data_type().clone()) + .unwrap_or(DataType::Null), + "tail" => match &types[0] { + DataType::FixedSizeList(f, _) => DataType::List(f.clone()), + t => t.clone(), + }, + "reverse" if list_field(&types[0]).is_some() => types[0].clone(), + _ => DataType::Utf8, + }) + } + + fn invoke_with_args(&self, mut args: ScalarFunctionArgs) -> Result { + let output_type = args.return_type().clone(); + if args.args[0].data_type().is_null() && self.name != "range" { + return Ok(ColumnarValue::Array(new_null_array( + &output_type, + args.number_rows, + ))); + } + if self.name == "range" { + for arg in &mut args.args { + *arg = + ColumnarValue::Array(cast(&arg.to_array(args.number_rows)?, &DataType::Int64)?); + } + return range::gen_series_udf().invoke_with_args(args); + } + if list_field(&args.args[0].data_type()).is_some() { + let udf = match self.name.as_str() { + "size" => length::array_length_udf(), + "reverse" => reverse::array_reverse_udf(), + "last" => { + if let DataType::FixedSizeList(field, _) = args.args[0].data_type() { + args.args[0] = ColumnarValue::Array(cast( + &args.args[0].to_array(args.number_rows)?, + &DataType::List(field), + )?); + } + args.args + .push(ColumnarValue::Scalar(ScalarValue::Int64(Some(-1)))); + extract::array_element_udf() + } + "tail" => { + args.args[0] = ColumnarValue::Array(cast( + &args.args[0].to_array(args.number_rows)?, + &output_type, + )?); + args.args.extend( + [2, -1].map(|n| ColumnarValue::Scalar(ScalarValue::Int64(Some(n)))), + ); + extract::array_slice_udf() + } + _ => unreachable!("validated list function"), + }; + let rows = args.number_rows; + args.arg_fields = args + .args + .iter() + .map(|arg| Arc::new(Field::new("arg", arg.data_type(), true))) + .collect(); + args.return_field = Arc::new(Field::new( + udf.name(), + udf.return_type( + &args + .args + .iter() + .map(ColumnarValue::data_type) + .collect::>(), + )?, + true, + )); + let result = udf.invoke_with_args(args)?; + return if self.name == "size" { + Ok(ColumnarValue::Array(cast( + &result.to_array(rows)?, + &DataType::Int64, + )?)) + } else { + Ok(result) + }; + } + if matches!( + self.name.as_str(), + "abs" | "round" | "floor" | "ceil" | "sqrt" + ) { + let input = cast(&args.args[0].to_array(args.number_rows)?, &output_type)?; + let result: ArrayRef = if output_type == DataType::Int64 { + Arc::new( + input + .as_primitive::() + .iter() + .map(|n| { + n.map(|n| { + n.checked_abs().ok_or_else(|| { + datafusion::common::exec_datafusion_err!("abs overflow") + }) + }) + .transpose() + }) + .collect::>()?, + ) + } else { + Arc::new( + input + .as_primitive::() + .iter() + .map(|n| { + n.map(|n| match self.name.as_str() { + "abs" => n.abs(), + // Cypher rounds ties toward positive infinity and returns positive zero. + "round" => { + let rounded = if n.fract() == -0.5 { + n.ceil() + } else { + n.round() + }; + if rounded == 0.0 { + 0.0 + } else { + rounded + } + } + "floor" => n.floor(), + "ceil" => n.ceil(), + "sqrt" => n.sqrt(), + _ => unreachable!(), + }) + }) + .collect::(), + ) + }; + return Ok(ColumnarValue::Array(result)); + } + let arrays = args + .args + .iter() + .enumerate() + .map(|(i, arg)| { + let t = if i > 0 && matches!(self.name.as_str(), "left" | "right" | "substring") { + DataType::Int64 + } else { + DataType::Utf8 + }; + cast(&arg.to_array(args.number_rows)?, &t) + }) + .collect::, _>>()?; + let strings = arrays[0].as_string::(); + if self.name == "size" { + return Ok(ColumnarValue::Array(Arc::new( + strings + .iter() + .map(|s| s.map(|s| s.chars().count() as i64)) + .collect::(), + ))); + } + if self.name == "split" { + let delimiters = arrays[1].as_string::(); + let mut result = ListBuilder::new(StringBuilder::new()); + for (s, delimiter) in strings.iter().zip(delimiters.iter()) { + match (s, delimiter) { + (Some(s), Some(delimiter)) => { + if s.is_empty() { + result.values().append_value(""); + } else if delimiter.is_empty() { + for c in s.chars() { + result.values().append_value(c.to_string()); + } + } else { + for part in s.split(delimiter) { + result.values().append_value(part); + } + } + result.append(true); + } + _ => result.append(false), + } + } + return Ok(ColumnarValue::Array(Arc::new(result.finish()))); + } + let result = strings + .iter() + .enumerate() + .map(|(row, s)| { + let Some(s) = s else { return Ok(None) }; + let value = match self.name.as_str() { + "trim" => s.trim().to_string(), + "ltrim" => s.trim_start().to_string(), + "rtrim" => s.trim_end().to_string(), + "reverse" => s.chars().rev().collect(), + "replace" => { + if arrays[1].is_null(row) || arrays[2].is_null(row) { + return Ok(None); + } + s.replace( + arrays[1].as_string::().value(row), + arrays[2].as_string::().value(row), + ) + } + "left" | "right" | "substring" => { + let index = |i: usize| -> Result { + if arrays[i].is_null(row) { + return exec_err!( + "{} requires a non-null integer for argument {}", + self.name, + i + 1 + ); + } + let n = arrays[i] + .as_primitive::() + .value(row); + if n < 0 { + return exec_err!( + "{} requires a non-negative integer for argument {}", + self.name, + i + 1 + ); + } + Ok(n as usize) + }; + let n = index(1)?; + let (start, count) = match self.name.as_str() { + "left" => (0, n), + "right" => (s.chars().count().saturating_sub(n), n), + _ => ( + n, + if arrays.len() == 3 { + index(2)? + } else { + usize::MAX + }, + ), + }; + s.chars().skip(start).take(count).collect() + } + _ => unreachable!("validated string function"), + }; + Ok(Some(value)) + }) + .collect::>()?; + Ok(ColumnarValue::Array(Arc::new(result))) + } +} + +pub(super) fn call(name: &str, args: Vec) -> Expr { + ScalarUDF::new_from_impl(CypherScalar { + name: name.to_lowercase(), + signature: Signature::any(args.len(), Volatility::Immutable), + }) + .call(args) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{config::GraphConfig, query::CypherQuery}; + use arrow::array::RecordBatch; + use std::collections::HashMap; + + async fn run(expression: &str) -> crate::error::Result { + let batch = + RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(vec![1])) as ArrayRef)]) + .unwrap(); + CypherQuery::new(&format!("RETURN {expression} AS value"))? + .with_config( + GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(), + ) + .execute(HashMap::from([("Person".to_string(), batch)]), None) + .await + } + + #[tokio::test] + async fn cypher_scalar_unicode_and_slicing() { + for (expression, expected) in [ + ("left('a猫😀z', 3)", "a猫😀"), + ("right('a猫😀z', 2)", "😀z"), + ("substring('a猫😀z', 1, 2)", "猫😀"), + ("substring('a猫😀z', 2)", "😀z"), + ("substring('x', 99, 2)", ""), + ("right('x', 0)", ""), + ("left('x', 99)", "x"), + ("reverse('a猫😀z')", "z😀猫a"), + ("trim(' 猫 ')", "猫"), + ("lTrim(' 猫 ')", "猫 "), + ("rTrim(' 猫 ')", " 猫"), + ("replace('a猫a', 'a', '😀')", "😀猫😀"), + ] { + let out = run(expression).await.unwrap(); + assert_eq!( + out.column(0).as_string::().value(0), + expected, + "{expression}" + ); + } + let out = run("size('a猫😀z')").await.unwrap(); + assert_eq!( + out.column(0) + .as_primitive::() + .value(0), + 4 + ); + } + + #[tokio::test] + async fn cypher_scalar_math_and_rounding() { + for (expression, expected) in [ + ("abs(-1.5)", 1.5), + ("round(-1.5)", -1.0), + ("round(1.5)", 2.0), + ("round(4503599627370495.5)", 4503599627370496.0), + ("round(4503599627370497.0)", 4503599627370497.0), + ("floor(-1.5)", -2.0), + ("ceil(-1.5)", -1.0), + ("sqrt(9)", 3.0), + ] { + let out = run(expression).await.unwrap(); + assert_eq!( + out.column(0) + .as_primitive::() + .value(0), + expected, + "{expression}" + ); + } + let out = run("abs(-5)").await.unwrap(); + assert_eq!( + out.column(0) + .as_primitive::() + .value(0), + 5 + ); + assert!(run("abs(-9223372036854775807 - 1)").await.is_err()); + } + + #[tokio::test] + async fn cypher_scalar_list_operations() { + for (expression, expected) in [ + ("reverse([1, 2, 3])", vec![3, 2, 1]), + ("tail([1, 2, 3])", vec![2, 3]), + ("tail([1])", vec![]), + ("range(5, 1, -2)", vec![5, 3, 1]), + ("range(1, 5, -2)", vec![]), + ("range(0, 4, 2)", vec![0, 2, 4]), + ] { + let out = run(expression).await.unwrap(); + let values = out.column(0).as_list::().value(0); + let values: Vec<_> = values + .as_primitive::() + .values() + .to_vec(); + assert_eq!(values, expected, "{expression}"); + } + for (expression, expected) in [("last([1, 2, 3])", 3), ("size([[1, 2], [3]])", 2)] { + let out = run(expression).await.unwrap(); + assert_eq!( + out.column(0) + .as_primitive::() + .value(0), + expected + ); + } + for (expression, expected) in [ + ("split('a,,b,', ',')", vec!["a", "", "b", ""]), + ("split('猫😀', '')", vec!["猫", "😀"]), + ("split('', ',')", vec![""]), + ] { + let out = run(expression).await.unwrap(); + let values = out.column(0).as_list::().value(0); + assert_eq!( + values + .as_string::() + .iter() + .map(Option::unwrap) + .collect::>(), + expected + ); + } + } + + #[tokio::test] + async fn cypher_scalar_nulls_and_invalid_arguments() { + for expression in [ + "abs(null)", + "round(null)", + "size(null)", + "reverse(null)", + "last(null)", + "tail(null)", + "last([])", + "left(null, null)", + "substring(null, -1)", + "replace('x', null, '')", + "split('x', null)", + "range(null, 3)", + ] { + let out = run(expression).await.unwrap(); + assert_eq!(out.column(0).logical_null_count(), 1, "{expression}"); + } + for expression in [ + "left('x', -1)", + "right('x', null)", + "substring('x', -1)", + "substring('x', 0, null)", + "range(1, 3, 0)", + "range(1.5, 3)", + "abs('1')", + "size(1)", + "reverse(1)", + "tail('abc')", + "last(2)", + "left('x', 1.5)", + "split('x', 1)", + "round(1, 2)", + "substring('x')", + "range(1)", + ] { + assert!(run(expression).await.is_err(), "{expression} should fail"); + } + } +} diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index 2f6035b..260240b 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -10,7 +10,7 @@ use crate::ast::*; use crate::error::{GraphError, Result}; use nom::{ branch::alt, - bytes::complete::{tag, tag_no_case, take_while1}, + bytes::complete::{tag, tag_no_case, take_while, take_while1}, character::complete::{char, digit0, digit1, multispace0, multispace1, one_of, satisfy}, combinator::{map, map_res, not, opt, peek, recognize}, multi::{many0, many1, separated_list0, separated_list1}, @@ -1382,14 +1382,14 @@ fn string_literal(input: &str) -> IResult<&str, String> { fn double_quoted_string(input: &str) -> IResult<&str, String> { let (input, _) = char('"')(input)?; - let (input, content) = take_while1(|c| c != '"')(input)?; + let (input, content) = take_while(|c| c != '"')(input)?; let (input, _) = char('"')(input)?; Ok((input, content.to_string())) } fn single_quoted_string(input: &str) -> IResult<&str, String> { let (input, _) = char('\'')(input)?; - let (input, content) = take_while1(|c| c != '\'')(input)?; + let (input, content) = take_while(|c| c != '\'')(input)?; let (input, _) = char('\'')(input)?; Ok((input, content.to_string())) } diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index 61679ea..a418d7d 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -483,7 +483,9 @@ impl SemanticAnalyzer { let function_name = name.to_lowercase(); // Validate arity and known functions match function_name.as_str() { - "tolower" | "lower" | "toupper" | "upper" => { + "tolower" | "lower" | "toupper" | "upper" | "abs" | "round" | "floor" + | "ceil" | "sqrt" | "trim" | "ltrim" | "rtrim" | "reverse" | "last" + | "tail" => { if args.len() != 1 { return Err(GraphError::PlanError { message: format!( @@ -495,10 +497,11 @@ impl SemanticAnalyzer { }); } } - "date_part" | "to_epoch_ms" | "length" | "nodes" | "size" | "range" - | "__index" => { + "date_part" | "to_epoch_ms" | "length" | "nodes" | "size" | "left" + | "right" | "split" | "replace" | "__index" => { let want = match function_name.as_str() { - "date_part" | "range" | "__index" => 2, + "date_part" | "__index" | "left" | "right" | "split" => 2, + "replace" => 3, _ => 1, }; if args.len() != want { @@ -513,6 +516,18 @@ impl SemanticAnalyzer { }); } } + "substring" | "range" => { + if !(2..=3).contains(&args.len()) { + return Err(GraphError::PlanError { + message: format!( + "{} requires 2 or 3 arguments, got {}", + name, + args.len() + ), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + } "__list" => {} "head" | "tointeger" | "tofloat" | "date" | "datetime" => { if args.len() != 1 { @@ -551,7 +566,7 @@ impl SemanticAnalyzer { // Unknown scalar function - reject early with helpful error return Err(GraphError::UnsupportedFeature { feature: format!( - "Cypher function '{}' is not implemented. Supported scalar functions: toLower, lower, toUpper, upper, toString, coalesce, date_part, to_epoch_ms, length, nodes, size, range. Supported aggregate functions: COUNT, SUM, AVG, MIN, MAX, COLLECT.", + "Cypher function '{}' is not implemented. Supported scalar functions: toLower, lower, toUpper, upper, toString, coalesce, date_part, to_epoch_ms, length, nodes, size, range, head, last, tail, reverse, left, right, substring, replace, split, trim, lTrim, rTrim, abs, round, floor, ceil, sqrt, toInteger, toFloat, date, datetime, duration.inSeconds. Supported aggregate functions: COUNT, SUM, AVG, MIN, MAX, COLLECT.", name ), location: snafu::Location::new(file!(), line!(), column!()), @@ -1642,7 +1657,7 @@ mod tests { #[test] fn test_unimplemented_scalar_function_fails_validation() { let expr = ValueExpression::ScalarFunction { - name: "replace".to_string(), + name: "not_a_function".to_string(), args: vec![ValueExpression::Property(PropertyRef::new("n", "name"))], }; let result = analyze_return_with_match("n", "Person", expr).unwrap(); diff --git a/docs/cypher.md b/docs/cypher.md index d8fda29..36cf70e 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -258,16 +258,47 @@ RETURN length(p), nodes(p) | Group | Functions | | --- | --- | | Aggregates | `count` (including `count(*)` and `count(DISTINCT x)`), `sum`, `avg`, `min`, `max`, `collect` (including `DISTINCT`; skips nulls) | -| Strings | `toLower` / `lower`, `toUpper` / `upper`, `toString` | +| Strings | `toLower` / `lower`, `toUpper` / `upper`, `toString`, `left`, `right`, `substring`, `replace`, `split`, `trim`, `lTrim`, `rTrim`, `reverse`, `size` | | Conversion | `toInteger`, `toFloat`, `coalesce` | -| Lists | `size`, `head`, `range(start, end)` (inclusive, no step) | +| Lists | `size`, `head`, `last`, `tail`, `reverse`, `range(start, end [, step])` (inclusive) | +| Math | `abs`, `round`, `floor`, `ceil`, `sqrt` | | Temporal | `date`, `datetime`, `date_part('month', d)`, `to_epoch_ms`, `duration.inSeconds(a, b).seconds` | | Paths | `length(p)`, `nodes(p)` on `shortestPath` variables | | Vectors | `vector_distance(a, b, metric)`, `vector_similarity(a, b, metric)`; see [Vector search](vector-search.md) | -Any other function raises `NotImplementedError: Cypher function 'name' is not implemented`. Notably -missing: math functions (`abs`, `round`, `sqrt`, ...), most string and list functions, `reduce`, the list -predicates `any` / `all` / `none` / `single`, and `relationships(p)`. +The string and list functions follow the [Cypher v9 reference](https://s3.amazonaws.com/artifacts.opencypher.org/openCypher9.pdf), +within the list representation described above. `size` counts Unicode characters in a string or the +outer elements of a list. `left(text, length)`, `right(text, length)` and +`substring(text, start [, length])` count characters; the substring start is zero-based. Lengths and +starts must be non-negative integers. A null input string returns null; a null start or length on a +non-null string raises an error. `trim`, `lTrim` and `rTrim` remove whitespace, and `replace` replaces +all literal matches. `split` preserves empty fields; an empty delimiter splits into characters. +Splitting an empty string returns `['']`. + +`last([])` returns null; `tail([])` and `reverse([])` return empty lists. `reverse` also accepts strings. +`range(start, end [, step])` includes the end when reached, defaults to step 1, and accepts negative +steps. A step pointing away from the end gives an empty list; zero steps and non-integer arguments +raise an error. Null arguments return null. + +`abs` preserves integer results; the other math functions return floating point values. Integer +absolute-value overflow raises an error. `round` accepts one argument, with ties toward positive +infinity (`round(-1.5)` is `-1.0`), as specified in the +[Neo4j 4.4 numeric reference](https://neo4j.com/docs/cypher-manual/4.4/functions/mathematical-numeric/). +`sqrt` of a negative number returns NaN. The math functions return null for null input and reject +non-numeric values. + +```python +result = engine.execute(""" + RETURN substring('a猫😀z', 1, 2) AS text, + reverse(range(1, 5, 2)) AS numbers, + round(-1.5) AS rounded +""") +print(result.to_pylist()) # [{'text': '猫😀', 'numbers': [5, 3, 1], 'rounded': -1.0}] +``` + +Any other function raises `ValueError` with `Cypher function 'name' is not implemented`. This is a +bounded function set, not a claim of full Cypher conformance. Remaining functions include `keys`, +`labels`, `reduce`, the list predicates `any` / `all` / `none` / `single`, and `relationships(p)`. When `size()` consumes an unmapped list comprehension, such as `size([x IN values WHERE x > 0])`, the planner counts retained elements directly. Each duplicate and each retained null element still contributes one; diff --git a/python/python/tests/test_scalar_functions.py b/python/python/tests/test_scalar_functions.py index 83eb1b5..196b76d 100644 --- a/python/python/tests/test_scalar_functions.py +++ b/python/python/tests/test_scalar_functions.py @@ -1,22 +1,201 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The Lance Authors +"""Cypher v9 scalar semantics, including values from Arrow-backed graph rows.""" + +import math + import pyarrow as pa import pytest -from lance_csr import CypherQuery, GraphConfig +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture(params=[False, True], ids=["join", "csr"]) +def engine(request): + config = GraphConfig.builder().node_table("Person", "id").build() + people = pa.table( + { + "id": [1, 2, 3], + "text": pa.array(["a猫😀z", None, ""], type=pa.large_string()), + "number": [-1.5, None, 2.0], + "items": pa.array([[1, None, 3], None, []], type=pa.list_(pa.int64())), + "pair": pa.array([[1, 2], None, [3, 4]], type=pa.list_(pa.int64(), 2)), + } + ) + return CypherEngine(config, {"Person": people}, use_csr=request.param) + + +@pytest.mark.parametrize( + ("expression", "expected"), + [ + ("left('a猫😀z', 3)", "a猫😀"), + ("right('a猫😀z', 2)", "😀z"), + ("substring('a猫😀z', 1, 2)", "猫😀"), + ("substring('a猫😀z', 2)", "😀z"), + ("substring('x', 99, 2)", ""), + ("substring('x', 0, 0)", ""), + ("left('x', 0)", ""), + ("right('x', 0)", ""), + ("left('x', 99)", "x"), + ("right('x', 99)", "x"), + ("size('a猫😀z')", 4), + ("size('')", 0), + ('size("")', 0), + ("reverse('a猫😀z')", "z😀猫a"), + ("replace('a猫a', 'a', '😀')", "😀猫😀"), + ("replace('abc', 'x', 'y')", "abc"), + ("split('a,,b,', ',')", ["a", "", "b", ""]), + ("split('猫😀', '')", ["猫", "😀"]), + ("split('', ',')", [""]), + ("split('', '')", [""]), + ("trim(' 猫 ')", "猫"), + ("lTrim(' 猫 ')", "猫 "), + ("rTrim(' 猫 ')", " 猫"), + ("abs(-5)", 5), + ("abs(-1.5)", 1.5), + ("round(-1.5)", -1.0), + ("round(1.5)", 2.0), + ("round(4503599627370497.0)", 4503599627370497.0), + ("floor(-1.5)", -2.0), + ("ceil(-1.5)", -1.0), + ("sqrt(9)", 3.0), + ("last([1, null, 3])", 3), + ("last([1, null])", None), + ("last([])", None), + ("tail([1, null, 3])", [None, 3]), + ("tail([1])", []), + ("tail([])", []), + ("reverse([1, null, 3])", [3, None, 1]), + ("reverse([])", []), + ("reverse([[1, 2], [3]])", [[3], [1, 2]]), + ("size([[1, 2], [3]])", 2), + ("range(0, 4)", [0, 1, 2, 3, 4]), + ("range(0, 4, 2)", [0, 2, 4]), + ("range(5, 1, -2)", [5, 3, 1]), + ("range(1, 5, -2)", []), + ("range(5, 1, 2)", []), + ("range(1, 1, -2)", [1]), + ("range(null, 4)", None), + ("range(1, null)", None), + ("range(1, 4, null)", None), + ("left(null, null)", None), + ("substring(null, -1)", None), + ("replace('x', null, '')", None), + ("split('x', null)", None), + ] + + [ + (f"{name}(null)", None) + for name in ( + "abs", + "round", + "floor", + "ceil", + "sqrt", + "trim", + "lTrim", + "rTrim", + "reverse", + "size", + "last", + "tail", + ) + ], +) +def test_scalar_value(engine, expression, expected): + result = engine.execute(f"RETURN {expression} AS value") + assert result.to_pylist() == [{"value": expected}] + + +@pytest.mark.parametrize( + "expression", + [ + "left('x', -1)", + "right('x', null)", + "substring('x', -1)", + "substring('x', 0, null)", + "range(1, 3, 0)", + "range(1.5, 3)", + "abs('1')", + "abs(true)", + "abs(-9223372036854775807 - 1)", + "size(1)", + "reverse(1)", + "tail('abc')", + "last(2)", + "left('x', 1.5)", + "split('x', 1)", + "round(1, 2)", + "substring('x')", + "range(1)", + ], +) +def test_scalar_invalid_arguments(engine, expression): + with pytest.raises((ValueError, RuntimeError)): + engine.execute(f"RETURN {expression} AS value") + + +def test_scalar_arrow_columns_and_nulls(engine): + result = engine.execute( + "MATCH (p:Person) RETURN p.id AS id, " + "substring(p.text, 1, 2) AS sub, size(p.text) AS size, " + "round(p.number) AS rounded, reverse(p.items) AS reversed, " + "tail(p.pair) AS tail, last(p.pair) AS last ORDER BY id" + ) + assert result.to_pylist() == [ + { + "id": 1, + "sub": "猫😀", + "size": 4, + "rounded": -1.0, + "reversed": [3, None, 1], + "tail": [2], + "last": 2, + }, + { + "id": 2, + "sub": None, + "size": None, + "rounded": None, + "reversed": None, + "tail": None, + "last": None, + }, + { + "id": 3, + "sub": "", + "size": 0, + "rounded": 2.0, + "reversed": [], + "tail": [4], + "last": 4, + }, + ] + assert result.schema.field("rounded").type == pa.float64() + assert result.schema.field("size").type == pa.int64() + +def test_scalar_whitespace_and_nan(engine): + result = engine.execute( + "RETURN trim($text) AS text, sqrt(-1) AS root, round(-0.5) AS zero", + {"text": "\t\n 猫\u2003"}, + ).to_pylist()[0] + assert result["text"] == "猫" + assert math.isnan(result["root"]) + assert math.copysign(1, result["zero"]) == 1 -def test_unimplemented_scalar_function_errors() -> None: - cfg = GraphConfig.builder().node_table("Person", "name").build() - datasets = {"Person": pa.table({"name": ["Alice", "BOB", "CaSeY"]})} - query = CypherQuery( - "MATCH (p:Person) " - "RETURN p.name AS name, replace(p.name, 'A', 'a') AS replaced " - "ORDER BY name", - ).with_config(cfg) +def test_unimplemented_scalar_function_errors(engine): + with pytest.raises(ValueError, match="not_a_function.*is not implemented"): + engine.execute("RETURN not_a_function('x') AS value") - with pytest.raises(Exception) as excinfo: - query.execute(datasets) - assert "replace" in str(excinfo.value).lower() +def test_scalar_varying_arguments(engine): + result = engine.execute( + "MATCH (p:Person) RETURN p.id AS id, left(p.text, p.id) AS prefix, " + "range(0, p.id, 2) AS numbers ORDER BY id" + ) + assert result.to_pylist() == [ + {"id": 1, "prefix": "a", "numbers": [0]}, + {"id": 2, "prefix": None, "numbers": [0, 2]}, + {"id": 3, "prefix": "", "numbers": [0, 2]}, + ] From fbff9707abfc739fb5e2cc8d7df58306ba59848f Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:27:34 -0700 Subject: [PATCH 15/66] refactor(cypher): reuse templates for stable unaliased columns Source and runtime checks confirm literals and arithmetic already use the stable expr column name. Remove the unnecessary fallback and correct its regression expectation and docs. Fresh-versus-cached Arrow schema/result tests pass for unaliased scalar parameters. Part of #14. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/src/query.rs | 43 +++----------------------- docs/api-reference.md | 2 +- docs/lance.md | 5 +-- python/python/tests/test_plan_cache.py | 6 ++-- 4 files changed, 10 insertions(+), 46 deletions(-) diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 62c97dc..d8c3eb1 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -253,10 +253,7 @@ impl CypherQuery { .filter_map(|stage| stage.with_clause.limit.as_ref()), ) .any(|count| matches!(count, crate::ast::RowCount::Parameter(_))); - if parameter_counts - || self.scalar_parameters().is_none() - || self.has_parameter_named_columns() - { + if parameter_counts || self.scalar_parameters().is_none() { return Ok(false); } let (graph, datafusion) = self.build_logical_plans(catalog.clone(), true)?; @@ -274,38 +271,6 @@ impl CypherQuery { Ok(true) } - // Unaliased expressions use their bound values as column names. Reusing a - // placeholder-named projection would change the observable Arrow schema. - fn has_parameter_named_columns(&self) -> bool { - self.ast - .return_clause - .items - .iter() - .chain(self.ast.with_clause.iter().flat_map(|clause| &clause.items)) - .chain( - self.ast - .further_with_stages - .iter() - .flat_map(|stage| &stage.with_clause.items), - ) - .any(|item| { - if item.alias.is_some() { - return false; - } - let mut found = false; - crate::ast::transform_value(&mut item.expression.clone(), &mut |value| { - found |= matches!( - value, - crate::ast::ValueExpression::Parameter(_) - | crate::ast::ValueExpression::Literal( - crate::ast::PropertyValue::Parameter(_) - ) - ); - }); - found - }) - } - fn scalar_parameters(&self) -> Option> { use datafusion::common::ScalarValue; self.parameters @@ -1587,8 +1552,10 @@ mod tests { .await .unwrap(); let catalog: Arc = Arc::new(catalog); - let cached = template.prepare_cached_plan(catalog.clone()).unwrap(); - assert_eq!(cached, !template.has_parameter_named_columns(), "{text}"); + assert!( + template.prepare_cached_plan(catalog.clone()).unwrap(), + "{text}" + ); for value in values { let cached = template.clone().with_parameter("value", value.clone()); let fresh = CypherQuery::new(text) diff --git a/docs/api-reference.md b/docs/api-reference.md index 7764763..8d471e5 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -89,7 +89,7 @@ CypherEngine(config: GraphConfig, datasets: dict, use_csr: bool = False) Registers `datasets` once and reuses the catalog, DataFusion session and (with `use_csr=True`) CSR index across queries. Safe to share between threads. `execute()` and `execute_batches()` reuse up to 128 unoptimized query templates per snapshot, keyed by exact text and scalar parameter names/types. Values are -bound before optimization on every call. Containers, parameterized `SKIP`/`LIMIT`, value-dependent column names and unsupported template +bound before optimization on every call. Containers, parameterized `SKIP`/`LIMIT` and unsupported template shapes use ordinary planning. `explain()` always plans directly; successful `refresh()` clears the cache. See [Repeated queries](lance.md#repeated-queries) for the scope and limits. diff --git a/docs/lance.md b/docs/lance.md index 9aaec50..4e09671 100644 --- a/docs/lance.md +++ b/docs/lance.md @@ -91,10 +91,7 @@ rows and results are never cached. Query text and parameter names/types identify separate from non-null values. Least recently used templates are evicted at the limit. Lists, maps and parameterized `SKIP`/`LIMIT` counts use ordinary planning because their shape or value can -change the plan. Unaliased parameter expressions -also use ordinary planning to preserve their value-dependent column names; use `RETURN $value AS value` -when you want a stable name. Queries whose operators cannot produce a reusable template fall back to ordinary -planning. `explain()` plans the current values directly. A successful `refresh()` starts an empty cache for the +change the plan. Queries whose operators cannot produce a reusable template fall back to ordinary planning. `explain()` plans the current values directly. A successful `refresh()` starts an empty cache for the new snapshot, including a refresh whose datasets did not change; active readers retain their existing plans. These rules apply with and without CSR. They do not promise a speedup for a particular workload. diff --git a/python/python/tests/test_plan_cache.py b/python/python/tests/test_plan_cache.py index a0b2ea5..30aa404 100644 --- a/python/python/tests/test_plan_cache.py +++ b/python/python/tests/test_plan_cache.py @@ -51,12 +51,12 @@ def test_cached_parameter_types_and_container_fallback(engine): ] -def test_unaliased_parameters_keep_bound_column_names(engine): +def test_unaliased_parameters_keep_existing_column_names(engine): query = "MATCH (p:Person) WHERE p.id = 1 RETURN $value" first = engine.execute(query, {"value": 1}) second = engine.execute(query, {"value": 2}) - assert first.column_names == ["1"] - assert second.column_names == ["2"] + assert first.column_names == ["expr"] + assert second.column_names == ["expr"] assert first.column(0).to_pylist() == [1] assert second.column(0).to_pylist() == [2] From b11856218adf73ed16b34ae13f8d081276715f1b Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:32:45 -0700 Subject: [PATCH 16/66] test(cypher): keep unsupported-function coverage after replace support Use an explicitly unknown function in the rejection regression now that replace() is supported. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/tests/test_datafusion_pipeline.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/crates/lance-graph/tests/test_datafusion_pipeline.rs b/crates/lance-graph/tests/test_datafusion_pipeline.rs index 2b2041f..444c67f 100644 --- a/crates/lance-graph/tests/test_datafusion_pipeline.rs +++ b/crates/lance-graph/tests/test_datafusion_pipeline.rs @@ -4747,7 +4747,7 @@ async fn test_unimplemented_scalar_function_errors() { .unwrap(); let query = CypherQuery::new( - "MATCH (p:Person) RETURN p.name AS name, replace(p.name, 'A', 'a') AS replaced", + "MATCH (p:Person) RETURN p.name AS name, unsupported_function(p.name) AS value", ) .unwrap() .with_config(config); @@ -4758,10 +4758,13 @@ async fn test_unimplemented_scalar_function_errors() { let err = query .execute(datasets, Some(ExecutionStrategy::DataFusion)) .await - .expect_err("replace() should error until implemented"); + .expect_err("unknown functions should report an unsupported feature"); let message = err.to_string().to_lowercase(); - assert!(message.contains("replace"), "unexpected error: {err}"); + assert!( + message.contains("unsupported_function"), + "unexpected error: {err}" + ); assert!( message.contains("not implemented") || message.contains("unsupported"), "unexpected error: {err}" From dbef77dfece50eebcb595a19b293411db3f61d82 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:03:02 -0700 Subject: [PATCH 17/66] feat(query): execute read-only correlated CALL subqueries Address the read-only CALL portion of #20 with explicit imports and per-row streaming Apply execution. Preserve entity aliases, child aggregates and top-k semantics, and keep child variables scoped. Bypass snapshot templates for CALL and yield during long runs of empty children. Validation: workspace Rust tests, focused CALL regressions, 724 Python tests on the rebuilt extension, Clippy, executable docs, strict docs build, Ruff, rustfmt and typos 1.26.0. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/src/apply.rs | 439 ++++++++++++++++++ crates/lance-graph/src/ast.rs | 61 ++- crates/lance-graph/src/csr_expand.rs | 24 +- .../src/datafusion_planner/analysis.rs | 17 + .../datafusion_planner/builder/call_ops.rs | 78 ++++ .../src/datafusion_planner/builder/helpers.rs | 14 + .../src/datafusion_planner/builder/mod.rs | 11 + crates/lance-graph/src/lib.rs | 1 + crates/lance-graph/src/logical_plan.rs | 101 +++- .../lance-graph/src/parameter_substitution.rs | 1 + crates/lance-graph/src/parser.rs | 22 +- crates/lance-graph/src/query.rs | 44 +- crates/lance-graph/src/semantic.rs | 65 +++ docs/api-reference.md | 2 +- docs/cypher.md | 40 +- docs/execution.md | 14 +- python/python/tests/test_call_subqueries.py | 203 ++++++++ 17 files changed, 1114 insertions(+), 23 deletions(-) create mode 100644 crates/lance-graph/src/apply.rs create mode 100644 crates/lance-graph/src/datafusion_planner/builder/call_ops.rs create mode 100644 python/python/tests/test_call_subqueries.py diff --git a/crates/lance-graph/src/apply.rs b/crates/lance-graph/src/apply.rs new file mode 100644 index 0000000..e93c05d --- /dev/null +++ b/crates/lance-graph/src/apply.rs @@ -0,0 +1,439 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Read-only CALL: one independently bound child stream per outer row. + +use std::cmp::Ordering; +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +use arrow_array::RecordBatch; +use arrow_schema::SchemaRef; +use datafusion::common::{DFSchemaRef, Result, ScalarValue}; +use datafusion::execution::context::SessionState; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_expr::EquivalenceProperties; +use datafusion::physical_plan::execution_plan::EmissionType; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + execute_stream, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + Partitioning, PlanProperties, +}; +use futures::StreamExt; + +/// The child is deliberately unoptimized until its imports have concrete values. +/// It is not an ordinary input: planning it before binding would give global top-k +/// and aggregate semantics instead of evaluating them once per outer row. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub(crate) struct ApplyNode { + input: LogicalPlan, + inner: LogicalPlan, + /// (outer column name, private placeholder name). + imports: Vec<(String, String)>, + schema: DFSchemaRef, +} + +impl ApplyNode { + pub(crate) fn try_new( + input: LogicalPlan, + inner: LogicalPlan, + imports: Vec<(String, String)>, + ) -> Result { + let schema = Arc::new(input.schema().join(inner.schema())?); + Ok(Self { + input, + inner, + imports, + schema, + }) + } +} + +impl PartialOrd for ApplyNode { + fn partial_cmp(&self, other: &Self) -> Option { + self.input + .partial_cmp(&other.input) + .and_then(|order| match order { + Ordering::Equal => self.inner.partial_cmp(&other.inner), + order => Some(order), + }) + .filter(|order| *order != Ordering::Equal || self == other) + } +} + +impl UserDefinedLogicalNodeCore for ApplyNode { + fn name(&self) -> &str { + "Apply" + } + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + fn expressions(&self) -> Vec { + vec![] + } + fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "Apply: imports={:?}, child={}", + self.imports + .iter() + .map(|(name, _)| name) + .collect::>(), + self.inner.display_indent() + ) + } + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> Result { + Self::try_new(inputs.remove(0), self.inner.clone(), self.imports.clone()) + } +} + +pub(crate) struct ApplyExec { + input: Arc, + inner: LogicalPlan, + imports: Vec<(String, String)>, + state: Arc, + schema: SchemaRef, + properties: Arc, +} + +impl fmt::Debug for ApplyExec { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("ApplyExec") + .field("imports", &self.imports) + .finish_non_exhaustive() + } +} + +impl ApplyExec { + fn new( + input: Arc, + inner: LogicalPlan, + imports: Vec<(String, String)>, + state: Arc, + schema: SchemaRef, + ) -> Self { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + Partitioning::UnknownPartitioning(input.output_partitioning().partition_count()), + EmissionType::Incremental, + input.boundedness(), + )); + Self { + input, + inner, + imports, + state, + schema, + properties, + } + } +} + +impl DisplayAs for ApplyExec { + fn fmt_as(&self, _format: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "ApplyExec: per-row CALL, {} imported columns", + self.imports.len() + ) + } +} + +impl ExecutionPlan for ApplyExec { + fn name(&self) -> &str { + "ApplyExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self::new( + children.remove(0), + self.inner.clone(), + self.imports.clone(), + self.state.clone(), + self.schema.clone(), + ))) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let input = self.input.execute(partition, context.clone())?; + let indices = self + .imports + .iter() + .map(|(column, name)| Ok((input.schema().index_of(column)?, name.clone()))) + .collect::>>()?; + let state = ApplyStreamState { + input, + batch: None, + next_row: 0, + rows_since_yield: 0, + current: vec![], + child: None, + inner: self.inner.clone(), + imports: indices, + session: self.state.clone(), + context, + schema: self.schema.clone(), + }; + let stream = futures::stream::try_unfold(state, |mut state| async move { + loop { + if let Some(child) = &mut state.child { + if let Some(batch) = child.next().await { + let batch = batch?; + if batch.num_rows() == 0 { + continue; + } + let mut columns = state + .current + .iter() + .map(|value| value.to_array_of_size(batch.num_rows())) + .collect::>>()?; + columns.extend_from_slice(batch.columns()); + let output = RecordBatch::try_new(state.schema.clone(), columns)?; + return Ok(Some((output, state))); + } + state.child = None; + } + if state + .batch + .as_ref() + .is_none_or(|batch| state.next_row >= batch.num_rows()) + { + let Some(batch) = state.input.next().await else { + return Ok(None); + }; + state.batch = Some(batch?); + state.next_row = 0; + if state + .batch + .as_ref() + .is_some_and(|batch| batch.num_rows() == 0) + { + continue; + } + } + // Empty children can complete synchronously for many outer rows. + // Bound that work so cancellation and other tasks can be polled. + if state.rows_since_yield == 64 { + crate::csr_expand::cooperative_yield().await; + state.rows_since_yield = 0; + } + state.rows_since_yield += 1; + let batch = state.batch.as_ref().expect("a nonempty outer batch"); + state.current = batch + .columns() + .iter() + .map(|array| ScalarValue::try_from_array(array, state.next_row)) + .collect::>>()?; + state.next_row += 1; + let parameters: HashMap<_, _> = state + .imports + .iter() + .map(|(index, name)| (name.clone(), state.current[*index].clone())) + .collect(); + let bound = state.inner.clone().with_param_values(parameters)?; + let physical = state.session.create_physical_plan(&bound).await?; + state.child = Some(execute_stream(physical, state.context.clone())?); + } + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + stream, + ))) + } +} + +struct ApplyStreamState { + input: SendableRecordBatchStream, + batch: Option, + next_row: usize, + rows_since_yield: usize, + current: Vec, + child: Option, + inner: LogicalPlan, + imports: Vec<(usize, String)>, + session: Arc, + context: Arc, + schema: SchemaRef, +} + +pub(crate) fn plan_apply( + node: &ApplyNode, + inputs: &[Arc], + state: &SessionState, +) -> Arc { + Arc::new(ApplyExec::new( + inputs[0].clone(), + node.inner.clone(), + node.imports.clone(), + Arc::new(state.clone()), + Arc::new(node.schema.as_arrow().clone()), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{CypherQuery, GraphConfig}; + use arrow_array::{Int64Array, StringArray}; + + fn graph() -> (GraphConfig, HashMap) { + let config = GraphConfig::builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + .unwrap(); + let people = RecordBatch::try_from_iter(vec![ + ( + "id", + Arc::new(Int64Array::from(vec![10, 20, 30, 40])) as arrow_array::ArrayRef, + ), + ( + "name", + Arc::new(StringArray::from(vec!["A", "B", "C", "D"])), + ), + ("age", Arc::new(Int64Array::from(vec![20, 30, 40, 50]))), + ]) + .unwrap(); + let edges = RecordBatch::try_from_iter(vec![ + ( + "src", + Arc::new(Int64Array::from(vec![10, 10, 10, 20])) as arrow_array::ArrayRef, + ), + ("dst", Arc::new(Int64Array::from(vec![20, 20, 30, 40]))), + ]) + .unwrap(); + ( + config, + HashMap::from([("Person".into(), people), ("KNOWS".into(), edges)]), + ) + } + + async fn run(text: &str) -> crate::error::Result { + let (config, data) = graph(); + CypherQuery::new(text)? + .with_config(config) + .execute(data, None) + .await + } + + fn ints(batch: &RecordBatch, name: &str) -> Vec { + batch + .column_by_name(name) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + } + + #[tokio::test] + async fn call_per_row_top_k_and_entity_export() { + let result = run("MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) RETURN f AS friend ORDER BY friend.age DESC LIMIT 1 } RETURN p.id AS id, friend.id AS friend ORDER BY id").await.unwrap(); + assert_eq!(ints(&result, "id"), [10, 20]); + assert_eq!(ints(&result, "friend"), [30, 40]); + } + + #[tokio::test] + async fn call_aggregate_keeps_empty_child_and_outer_multiplicity() { + let result = run("MATCH (p:Person) UNWIND [1, 1] AS duplicate CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) RETURN count(f) AS n } RETURN p.id AS id, n ORDER BY id").await.unwrap(); + assert_eq!(ints(&result, "id"), [10, 10, 20, 20, 30, 30, 40, 40]); + assert_eq!(ints(&result, "n"), [3, 3, 1, 1, 0, 0, 0, 0]); + } + + #[tokio::test] + async fn call_nested_and_sibling_imports_have_separate_bindings() { + let (config, data) = graph(); + let result = CypherQuery::new("UNWIND [1, 2] AS x CALL { WITH x CALL { WITH x RETURN x + $__call_1_0 AS y } RETURN y + 1 AS z } CALL { WITH z RETURN z * 2 AS answer } RETURN x, answer ORDER BY x") + .unwrap().with_parameter("__call_1_0", 10).with_config(config).execute(data, None).await.unwrap(); + assert_eq!(ints(&result, "x"), [1, 2]); + assert_eq!(ints(&result, "answer"), [24, 26]); + } + + #[tokio::test] + async fn call_empty_outer_never_evaluates_child() { + let result = + run("MATCH (p:Person) WHERE p.id < 0 CALL { RETURN 1 / 0 AS invalid } RETURN invalid") + .await + .unwrap(); + assert_eq!(result.num_rows(), 0); + } + + #[tokio::test] + async fn call_inner_names_are_lexical_and_exports_explicit() { + let result = run("UNWIND [1, 2] AS x CALL { UNWIND [10, 20] AS x RETURN x AS y } RETURN x, y ORDER BY x, y").await.unwrap(); + assert_eq!(ints(&result, "x"), [1, 1, 2, 2]); + assert_eq!(ints(&result, "y"), [10, 20, 10, 20]); + for (query, message) in [ + ( + "UNWIND [1] AS x CALL { RETURN x AS y } RETURN y", + "Undefined variable", + ), + ( + "UNWIND [1] AS x CALL { RETURN 2 AS x } RETURN x", + "collides", + ), + ("CALL { RETURN 1 + 2 } RETURN 1", "need an alias"), + ( + "CALL { WITH missing RETURN missing AS x } RETURN x", + "undefined variable", + ), + ( + "UNWIND [1] AS x CALL { WITH x + 1 AS y RETURN y } RETURN y", + "importing WITH", + ), + ( + "CALL { UNWIND [1] AS local RETURN 2 AS y } RETURN local", + "Undefined variable", + ), + ] { + let error = run(query).await.unwrap_err().to_string(); + assert!(error.contains(message), "{query}: {error}"); + } + } + + #[test] + fn call_bypasses_snapshot_templates_in_every_reading_stage() { + let catalog = Arc::new(lance_graph_catalog::InMemoryCatalog::new()); + for text in [ + "CALL { RETURN 1 AS n } RETURN n", + "WITH 1 AS x CALL { WITH x RETURN x AS n } RETURN n", + "WITH 1 AS x WITH x AS y CALL { WITH y RETURN y AS n } RETURN n", + ] { + let mut query = CypherQuery::new(text).unwrap(); + assert!(!query.prepare_cached_plan(catalog.clone()).unwrap()); + } + } + + #[test] + fn call_parser_requires_returning_read_only_body() { + for query in [ + "CALL db.labels() RETURN 1", + "CALL { CREATE (n:Person) RETURN n } RETURN n", + "CALL { MATCH (n:Person) } RETURN n", + "CALL { RETURN 1 AS x } YIELD x RETURN x", + ] { + assert!(CypherQuery::new(query).is_err(), "{query}"); + } + } +} diff --git a/crates/lance-graph/src/ast.rs b/crates/lance-graph/src/ast.rs index 7d4e183..24b0304 100644 --- a/crates/lance-graph/src/ast.rs +++ b/crates/lance-graph/src/ast.rs @@ -59,11 +59,29 @@ pub struct CypherQuery { } impl CypherQuery { + pub(crate) fn all_reading_clauses(&self) -> impl Iterator { + self.reading_clauses + .iter() + .chain(&self.post_with_reading_clauses) + .chain( + self.further_with_stages + .iter() + .flat_map(|stage| &stage.reading_clauses), + ) + } + /// Extract all node labels referenced in the query pub fn get_node_labels(&self) -> Vec { let mut labels = Vec::new(); // Iterate all match clauses directly - for clause in &self.reading_clauses { + for clause in self.all_reading_clauses() { + if let ReadingClause::Call(inner) = clause { + for label in inner.get_node_labels() { + if !labels.contains(&label) { + labels.push(label); + } + } + } if let ReadingClause::Match(match_clause) = clause { for pattern in &match_clause.patterns { match pattern { @@ -98,7 +116,14 @@ impl CypherQuery { /// Extract all relationship types referenced in the query pub fn get_relationship_types(&self) -> Vec { let mut types = Vec::new(); - for clause in &self.reading_clauses { + for clause in self.all_reading_clauses() { + if let ReadingClause::Call(inner) = clause { + for kind in inner.get_relationship_types() { + if !types.contains(&kind) { + types.push(kind); + } + } + } if let ReadingClause::Match(match_clause) = clause { for pattern in &match_clause.patterns { self.collect_relationship_types_from_pattern(pattern, &mut types); @@ -130,6 +155,38 @@ impl CypherQuery { pub enum ReadingClause { Match(MatchClause), Unwind(UnwindClause), + /// Read-only subquery, with an optional explicit importing WITH. + Call(Box), +} + +impl CypherQuery { + /// Names explicitly imported by a leading plain WITH in a CALL body. + pub(crate) fn call_imports(&self) -> crate::error::Result> { + let Some(clause) = self + .with_clause + .as_ref() + .filter(|_| self.reading_clauses.is_empty()) + else { + return Ok(vec![]); + }; + if clause.distinct + || clause.order_by.is_some() + || clause.limit.is_some() + || clause.where_clause.is_some() + { + return Err(crate::error::GraphError::PlanError { + message: "CALL importing WITH must contain only unaliased variable names; use a second WITH for expressions or filters".into(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + clause.items.iter().map(|item| match (&item.expression, &item.alias) { + (ValueExpression::Variable(name), None) => Ok(name.to_lowercase()), + _ => Err(crate::error::GraphError::PlanError { + message: "CALL importing WITH must contain only unaliased variable names; use a second WITH for expressions".into(), + location: snafu::Location::new(file!(), line!(), column!()), + }), + }).collect() + } } /// A MATCH clause containing graph patterns diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index b8b88f9..010023e 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -61,6 +61,21 @@ use futures::StreamExt; use lance::dataset::{Dataset, ProjectionRequest}; use lance_table::format::Fragment; +/// Yield after bounded CPU work without requiring a particular async runtime. +pub(crate) async fn cooperative_yield() { + let mut yielded = false; + std::future::poll_fn(|cx| { + if yielded { + std::task::Poll::Ready(()) + } else { + yielded = true; + cx.waker().wake_by_ref(); + std::task::Poll::Pending + } + }) + .await +} + /// Lance's row-address column (`lance_core::ROW_ADDR`, not re-exported by `lance`). const ROW_ADDR: &str = "_rowaddr"; @@ -1373,6 +1388,13 @@ impl ExtensionPlanner for CsrExtensionPlanner { physical_inputs: &[Arc], session_state: &SessionState, ) -> datafusion::common::Result>> { + if let Some(node) = node.as_any().downcast_ref::() { + return Ok(Some(crate::apply::plan_apply( + node, + physical_inputs, + session_state, + ))); + } if let Some(node) = node .as_any() .downcast_ref::() @@ -1402,7 +1424,7 @@ impl ExtensionPlanner for CsrExtensionPlanner { } /// Query planner that adds [`CsrExtensionPlanner`]; install it on any -/// `SessionContext` that executes plans built with an [`AdjacencyCatalog`]. +/// `SessionContext` that executes plans built with an [`AdjacencyCatalog`] or CALL subqueries. #[derive(Debug, Default)] pub struct CsrQueryPlanner; diff --git a/crates/lance-graph/src/datafusion_planner/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index 3329517..aa9a967 100644 --- a/crates/lance-graph/src/datafusion_planner/analysis.rs +++ b/crates/lance-graph/src/datafusion_planner/analysis.rs @@ -39,6 +39,7 @@ 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) call_input: Option, pub(crate) comprehension_elements: Vec, } @@ -48,6 +49,7 @@ impl<'a> PlanningContext<'a> { analysis, relationship_instance_idx: HashMap::new(), comprehension_elements: Vec::new(), + call_input: None, } } @@ -256,6 +258,21 @@ fn analyze_operator( } analyze_operator(element, analysis, rel_counter)?; } + LogicalOperator::Apply { + input, + inner, + labels, + .. + } => { + analyze_operator(input, analysis, rel_counter)?; + analysis + .required_datasets + .extend(analyze(inner)?.required_datasets); + analysis.var_to_label.extend(labels.clone()); + } + LogicalOperator::CallInput { labels } => { + analysis.var_to_label.extend(labels.clone()); + } LogicalOperator::ComprehensionElement | LogicalOperator::SingleRow => {} } Ok(()) diff --git a/crates/lance-graph/src/datafusion_planner/builder/call_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/call_ops.rs new file mode 100644 index 0000000..15d9a77 --- /dev/null +++ b/crates/lance-graph/src/datafusion_planner/builder/call_ops.rs @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +use super::super::{analysis, DataFusionPlanner, PlanningContext}; +use crate::apply::ApplyNode; +use crate::error::Result; +use crate::logical_plan::LogicalOperator; +use datafusion::logical_expr::{ + col, expr::Placeholder, Expr, Extension, LogicalPlan, LogicalPlanBuilder, +}; +use std::sync::Arc; + +impl DataFusionPlanner { + pub(crate) fn build_apply( + &self, + ctx: &mut PlanningContext, + input: &LogicalOperator, + inner: &LogicalOperator, + names: &[String], + id: usize, + ) -> Result { + let input = self.build_operator(ctx, input)?; + // A scalar named x wins over entity-shaped names such as x__property. + let entity_prefixes: Vec<_> = names + .iter() + .filter(|name| { + !input + .schema() + .fields() + .iter() + .any(|field| field.name() == *name) + }) + .map(|name| format!("{name}__")) + .collect(); + let mut imports = Vec::new(); + let mut values = Vec::new(); + for (index, (_, field)) in input.schema().iter().enumerate() { + if names.contains(field.name()) + || entity_prefixes + .iter() + .any(|prefix| field.name().starts_with(prefix)) + { + // User parameters have already become literals. The monotonically assigned + // CALL id also separates nested and sibling import placeholders. + let parameter = format!("__call_{id}_{index}"); + imports.push((field.name().clone(), parameter.clone())); + values.push(Expr::Placeholder(Placeholder::new_with_field( + format!("${parameter}"), + Some(field.clone()), + ))); + } + } + // DataFusion reconstructs VALUES with column1, column2, ... when binding + // placeholders. Keep those names and restore the imported names in a projection. + let seed = + if imports.is_empty() { + LogicalPlanBuilder::empty(true).build() + } else { + LogicalPlanBuilder::values(vec![values]).and_then(|builder| { + builder + .project(imports.iter().enumerate().map(|(index, (name, _))| { + col(format!("column{}", index + 1)).alias(name) + }))? + .build() + }) + } + .map_err(|error| self.plan_error("CALL imports", error))?; + let analysis = analysis::analyze(inner)?; + let mut child_ctx = PlanningContext::new(&analysis); + child_ctx.call_input = Some(seed); + let child = self.build_operator(&mut child_ctx, inner)?; + let node = ApplyNode::try_new(input, child, imports) + .map_err(|error| self.plan_error("CALL output", error))?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(node), + })) + } +} diff --git a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs index 34647d5..0df0bf7 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs @@ -112,6 +112,11 @@ impl DataFusionPlanner { LogicalOperator::ListComprehension { input, .. } => { Self::collect_variables(input, vars) } + LogicalOperator::Apply { input, labels, .. } => { + Self::collect_variables(input, vars); + vars.extend(labels.keys().cloned()); + } + LogicalOperator::CallInput { labels } => vars.extend(labels.keys().cloned()), LogicalOperator::ComprehensionElement | LogicalOperator::SingleRow => {} } } @@ -176,6 +181,15 @@ impl DataFusionPlanner { LogicalOperator::ShortestPath { input, .. } => Self::collect_labels(input, labels), LogicalOperator::Union { .. } => {} LogicalOperator::ListComprehension { input, .. } => Self::collect_labels(input, labels), + LogicalOperator::Apply { + input, + labels: exports, + .. + } => { + Self::collect_labels(input, labels); + labels.extend(exports.clone()); + } + LogicalOperator::CallInput { labels: imports } => labels.extend(imports.clone()), LogicalOperator::ComprehensionElement | LogicalOperator::SingleRow => {} } } diff --git a/crates/lance-graph/src/datafusion_planner/builder/mod.rs b/crates/lance-graph/src/datafusion_planner/builder/mod.rs index 5f07235..2e6460a 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 call_ops; mod comprehension_ops; mod expand_ops; mod helpers; @@ -34,6 +35,16 @@ impl DataFusionPlanner { op: &LogicalOperator, ) -> Result { match op { + LogicalOperator::Apply { + input, + inner, + imports, + id, + .. + } => self.build_apply(ctx, input, inner, imports, *id), + LogicalOperator::CallInput { .. } => ctx.call_input.clone().ok_or_else(|| { + self.plan_error("CALL input", "import rows must be built inside CALL") + }), LogicalOperator::SingleRow => LogicalPlanBuilder::empty(true) .build() .map_err(|e| self.plan_error("Failed to build singleton row", e)), diff --git a/crates/lance-graph/src/lib.rs b/crates/lance-graph/src/lib.rs index 37083e7..f7746e0 100644 --- a/crates/lance-graph/src/lib.rs +++ b/crates/lance-graph/src/lib.rs @@ -35,6 +35,7 @@ //! # } //! ``` +mod apply; pub mod ast; pub mod case_insensitive; pub mod config; diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index e8ccc63..ad7e8d1 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -177,6 +177,17 @@ pub enum LogicalOperator { count_only: bool, }, + /// Evaluate a read-only subquery once per input row, extending it by each result. + Apply { + input: Box, + inner: Box, + imports: Vec, + labels: HashMap, + id: usize, + }, + /// The single imported row of a CALL body; its values are bound at execution. + CallInput { labels: HashMap }, + /// The element rows of the enclosing [`LogicalOperator::ListComprehension`]. ComprehensionElement, } @@ -221,6 +232,7 @@ pub struct LogicalPlanner<'a> { node_lists: HashMap, /// List comprehensions planned so far (output columns `__list_`). comprehensions: usize, + calls: 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 @@ -238,6 +250,7 @@ impl<'a> LogicalPlanner<'a> { exists_marks: 0, node_lists: HashMap::new(), comprehensions: 0, + calls: 0, path_labels: HashMap::new(), seed_source: None, } @@ -316,9 +329,18 @@ impl<'a> LogicalPlanner<'a> { /// Convert a Cypher AST to a logical plan pub fn plan(&mut self, query: &CypherQuery) -> Result { + self.plan_query(query, None, false) + } + + fn plan_query( + &mut self, + query: &CypherQuery, + input: Option, + preserve_entities: bool, + ) -> Result { // Plan main MATCH clauses; the WHERE clause guides where traversal starts. let where_expr = query.where_clause.as_ref().map(|w| &w.expression); - let mut plan = self.plan_reading_clauses(None, &query.reading_clauses, where_expr)?; + let mut plan = self.plan_reading_clauses(input, &query.reading_clauses, where_expr)?; // Apply WHERE clause if present (before WITH) if let Some(where_clause) = &query.where_clause { @@ -352,7 +374,20 @@ impl<'a> LogicalPlanner<'a> { } // Apply RETURN clause - plan = self.plan_return_clause(&query.return_clause, plan)?; + plan = if preserve_entities { + self.plan_with_clause( + &WithClause { + distinct: query.return_clause.distinct, + items: query.return_clause.items.clone(), + order_by: None, + limit: None, + where_clause: None, + }, + plan, + )? + } else { + self.plan_return_clause(&query.return_clause, plan)? + }; // Apply ORDER BY if present if let Some(order_by) = &query.order_by { @@ -411,6 +446,7 @@ impl<'a> LogicalPlanner<'a> { filter: Option<&BooleanExpression>, ) -> Result { match clause { + ReadingClause::Call(query) => self.plan_call(base, query), ReadingClause::Match(match_clause) if match_clause.optional => { self.plan_optional_match(base, match_clause, None) } @@ -427,6 +463,56 @@ impl<'a> LogicalPlanner<'a> { } } + fn plan_call( + &mut self, + input: Option, + query: &CypherQuery, + ) -> Result { + let imports = query.call_imports()?; + self.calls += 1; + let id = self.calls; + let mut inner = Self::new(self.config).with_fan_out(self.fan_out); + inner.calls = self.calls; + for name in &imports { + if let Some(label) = self.variables.get(name) { + inner.variables.insert(name.clone(), label.clone()); + } + if let Some(label) = self.node_lists.get(name) { + inner.node_lists.insert(name.clone(), label.clone()); + } + if let Some(label) = self.path_labels.get(name) { + inner.path_labels.insert(name.clone(), label.clone()); + } + } + let input_labels = inner.variables.clone(); + let plan = inner.plan_query( + query, + Some(LogicalOperator::CallInput { + labels: input_labels, + }), + true, + )?; + self.calls = inner.calls; + let labels = inner.variables.clone(); + self.variables.extend(inner.variables); + self.node_lists.extend(inner.node_lists); + for item in &query.return_clause.items { + if let ValueExpression::Variable(name) = &item.expression { + if let Some(label) = inner.path_labels.get(name) { + self.path_labels + .insert(item.alias.as_ref().unwrap_or(name).clone(), label.clone()); + } + } + } + Ok(LogicalOperator::Apply { + input: Box::new(input.unwrap_or(LogicalOperator::SingleRow)), + inner: Box::new(plan), + imports, + labels, + id, + }) + } + /// Plan an UNWIND clause fn plan_unwind_clause_with_base( &mut self, @@ -1287,10 +1373,13 @@ impl<'a> LogicalPlanner<'a> { #[allow(clippy::only_used_in_recursion)] fn extract_variable_from_plan(&self, plan: &LogicalOperator) -> Result { match plan { - LogicalOperator::SingleRow => Err(GraphError::PlanError { - message: "a singleton row binds no pattern variable".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - }), + LogicalOperator::Apply { input, .. } => self.extract_variable_from_plan(input), + LogicalOperator::CallInput { .. } | LogicalOperator::SingleRow => { + Err(GraphError::PlanError { + message: "a singleton row binds no pattern variable".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }) + } LogicalOperator::ScanByLabel { variable, .. } => Ok(variable.clone()), LogicalOperator::Unwind { alias, .. } => Ok(alias.clone()), LogicalOperator::Expand { diff --git a/crates/lance-graph/src/parameter_substitution.rs b/crates/lance-graph/src/parameter_substitution.rs index d0a0df4..ddb8a88 100644 --- a/crates/lance-graph/src/parameter_substitution.rs +++ b/crates/lance-graph/src/parameter_substitution.rs @@ -77,6 +77,7 @@ fn substitute_in_reading_clause( parameters: &HashMap, ) -> Result<()> { match clause { + ReadingClause::Call(query) => substitute_parameters(query, parameters)?, ReadingClause::Match(match_clause) => { for pattern in &mut match_clause.patterns { substitute_in_graph_pattern(pattern, parameters)?; diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index 260240b..5634a55 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -36,6 +36,11 @@ pub fn parse_cypher_query(input: &str) -> Result { } let mut query = query; + normalize_inline_properties(&mut query); + Ok(query) +} + +fn normalize_inline_properties(query: &mut CypherQuery) { let stages = query .further_with_stages .iter_mut() @@ -47,11 +52,12 @@ pub fn parse_cypher_query(input: &str) -> Result { .chain(&mut query.post_with_reading_clauses) .chain(stages) { - if let ReadingClause::Match(m) = clause { - move_variable_properties_to_where(m, &mut anonymous); + match clause { + ReadingClause::Match(m) => move_variable_properties_to_where(m, &mut anonymous), + ReadingClause::Call(inner) => normalize_inline_properties(inner), + ReadingClause::Unwind(_) => {} } } - Ok(query) } /// Inline node properties whose value is a variable or property from an earlier clause @@ -153,9 +159,17 @@ fn reading_clause(input: &str) -> IResult<&str, ReadingClause> { alt(( map(match_clause, ReadingClause::Match), map(unwind_clause, ReadingClause::Unwind), + map(call_clause, |query| ReadingClause::Call(Box::new(query))), ))(input) } +fn call_clause(input: &str) -> IResult<&str, CypherQuery> { + let (input, _) = multispace0(input)?; + let (input, _) = tag_no_case("CALL")(input)?; + let (input, _) = multispace0(input)?; + delimited(char('{'), cypher_query, preceded(multispace0, char('}')))(input) +} + // Parse a MATCH clause fn match_clause(input: &str) -> IResult<&str, MatchClause> { let (input, _) = multispace0(input)?; @@ -2499,7 +2513,7 @@ mod tests { .iter() .map(|c| match c { ReadingClause::Match(m) => m, - ReadingClause::Unwind(_) => unreachable!(), + ReadingClause::Unwind(_) | ReadingClause::Call(_) => unreachable!(), }) .collect(); assert_eq!(clauses.len(), 3); diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index d8c3eb1..fc9b47a 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -253,7 +253,14 @@ impl CypherQuery { .filter_map(|stage| stage.with_clause.limit.as_ref()), ) .any(|count| matches!(count, crate::ast::RowCount::Parameter(_))); - if parameter_counts || self.scalar_parameters().is_none() { + if parameter_counts + || self.scalar_parameters().is_none() + || self + .ast + .all_reading_clauses() + .any(|clause| matches!(clause, ReadingClause::Call(_))) + { + // CALL owns an inner template that is bound independently per outer row. return Ok(false); } let (graph, datafusion) = self.build_logical_plans(catalog.clone(), true)?; @@ -771,7 +778,12 @@ impl CypherQuery { } // Create session context and catalog - let ctx = SessionContext::new(); + let ctx = SessionContext::new_with_state( + datafusion::execution::session_state::SessionStateBuilder::new() + .with_default_features() + .with_query_planner(Arc::new(crate::csr_expand::CsrQueryPlanner)) + .build(), + ); let mut catalog = InMemoryCatalog::new(); // Register all datasets as tables @@ -851,7 +863,12 @@ impl CypherQuery { }); } - let ctx = SessionContext::new(); + let ctx = SessionContext::new_with_state( + datafusion::execution::session_state::SessionStateBuilder::new() + .with_default_features() + .with_query_planner(Arc::new(crate::csr_expand::CsrQueryPlanner)) + .build(), + ); let mut catalog = InMemoryCatalog::new(); let mut providers: HashMap> = HashMap::new(); @@ -1172,7 +1189,10 @@ impl CypherQuery { pub fn referenced_node_labels(&self) -> Vec { let mut labels = Vec::new(); - for clause in &self.ast.reading_clauses { + for clause in self.ast.all_reading_clauses() { + if let ReadingClause::Call(inner) = clause { + labels.extend(inner.get_node_labels()); + } if let ReadingClause::Match(match_clause) = clause { for pattern in &match_clause.patterns { self.collect_node_labels_from_pattern(pattern, &mut labels); @@ -1189,7 +1209,10 @@ impl CypherQuery { pub fn referenced_relationship_types(&self) -> Vec { let mut types = Vec::new(); - for clause in &self.ast.reading_clauses { + for clause in self.ast.all_reading_clauses() { + if let ReadingClause::Call(inner) = clause { + types.extend(inner.get_relationship_types()); + } if let ReadingClause::Match(match_clause) = clause { for pattern in &match_clause.patterns { self.collect_relationship_types_from_pattern(pattern, &mut types); @@ -1206,8 +1229,17 @@ impl CypherQuery { pub fn variables(&self) -> Vec { let mut variables = Vec::new(); - for clause in &self.ast.reading_clauses { + for clause in self.ast.all_reading_clauses() { match clause { + ReadingClause::Call(inner) => { + variables.extend(inner.call_imports().unwrap_or_default()); + variables.extend(inner.return_clause.items.iter().filter_map(|item| { + item.alias.clone().or_else(|| match &item.expression { + crate::ast::ValueExpression::Variable(name) => Some(name.clone()), + _ => None, + }) + })); + } ReadingClause::Match(match_clause) => { for pattern in &match_clause.patterns { self.collect_variables_from_pattern(pattern, &mut variables); diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index a418d7d..507e544 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -109,6 +109,11 @@ impl SemanticAnalyzer { self.current_scope = ScopeType::Match; for clause in &analyzed_query.reading_clauses { match clause { + ReadingClause::Call(query) => { + if let Err(error) = self.analyze_call(query) { + errors.push(format!("CALL clause error: {error}")); + } + } ReadingClause::Match(match_clause) => { if let Err(e) = self.analyze_match_clause(match_clause) { errors.push(format!("MATCH clause error: {}", e)); @@ -142,6 +147,11 @@ impl SemanticAnalyzer { self.current_scope = ScopeType::Match; for clause in &analyzed_query.post_with_reading_clauses { match clause { + ReadingClause::Call(query) => { + if let Err(error) = self.analyze_call(query) { + errors.push(format!("CALL clause error: {error}")); + } + } ReadingClause::Match(match_clause) => { if let Err(e) = self.analyze_match_clause(match_clause) { errors.push(format!("Post-WITH MATCH clause error: {}", e)); @@ -172,6 +182,7 @@ impl SemanticAnalyzer { self.current_scope = ScopeType::Match; for clause in &stage.reading_clauses { let result = match clause { + ReadingClause::Call(query) => self.analyze_call(query), ReadingClause::Match(match_clause) => self.analyze_match_clause(match_clause), ReadingClause::Unwind(unwind_clause) => { self.analyze_unwind_clause(unwind_clause) @@ -217,6 +228,60 @@ impl SemanticAnalyzer { }) } + fn analyze_call(&mut self, query: &CypherQuery) -> Result<()> { + let error = |message| GraphError::PlanError { + message, + location: snafu::Location::new(file!(), line!(), column!()), + }; + let mut inner = Self::new(self.config.clone()); + inner.catalog = self.catalog.clone(); + for name in query.call_imports()? { + let variable = self + .variables + .get(&name) + .ok_or_else(|| error(format!("CALL imports undefined variable '{name}'")))?; + inner.variables.insert(name.clone(), variable.clone()); + if let Some(value_type) = self.expression_types.get(&name) { + inner.expression_types.insert(name, value_type.clone()); + } + } + // Outer analysis already substituted all user parameters recursively. + let analyzed = inner.analyze(query, &HashMap::new())?; + if !analyzed.errors.is_empty() { + return Err(error(analyzed.errors.join("\n"))); + } + let mut exports = HashSet::new(); + for item in &query.return_clause.items { + let name = match (&item.alias, &item.expression) { + (Some(alias), _) => alias.to_lowercase(), + (None, ValueExpression::Variable(name)) => name.to_lowercase(), + _ => return Err(error("CALL RETURN expressions need an alias".into())), + }; + if self.variables.contains_key(&name) || !exports.insert(name.clone()) { + return Err(error(format!( + "CALL export '{name}' collides with an existing variable" + ))); + } + let value_type = inner.expression_type(&item.expression); + let (variable_type, labels) = match &value_type { + ExpressionType::Entity(kind, labels) => (kind.clone(), labels.clone()), + _ => (VariableType::Property, vec![]), + }; + self.variables.insert( + name.clone(), + VariableInfo { + name: name.clone(), + variable_type, + labels, + properties: HashSet::new(), + defined_in: self.current_scope.clone(), + }, + ); + self.expression_types.insert(name, value_type); + } + Ok(()) + } + /// Analyze MATCH clause and discover variables fn analyze_match_clause(&mut self, match_clause: &MatchClause) -> Result<()> { for pattern in &match_clause.patterns { diff --git a/docs/api-reference.md b/docs/api-reference.md index 8d471e5..cb5a38b 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -89,7 +89,7 @@ CypherEngine(config: GraphConfig, datasets: dict, use_csr: bool = False) Registers `datasets` once and reuses the catalog, DataFusion session and (with `use_csr=True`) CSR index across queries. Safe to share between threads. `execute()` and `execute_batches()` reuse up to 128 unoptimized query templates per snapshot, keyed by exact text and scalar parameter names/types. Values are -bound before optimization on every call. Containers, parameterized `SKIP`/`LIMIT` and unsupported template +bound before optimization on every call. Containers, `CALL` subqueries, parameterized `SKIP`/`LIMIT` and unsupported template shapes use ordinary planning. `explain()` always plans directly; successful `refresh()` clears the cache. See [Repeated queries](lance.md#repeated-queries) for the scope and limits. diff --git a/docs/cypher.md b/docs/cypher.md index 36cf70e..2eec634 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -190,12 +190,13 @@ read-only subset of openCypher, enough to run all 14 (IC14 without `reduce()`). | `WHERE` | After `MATCH`, `OPTIONAL MATCH` (filters inside the optional part) or `WITH`. | | `WITH [DISTINCT]` | With `ORDER BY`, `LIMIT` and `WHERE`; chain as many as you need. No `SKIP` inside `WITH`. | | `UNWIND list AS x` | | +| `CALL { ... RETURN ... }` | Read-only subquery. Import outer variables with a leading plain `WITH`; see below. | | `RETURN [DISTINCT]` | With `AS` aliases. Without reading clauses, expressions run on one implicit row. | | `ORDER BY`, `SKIP`, `LIMIT` | `SKIP` and `LIMIT` take nonnegative integer literals or parameters. | !!! failure "Not supported" - Writes (`CREATE`, `MERGE`, `SET`, `DELETE`, `REMOVE`), `UNION`, `CALL`, and `FOREACH`. + Writes (`CREATE`, `MERGE`, `SET`, `DELETE`, `REMOVE`), `UNION`, procedure calls (`CALL name(...)` / `YIELD`), and `FOREACH`. Integer list literals retain signed 64-bit values: `RETURN [1, 2, 3] AS ids` returns integers, and `UNWIND [1131, 19791209300775] AS id RETURN id` preserves both IDs exactly. Numeric literals with a decimal @@ -203,6 +204,43 @@ point or exponent, such as `[0.1, 1e-3]`, form float32 vectors. Whole-number lis vectors; conversion to float32 happens inside the vector operation. A configured engine is still required for standalone `RETURN` queries. +### Read-only subqueries + +Use `CALL { ... }` when a match, aggregate or limit should run separately for each incoming row: + +```python +print(engine.execute(""" + MATCH (p:Person) + CALL { + WITH p + MATCH (p)-[:KNOWS]->(f:Person) + RETURN f AS friend ORDER BY friend.age DESC LIMIT 1 + } + RETURN p.name AS person, friend.name AS friend ORDER BY person +""").to_pylist()) +# [{'person': 'Alice', 'friend': 'Bob'}, {'person': 'Bob', 'friend': 'David'}, +# {'person': 'Carol', 'friend': 'David'}] +``` + +The leading importing `WITH` accepts only existing variable names, without aliases, expressions, +`DISTINCT`, filters or pagination. Use a second `WITH` to transform or filter them. A body without an +importing `WITH` has its own independent scope. Calls can nest and can appear after `WITH` or `UNWIND`. + +Every body must end with `RETURN`. Give expressions an `AS` alias; bare variables keep their names. +Returned names must be distinct and must not already exist outside the call. Returning a node or +relationship retains its properties, so later clauses can use the returned entity. Names that the body +only uses internally do not become visible outside it. + +Each child result row is combined with its incoming row. Duplicate results remain duplicates. An empty +child removes that incoming row; an aggregate such as `count()` without grouping still returns a row +with zero. An empty outer input never executes the child. `ORDER BY` and `LIMIT` inside the body apply +separately to each incoming row. + +CALL currently binds, optimizes and physically plans a child for every incoming row. It streams the +child's batches before advancing that row, but operators such as sorting and aggregation can still +buffer their inputs. CALL queries bypass the query-template cache. For large outer inputs, measure the +query on representative data before choosing this form. Procedures, `YIELD`, and writes are unsupported. + ### Patterns ```cypher diff --git a/docs/execution.md b/docs/execution.md index f96bd1a..b7969b5 100644 --- a/docs/execution.md +++ b/docs/execution.md @@ -23,15 +23,25 @@ Source paths below are in `crates/lance-graph/src/`. only the order of work is. 4. **Plan the graph** (`logical_plan.rs`). Graph operators describe the query: `ScanByLabel`, `Expand` for a hop, `VariableLengthExpand` (unrolled into one branch per length), `OptionalMatch` as a left join, - `ShortestPath`. + `ShortestPath`, and `Apply` for a read-only CALL subquery. 5. **Plan for DataFusion** (`datafusion_planner/`). Each graph operator becomes DataFusion relational operators. A hop becomes either two hash joins or, when the engine has a [CSR](csr.md) for it, a `CsrExpand` extension node. 6. **Execute.** DataFusion optimizes and runs the physical plan, vectorized and multi-threaded across cores, - with `CsrExpandExec` and the shortest-path search as extension operators. + with `CsrExpandExec`, the shortest-path search and `ApplyExec` as extension operators. `engine.explain(query)` prints the plans from steps 4 to 6. +A CALL body is kept as an unoptimized child template. For each incoming row, `ApplyExec` binds its +explicit imports, plans the child, and streams that child's batches. Child `LIMIT` and aggregates +therefore apply separately to each incoming row. The operator retains one outer batch and one child +stream per partition; its children may have their own buffering requirements. Dropping the result +stream drops both owned streams. CALL queries use ordinary planning instead of the engine's template +cache. + +Rust callers supplying their own `SessionContext` must install `CsrQueryPlanner` to execute CALL or +other graph extension nodes. Contexts created by `CypherQuery` and `CypherEngine` install it themselves. + ## Joins versus CSR expansion ```mermaid diff --git a/python/python/tests/test_call_subqueries.py b/python/python/tests/test_call_subqueries.py new file mode 100644 index 0000000..592d297 --- /dev/null +++ b/python/python/tests/test_call_subqueries.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Read-only CALL uses explicit imports and evaluates its body per outer row.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, CypherQuery, GraphConfig + + +@pytest.fixture(params=[False, True]) +def graph(request): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + data = { + "Person": pa.table( + { + "id": [10, 20, 30, 40], + "name": ["A", "B", "C", "D"], + "age": [20, 30, 40, 50], + } + ), + "KNOWS": pa.table({"src": [10, 10, 10, 20], "dst": [20, 20, 30, 40]}), + } + return CypherEngine(config, data, use_csr=request.param), config, data + + +def test_call_top_k_preserves_node_alias_and_properties(graph): + engine, config, data = graph + query = ( + "MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) " + "RETURN f AS friend ORDER BY friend.age DESC LIMIT 1 } " + "RETURN p.id AS id, friend.name AS name ORDER BY id" + ) + expected = [{"id": 10, "name": "C"}, {"id": 20, "name": "D"}] + assert engine.execute(query).to_pylist() == expected + assert engine.execute_batches(query).read_all().to_pylist() == expected + assert CypherQuery(query).with_config(config).execute(data).to_pylist() == expected + assert "Apply" in engine.explain(query) + + +def test_call_aggregates_and_duplicates(graph): + engine, _, _ = graph + query = ( + "MATCH (p:Person) UNWIND [1, 1] AS duplicate CALL { WITH p " + "MATCH (p)-[:KNOWS]->(f:Person) RETURN count(f) AS n } " + "RETURN p.id AS id, n ORDER BY id" + ) + assert engine.execute(query).to_pylist() == [ + {"id": person, "n": count} + for person, count in [(10, 3), (20, 1), (30, 0), (40, 0)] + for _ in range(2) + ] + query = ( + "MATCH (p:Person) WHERE p.id = 10 CALL { WITH p " + "MATCH (p)-[:KNOWS]->(f:Person) RETURN f.id AS id } RETURN id ORDER BY id" + ) + assert engine.execute(query).to_pylist() == [{"id": 20}, {"id": 20}, {"id": 30}] + + +def test_call_nested_parameters_are_fresh_and_do_not_collide(graph): + engine, _, _ = graph + query = ( + "UNWIND [1, 2] AS x CALL { WITH x CALL { WITH x " + "RETURN x + $__call_1_0 AS y } RETURN y + 1 AS z } " + "CALL { WITH z RETURN z * 2 AS answer } RETURN x, answer ORDER BY x" + ) + for value in [10, 20, 10]: + assert engine.execute(query, {"__call_1_0": value}).to_pylist() == [ + {"x": x, "answer": (x + value + 1) * 2} for x in [1, 2] + ] + + +def test_call_nullable_scalar_import_and_second_with(graph): + engine, _, _ = graph + query = ( + "UNWIND [null, 2] AS x CALL { WITH x WITH coalesce(x, 7) AS value " + "RETURN value + 1 AS answer } RETURN answer ORDER BY answer" + ) + assert engine.execute(query).to_pylist() == [{"answer": 3}, {"answer": 8}] + + +def test_call_export_can_start_outer_traversal(graph): + engine, _, _ = graph + query = ( + "CALL { MATCH (p:Person) WHERE p.id = 10 RETURN p AS person } " + "MATCH (person)-[:KNOWS]->(f:Person) RETURN f.name AS name ORDER BY name" + ) + assert engine.execute(query).to_pylist() == [ + {"name": "B"}, + {"name": "B"}, + {"name": "C"}, + ] + + +def test_call_empty_outer_and_early_reader_close(graph): + engine, _, _ = graph + assert ( + engine.execute( + "MATCH (p:Person) WHERE p.id < 0 " + "CALL { RETURN 1 / 0 AS invalid } RETURN invalid" + ).num_rows + == 0 + ) + reader = engine.execute_batches( + "MATCH (p:Person) CALL { WITH p UNWIND range(1, 10000) AS n RETURN n } RETURN n" + ) + assert reader.read_next_batch().num_rows > 0 + reader.close() + assert engine.execute("RETURN 1 AS ok").to_pylist() == [{"ok": 1}] + + +@pytest.mark.parametrize( + ("query", "message"), + [ + ("UNWIND [1] AS x CALL { RETURN x AS y } RETURN y", "Undefined variable"), + ("UNWIND [1] AS x CALL { RETURN 2 AS x } RETURN x", "collides"), + ("CALL { RETURN 1 + 2 } RETURN 1", "need an alias"), + ("CALL { WITH missing RETURN missing AS x } RETURN x", "undefined variable"), + ( + "UNWIND [1] AS x CALL { WITH x + 1 AS y RETURN y } RETURN y", + "importing WITH", + ), + ( + "CALL { UNWIND [1] AS local RETURN 2 AS y } RETURN local", + "Undefined variable", + ), + ], +) +def test_call_scope_errors(graph, query, message): + engine, _, _ = graph + with pytest.raises(ValueError, match=message): + engine.execute(query) + + +def test_call_optional_entity_import_and_relationship_export(graph): + engine, _, _ = graph + query = ( + "MATCH (p:Person) OPTIONAL MATCH (p)-[:KNOWS]->(f:Person) " + "CALL { WITH f RETURN f.name AS name } " + "RETURN p.id AS id, name ORDER BY id, name" + ) + assert engine.execute(query).to_pylist() == [ + {"id": 10, "name": "B"}, + {"id": 10, "name": "B"}, + {"id": 10, "name": "C"}, + {"id": 20, "name": "D"}, + {"id": 30, "name": None}, + {"id": 40, "name": None}, + ] + query = ( + "CALL { MATCH (p:Person)-[r:KNOWS]->(f:Person) RETURN r AS edge } " + "CALL { WITH edge RETURN edge.dst AS destination } " + "RETURN destination ORDER BY destination" + ) + assert engine.execute(query).to_pylist() == [ + {"destination": 20}, + {"destination": 20}, + {"destination": 30}, + {"destination": 40}, + ] + + +def test_call_preserves_container_imports(graph): + engine, _, _ = graph + assert engine.execute( + "UNWIND [{values: [1, 2]}, {values: [3]}] AS item " + "CALL { WITH item UNWIND item.values AS n RETURN sum(n) AS total } " + "RETURN total ORDER BY total" + ).to_pylist() == [{"total": 3}, {"total": 3}] + + +def test_call_scalar_import_does_not_import_entity_shaped_names(graph): + engine, _, _ = graph + assert engine.execute( + "WITH 1 AS x, 2 AS x__other CALL { WITH x RETURN x AS y } RETURN y" + ).to_pylist() == [{"y": 1}] + + +def test_call_parameterized_child_limit_is_per_row_and_fresh(graph): + engine, _, _ = graph + query = ( + "MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) " + "RETURN f.id AS friend ORDER BY friend LIMIT $take } " + "RETURN p.id AS id, friend ORDER BY id, friend" + ) + for take, expected in [ + (1, [{"id": 10, "friend": 20}, {"id": 20, "friend": 40}]), + (0, []), + ( + 2, + [ + {"id": 10, "friend": 20}, + {"id": 10, "friend": 20}, + {"id": 20, "friend": 40}, + ], + ), + ]: + assert engine.execute(query, {"take": take}).to_pylist() == expected From c16fd54b8383c3be9cc1ded0f7800fd5401971e7 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:08:11 -0700 Subject: [PATCH 18/66] feat(cypher): execute union queries with isolated scopes Parse UNION and UNION ALL as independent branches, validate matching output columns and compatible Arrow types, and preserve distinct versus bag semantics. Include all branches in query metadata and document direct Rust AST migration. Part of #20. Submitted by David Hughes (@stugorf), with AI Collaboration. Full combined validation and SF1 measurements remain pending. --- crates/lance-graph/README.md | 10 ++ crates/lance-graph/src/ast.rs | 37 ++++++ .../src/datafusion_planner/analysis.rs | 3 +- .../datafusion_planner/builder/basic_ops.rs | 89 ++++++++++++++ .../src/datafusion_planner/builder/helpers.rs | 6 +- .../src/datafusion_planner/builder/mod.rs | 3 + crates/lance-graph/src/logical_plan.rs | 18 ++- .../lance-graph/src/parameter_substitution.rs | 3 + crates/lance-graph/src/parser.rs | 34 +++++ crates/lance-graph/src/query.rs | 116 +++++++++++++----- crates/lance-graph/src/semantic.rs | 36 ++++++ docs/api-reference.md | 2 +- docs/cypher.md | 12 +- python/python/tests/test_union.py | 96 +++++++++++++++ 14 files changed, 425 insertions(+), 40 deletions(-) create mode 100644 python/python/tests/test_union.py diff --git a/crates/lance-graph/README.md b/crates/lance-graph/README.md index 9750021..8a7c8cb 100644 --- a/crates/lance-graph/README.md +++ b/crates/lance-graph/README.md @@ -134,6 +134,16 @@ namespace utilities. Most APIs return `Result`. Errors include parsing failures, missing mappings, and execution issues surfaced from DataFusion. +## Constructing query ASTs directly + +Prefer `CypherQuery::new` or the query builder when you do not need direct AST access. Pagination fields +now use `Option`: wrap literal values with `RowCount::Literal(n)`, or use +`RowCount::Parameter(name)` before parameter substitution. AST struct literals also set +`union_branches: Vec::new()` for a single query. `UnionBranch` holds each following query and its `all` +flag; each branch has its own variable scope. Exhaustive matches on `LogicalOperator` must handle the +new `SingleRow` and `QueryUnion` variants. These are Rust source compatibility changes; Python method +signatures are unchanged by UNION. + ## Testing ```bash diff --git a/crates/lance-graph/src/ast.rs b/crates/lance-graph/src/ast.rs index 24b0304..d11233d 100644 --- a/crates/lance-graph/src/ast.rs +++ b/crates/lance-graph/src/ast.rs @@ -56,6 +56,16 @@ pub struct CypherQuery { pub order_by: Option, /// SKIP/OFFSET clause (optional) pub skip: Option, + /// Following UNION branches, evaluated in order with independent variable scopes. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub union_branches: Vec, +} + +/// One right-hand query of UNION (distinct) or UNION ALL (bag concatenation). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct UnionBranch { + pub all: bool, + pub query: CypherQuery, } impl CypherQuery { @@ -70,6 +80,19 @@ impl CypherQuery { ) } + /// Visit query scopes, including UNION branches and nested CALL bodies. + pub(crate) fn for_each_query(&self, visit: &mut dyn FnMut(&Self)) { + visit(self); + for clause in self.all_reading_clauses() { + if let ReadingClause::Call(query) = clause { + query.for_each_query(visit); + } + } + for branch in &self.union_branches { + branch.query.for_each_query(visit); + } + } + /// Extract all node labels referenced in the query pub fn get_node_labels(&self) -> Vec { let mut labels = Vec::new(); @@ -110,6 +133,13 @@ impl CypherQuery { } } } + for branch in &self.union_branches { + for label in branch.query.get_node_labels() { + if !labels.contains(&label) { + labels.push(label); + } + } + } labels } @@ -130,6 +160,13 @@ impl CypherQuery { } } } + for branch in &self.union_branches { + for rel_type in branch.query.get_relationship_types() { + if !types.contains(&rel_type) { + types.push(rel_type); + } + } + } types } diff --git a/crates/lance-graph/src/datafusion_planner/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index aa9a967..cda0dd7 100644 --- a/crates/lance-graph/src/datafusion_planner/analysis.rs +++ b/crates/lance-graph/src/datafusion_planner/analysis.rs @@ -207,7 +207,8 @@ fn analyze_operator( | LogicalOperator::Distinct { input } => { analyze_operator(input, analysis, rel_counter)?; } - LogicalOperator::Join { left, right, .. } => { + LogicalOperator::Join { left, right, .. } + | LogicalOperator::QueryUnion { left, right, .. } => { analyze_operator(left, analysis, rel_counter)?; analyze_operator(right, analysis, rel_counter)?; } 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 9a6c2d3..16a32bc 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs @@ -9,6 +9,34 @@ use crate::error::Result; use crate::logical_plan::*; use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, SortExpr}; +fn union_compatible(left: &arrow_schema::DataType, right: &arrow_schema::DataType) -> bool { + use arrow_schema::DataType; + if left == right + || left.is_null() + || right.is_null() + || (left.is_numeric() && right.is_numeric()) + { + return true; + } + let text = + |t: &DataType| matches!(t, DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View); + if text(left) && text(right) { + return true; + } + match (left, right) { + (DataType::List(left), DataType::List(right)) => { + union_compatible(left.data_type(), right.data_type()) + } + (DataType::Struct(left), DataType::Struct(right)) => { + left.len() == right.len() + && left.iter().zip(right).all(|(l, r)| { + l.name() == r.name() && union_compatible(l.data_type(), r.data_type()) + }) + } + _ => false, + } +} + impl DataFusionPlanner { /// UNION ALL of alternation branches. Each branch is planned with its own variable /// labels (the analysis keeps only one per variable), then all branches are projected @@ -79,6 +107,67 @@ impl DataFusionPlanner { .map_err(|e| self.plan_error("Failed to build UNION", e)) } + /// Language-level UNION validates the positional output contract before concatenation. + pub(crate) fn build_query_union( + &self, + ctx: &mut PlanningContext, + left: &LogicalOperator, + right: &LogicalOperator, + all: bool, + ) -> Result { + let mut plans = Vec::with_capacity(2); + for input in [left, right] { + let mut analysis = ctx.analysis.clone(); + Self::collect_labels(input, &mut analysis.var_to_label); + let mut branch = PlanningContext::new(&analysis); + branch.relationship_instance_idx = ctx.relationship_instance_idx.clone(); + plans.push(self.build_operator(&mut branch, input)?); + ctx.relationship_instance_idx = branch.relationship_instance_idx; + } + let right = plans.pop().unwrap(); + let left = plans.pop().unwrap(); + let l = left.schema(); + let r = right.schema(); + if l.fields().len() != r.fields().len() + || l.fields() + .iter() + .zip(r.fields()) + .any(|(l, r)| !l.name().eq_ignore_ascii_case(r.name())) + { + return Err(crate::error::GraphError::PlanError { + message: "UNION branches must return the same column names in the same order" + .to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + for (left, right) in l.fields().iter().zip(r.fields()) { + if !union_compatible(left.data_type(), right.data_type()) { + return Err(crate::error::GraphError::PlanError { + message: format!( + "UNION column '{}' has incompatible types {} and {}", + left.name(), + left.data_type(), + right.data_type() + ), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + } + let union = LogicalPlanBuilder::from(left) + .union(right) + .map_err(|e| self.plan_error("Failed to build query UNION", e))?; + let union = if all { + union + } else { + union + .distinct() + .map_err(|e| self.plan_error("Failed to deduplicate UNION", e))? + }; + union + .build() + .map_err(|e| self.plan_error("Failed to build query UNION", e)) + } + pub(crate) fn build_filter( &self, ctx: &mut PlanningContext, diff --git a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs index 0df0bf7..7044c3e 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs @@ -81,7 +81,8 @@ impl DataFusionPlanner { } } // Binary operator: recurse into both left and right - LogicalOperator::Join { left, right, .. } => { + LogicalOperator::Join { left, right, .. } + | LogicalOperator::QueryUnion { left, right, .. } => { Self::collect_variables(left, vars); Self::collect_variables(right, vars); } @@ -158,7 +159,8 @@ impl DataFusionPlanner { | LogicalOperator::Limit { input, .. } | LogicalOperator::Offset { input, .. } | LogicalOperator::Distinct { input } => Self::collect_labels(input, labels), - LogicalOperator::Join { left, right, .. } => { + LogicalOperator::Join { left, right, .. } + | LogicalOperator::QueryUnion { left, right, .. } => { Self::collect_labels(left, labels); Self::collect_labels(right, labels); } diff --git a/crates/lance-graph/src/datafusion_planner/builder/mod.rs b/crates/lance-graph/src/datafusion_planner/builder/mod.rs index 2e6460a..b738755 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/mod.rs @@ -119,6 +119,9 @@ impl DataFusionPlanner { label, } => self.build_unwind(ctx, input, expression, alias, label.as_deref()), LogicalOperator::Union { inputs } => self.build_union(ctx, inputs), + LogicalOperator::QueryUnion { left, right, all } => { + self.build_query_union(ctx, left, right, *all) + } LogicalOperator::ListComprehension { input, list, diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index ad7e8d1..d962c5d 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -20,6 +20,12 @@ use std::collections::HashMap; pub enum LogicalOperator { /// One row with no bound variables, for queries without reading clauses. SingleRow, + /// Language-level UNION; unlike pattern alternation, branches must share output schema. + QueryUnion { + left: Box, + right: Box, + all: bool, + }, /// Scan all nodes with a specific label ScanByLabel { variable: String, @@ -413,6 +419,15 @@ impl<'a> LogicalPlanner<'a> { }; } + for branch in &query.union_branches { + let mut planner = LogicalPlanner::new(self.config).with_fan_out(self.fan_out); + let right = planner.plan(&branch.query)?; + plan = LogicalOperator::QueryUnion { + left: Box::new(plan), + right: Box::new(right), + all: branch.all, + }; + } Ok(plan) } @@ -1394,7 +1409,8 @@ impl<'a> LogicalPlanner<'a> { LogicalOperator::Sort { input, .. } => self.extract_variable_from_plan(input), LogicalOperator::Offset { input, .. } => self.extract_variable_from_plan(input), LogicalOperator::Limit { input, .. } => self.extract_variable_from_plan(input), - LogicalOperator::Join { left, right, .. } => { + LogicalOperator::Join { left, right, .. } + | LogicalOperator::QueryUnion { left, right, .. } => { // Prefer the right branch's tail variable, else fall back to left self.extract_variable_from_plan(right) .or_else(|_| self.extract_variable_from_plan(left)) diff --git a/crates/lance-graph/src/parameter_substitution.rs b/crates/lance-graph/src/parameter_substitution.rs index ddb8a88..3808a58 100644 --- a/crates/lance-graph/src/parameter_substitution.rs +++ b/crates/lance-graph/src/parameter_substitution.rs @@ -7,6 +7,9 @@ pub fn substitute_parameters( query: &mut CypherQuery, parameters: &HashMap, ) -> Result<()> { + for branch in &mut query.union_branches { + substitute_parameters(&mut branch.query, parameters)?; + } substitute_row_count(&mut query.skip, parameters)?; substitute_row_count(&mut query.limit, parameters)?; diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index 5634a55..f9561c0 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -58,6 +58,9 @@ fn normalize_inline_properties(query: &mut CypherQuery) { ReadingClause::Unwind(_) => {} } } + for branch in &mut query.union_branches { + normalize_inline_properties(&mut branch.query); + } } /// Inline node properties whose value is a variable or property from an earlier clause @@ -112,6 +115,24 @@ fn move_variable_properties_to_where(clause: &mut MatchClause, anonymous: &mut u // Top-level parser for a complete Cypher query fn cypher_query(input: &str) -> IResult<&str, CypherQuery> { + let (input, mut query) = single_query(input)?; + let (input, branches) = many0(pair(union_header, single_query))(input)?; + query.union_branches = branches + .into_iter() + .map(|(all, query)| UnionBranch { all, query }) + .collect(); + Ok((input, query)) +} + +fn union_header(input: &str) -> IResult<&str, bool> { + let (input, _) = multispace0(input)?; + let (input, _) = tag_no_case("UNION")(input)?; + let (input, _) = multispace1(input)?; + let (input, all) = opt(terminated(tag_no_case("ALL"), multispace1))(input)?; + Ok((input, all.is_some())) +} + +fn single_query(input: &str) -> IResult<&str, CypherQuery> { let (input, _) = multispace0(input)?; let (input, reading_clauses) = many0(reading_clause)(input)?; let (input, pre_with_where) = opt(where_clause)(input)?; @@ -146,6 +167,7 @@ fn cypher_query(input: &str) -> IResult<&str, CypherQuery> { post_with_reading_clauses, post_with_where_clause: post_with_where, further_with_stages, + union_branches: Vec::new(), return_clause, limit, order_by, @@ -2309,6 +2331,18 @@ mod tests { } } + #[test] + fn unions_preserve_modes_and_quoted_keywords() { + let query = parse_cypher_query( + "RETURN 'UNION' AS n UNION ALL RETURN 'ALL' AS n UNION RETURN 'UNION ALL' AS n", + ) + .unwrap(); + assert_eq!(query.union_branches.len(), 2); + assert!(query.union_branches[0].all); + assert!(!query.union_branches[1].all); + assert!(parse_cypher_query("RETURN 1 AS n UNION ALL").is_err()); + } + #[test] fn pagination_accepts_parameters_without_losing_names() { let query = parse_cypher_query("RETURN 1 AS n SKIP $offset LIMIT $count").unwrap(); diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index fc9b47a..014e551 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -1189,16 +1189,15 @@ impl CypherQuery { pub fn referenced_node_labels(&self) -> Vec { let mut labels = Vec::new(); - for clause in self.ast.all_reading_clauses() { - if let ReadingClause::Call(inner) = clause { - labels.extend(inner.get_node_labels()); - } - if let ReadingClause::Match(match_clause) = clause { - for pattern in &match_clause.patterns { - self.collect_node_labels_from_pattern(pattern, &mut labels); + self.ast.for_each_query(&mut |ast| { + for clause in ast.all_reading_clauses() { + if let ReadingClause::Match(match_clause) = clause { + for pattern in &match_clause.patterns { + self.collect_node_labels_from_pattern(pattern, &mut labels); + } } } - } + }); labels.sort(); labels.dedup(); @@ -1209,16 +1208,15 @@ impl CypherQuery { pub fn referenced_relationship_types(&self) -> Vec { let mut types = Vec::new(); - for clause in self.ast.all_reading_clauses() { - if let ReadingClause::Call(inner) = clause { - types.extend(inner.get_relationship_types()); - } - if let ReadingClause::Match(match_clause) = clause { - for pattern in &match_clause.patterns { - self.collect_relationship_types_from_pattern(pattern, &mut types); + self.ast.for_each_query(&mut |ast| { + for clause in ast.all_reading_clauses() { + if let ReadingClause::Match(match_clause) = clause { + for pattern in &match_clause.patterns { + self.collect_relationship_types_from_pattern(pattern, &mut types); + } } } - } + }); types.sort(); types.dedup(); @@ -1229,27 +1227,29 @@ impl CypherQuery { pub fn variables(&self) -> Vec { let mut variables = Vec::new(); - for clause in self.ast.all_reading_clauses() { - match clause { - ReadingClause::Call(inner) => { - variables.extend(inner.call_imports().unwrap_or_default()); - variables.extend(inner.return_clause.items.iter().filter_map(|item| { - item.alias.clone().or_else(|| match &item.expression { - crate::ast::ValueExpression::Variable(name) => Some(name.clone()), - _ => None, - }) - })); - } - ReadingClause::Match(match_clause) => { - for pattern in &match_clause.patterns { - self.collect_variables_from_pattern(pattern, &mut variables); + self.ast.for_each_query(&mut |ast| { + for clause in ast.all_reading_clauses() { + match clause { + ReadingClause::Call(inner) => { + variables.extend(inner.call_imports().unwrap_or_default()); + variables.extend(inner.return_clause.items.iter().filter_map(|item| { + item.alias.clone().or_else(|| match &item.expression { + crate::ast::ValueExpression::Variable(name) => Some(name.clone()), + _ => None, + }) + })); + } + ReadingClause::Match(match_clause) => { + for pattern in &match_clause.patterns { + self.collect_variables_from_pattern(pattern, &mut variables); + } + } + ReadingClause::Unwind(unwind_clause) => { + variables.push(unwind_clause.alias.clone()); } - } - ReadingClause::Unwind(unwind_clause) => { - variables.push(unwind_clause.alias.clone()); } } - } + }); variables.sort(); variables.dedup(); @@ -1461,6 +1461,7 @@ impl CypherQueryBuilder { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: crate::ast::ReturnClause { distinct: self.distinct, items: self.return_items, @@ -1498,6 +1499,53 @@ mod tests { use super::*; use crate::config::GraphConfig; + #[tokio::test] + async fn union_distinct_all_and_schema_validation_execute() { + use arrow_array::{Int64Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])), + vec![Arc::new(Int64Array::from(vec![1, 2, 2]))], + ) + .unwrap(); + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + for (query, rows) in [ + ("MATCH (p:Person) RETURN p.id AS n UNION RETURN 2 AS n", 2), + ( + "MATCH (p:Person) RETURN p.id AS n UNION ALL RETURN 2 AS n", + 4, + ), + ( + "RETURN 1 AS n UNION ALL RETURN 1 AS n UNION RETURN 2 AS n", + 2, + ), + ("RETURN NULL AS n UNION RETURN NULL AS n", 1), + ] { + let result = CypherQuery::new(query) + .unwrap() + .with_config(config.clone()) + .execute(HashMap::from([("Person".to_string(), batch.clone())]), None) + .await + .unwrap(); + assert_eq!(result.num_rows(), rows, "{query}"); + } + for query in [ + "RETURN 1 AS x UNION RETURN 2 AS y", + "RETURN 1 AS x UNION RETURN 'a' AS x", + ] { + assert!(CypherQuery::new(query) + .unwrap() + .with_config(config.clone()) + .execute(HashMap::from([("Person".to_string(), batch.clone())]), None) + .await + .is_err()); + } + } + #[test] fn test_parse_simple_cypher_query() { let query = CypherQuery::new("MATCH (n:Person) RETURN n.name").unwrap(); diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index 507e544..dc8caa7 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -220,6 +220,19 @@ impl SemanticAnalyzer { // Phase 8: Type checking self.validate_types(&mut errors); + for branch in &mut analyzed_query.union_branches { + let mut analyzer = SemanticAnalyzer::new(self.config.clone()); + let analyzed = analyzer.analyze(&branch.query, parameters)?; + branch.query = analyzed.ast; + errors.extend( + analyzed + .errors + .into_iter() + .map(|e| format!("UNION branch: {e}")), + ); + warnings.extend(analyzed.warnings); + } + Ok(SemanticResult { ast: analyzed_query, variables: self.variables.clone(), @@ -1106,6 +1119,18 @@ mod tests { .unwrap() } + #[test] + fn union_branches_do_not_inherit_variables() { + let query = crate::parser::parse_cypher_query( + "MATCH (p:Person) RETURN p.id AS id UNION RETURN p.id AS id", + ) + .unwrap(); + let result = SemanticAnalyzer::new(test_config()) + .analyze(&query, &HashMap::new()) + .unwrap(); + assert!(result.errors.iter().any(|e| e.contains("UNION branch"))); + } + // Helper: analyze a query that only has a single RETURN expression fn analyze_return_expr(expr: ValueExpression) -> Result { let query = CypherQuery { @@ -1115,6 +1140,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![ReturnItem { @@ -1148,6 +1174,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![ReturnItem { @@ -1184,6 +1211,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![], @@ -1237,6 +1265,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![], @@ -1272,6 +1301,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![], @@ -1318,6 +1348,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![], @@ -1348,6 +1379,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), where_clause: None, with_clause: None, return_clause: ReturnClause { @@ -1391,6 +1423,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), where_clause: None, with_clause: None, return_clause: ReturnClause { @@ -1440,6 +1473,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), where_clause: None, with_clause: None, return_clause: ReturnClause { @@ -1511,6 +1545,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), where_clause: None, with_clause: None, return_clause: ReturnClause { @@ -1555,6 +1590,7 @@ mod tests { post_with_reading_clauses: vec![], post_with_where_clause: None, further_with_stages: vec![], + union_branches: Vec::new(), return_clause: ReturnClause { distinct: false, items: vec![ReturnItem { diff --git a/docs/api-reference.md b/docs/api-reference.md index cb5a38b..c1ae261 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -174,7 +174,7 @@ A parsed Cypher query. Raises `ValueError` if `query_text` doesn't parse. Builde : The query text and the bound parameters. `#!python variables() -> list[str]`, `#!python node_labels() -> list[str]`, `#!python relationship_types() -> list[str]` -: Names referenced by the query. +: Names referenced by the query, including independently scoped UNION branches. List ordering is not guaranteed. --- diff --git a/docs/cypher.md b/docs/cypher.md index 2eec634..74661d8 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -190,13 +190,14 @@ read-only subset of openCypher, enough to run all 14 (IC14 without `reduce()`). | `WHERE` | After `MATCH`, `OPTIONAL MATCH` (filters inside the optional part) or `WITH`. | | `WITH [DISTINCT]` | With `ORDER BY`, `LIMIT` and `WHERE`; chain as many as you need. No `SKIP` inside `WITH`. | | `UNWIND list AS x` | | +| `UNION`, `UNION ALL` | Combine independent queries with matching output names/order and compatible Arrow types. | | `CALL { ... RETURN ... }` | Read-only subquery. Import outer variables with a leading plain `WITH`; see below. | | `RETURN [DISTINCT]` | With `AS` aliases. Without reading clauses, expressions run on one implicit row. | | `ORDER BY`, `SKIP`, `LIMIT` | `SKIP` and `LIMIT` take nonnegative integer literals or parameters. | !!! failure "Not supported" - Writes (`CREATE`, `MERGE`, `SET`, `DELETE`, `REMOVE`), `UNION`, procedure calls (`CALL name(...)` / `YIELD`), and `FOREACH`. + Writes (`CREATE`, `MERGE`, `SET`, `DELETE`, `REMOVE`), procedure calls (`CALL name(...)` / `YIELD`), and `FOREACH`. Integer list literals retain signed 64-bit values: `RETURN [1, 2, 3] AS ids` returns integers, and `UNWIND [1131, 19791209300775] AS id RETURN id` preserves both IDs exactly. Numeric literals with a decimal @@ -204,6 +205,15 @@ point or exponent, such as `[0.1, 1e-3]`, form float32 vectors. Whole-number lis vectors; conversion to float32 happens inside the vector operation. A configured engine is still required for standalone `RETURN` queries. +`UNION` removes duplicate rows; `UNION ALL` preserves them. Branches have independent variable scopes and +keep their own `ORDER BY`, `SKIP` and `LIMIT`. Combined output has no guaranteed order. Mixed UNION modes +are evaluated from left to right. Numeric columns can promote to a common numeric type; incompatible +column counts, names or types raise `ValueError`. This follows the existing homogeneous Arrow value model. + +```python +engine.execute("RETURN 1 AS n UNION RETURN 1 AS n").to_pylist() # [{'n': 1}] +``` + ### Read-only subqueries Use `CALL { ... }` when a match, aggregate or limit should run separately for each incoming row: diff --git a/python/python/tests/test_union.py b/python/python/tests/test_union.py new file mode 100644 index 0000000..73ad036 --- /dev/null +++ b/python/python/tests/test_union.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Language UNION has independent scopes and a strict output-column contract.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, CypherQuery, GraphConfig + +CONFIG = ( + GraphConfig.builder().node_table("Person", "id").node_table("City", "id").build() +) +TABLES = {"Person": pa.table({"id": [1, 2, 2]}), "City": pa.table({"id": [2, 3]})} + + +@pytest.fixture(params=[False, True]) +def engine(request): + return CypherEngine(CONFIG, TABLES, use_csr=request.param) + + +@pytest.mark.parametrize( + ("query", "expected"), + [ + ("RETURN 1 AS n UNION RETURN 1 AS n", [1]), + ("RETURN 1 AS n UNION ALL RETURN 1 AS n", [1, 1]), + ("RETURN 1 AS n UNION ALL RETURN 1 AS n UNION RETURN 2 AS n", [1, 2]), + ("RETURN 1 AS n UNION RETURN 1 AS n UNION ALL RETURN 1 AS n", [1, 1]), + ( + "MATCH (p:Person) RETURN p.id AS n UNION MATCH (p:City) RETURN p.id AS n", + [1, 2, 3], + ), + ( + "MATCH (p:Person) RETURN p.id AS n UNION ALL " + "MATCH (p:City) RETURN p.id AS n", + [1, 2, 2, 2, 3], + ), + ("RETURN NULL AS n UNION RETURN NULL AS n", [None]), + ("RETURN NULL AS n UNION ALL RETURN 1 AS n", [None, 1]), + ("RETURN 1 AS n LIMIT 0 UNION RETURN 2 AS n", [2]), + ( + "UNWIND [1, 2] AS n RETURN n ORDER BY n DESC LIMIT 1 " + "UNION ALL RETURN 3 AS n", + [2, 3], + ), + ("RETURN 'UNION ALL' AS n UNION RETURN 'UNION' AS n", ["UNION", "UNION ALL"]), + ], +) +def test_union_results(engine, query, expected): + actual = engine.execute(query).column(0).to_pylist() + assert sorted(actual, key=lambda v: (v is not None, v)) == sorted( + expected, key=lambda v: (v is not None, v) + ) + assert engine.execute_batches(query).read_all().num_rows == len(expected) + + +def test_numeric_promotion_and_parameters(engine): + query = "RETURN $a AS n UNION ALL RETURN $b AS n" + assert sorted(engine.execute(query, {"a": 1, "b": 2.5}).column(0).to_pylist()) == [ + 1.0, + 2.5, + ] + one_shot = ( + CypherQuery(query) + .with_config(CONFIG) + .with_parameter("a", 1) + .with_parameter("b", 2) + ) + assert sorted(one_shot.execute(TABLES).column(0).to_pylist()) == [1, 2] + + +@pytest.mark.parametrize( + "query", + [ + "RETURN 1 AS x UNION RETURN 2 AS y", + "RETURN 1 AS x UNION RETURN 2 AS x, 3 AS y", + "RETURN 1 AS x UNION RETURN 'a' AS x", + "WITH 1 AS x RETURN x UNION RETURN x", + "RETURN 1 AS x UNION ALL", + ], +) +def test_invalid_union(engine, query): + with pytest.raises(ValueError): + engine.execute(query) + + +def test_union_query_metadata(): + query = CypherQuery( + "MATCH (p:Person) RETURN p.id AS n UNION MATCH (c:City) RETURN c.id AS n" + ) + assert set(query.node_labels()) == {"City", "Person"} + assert query.variables() == ["c", "p"] + relationships = CypherQuery( + "MATCH (p:Person)-[:X]->(c:City) RETURN p.id AS n " + "UNION MATCH (c:City)-[:Y]->(p:Person) RETURN p.id AS n" + ) + assert set(relationships.relationship_types()) == {"X", "Y"} From 123c8d3ef2b31ac387ff6cd893a5766acfdc7a0b Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:18:38 -0700 Subject: [PATCH 19/66] fix(query): preserve CALL scopes across UNION branches Seed branches from the same available imports while requiring each branch to declare its own WITH. Preserve entity outputs, scope catalog checks, and bypass cached shapes in every nested query. Use DataFusion UNION coercion before capturing child schemas and the final physical result schema when collecting batches. Validation: 10 focused Rust tests and 88 CALL/UNION Python cases pass, including streaming null/numeric coercion and repeated parameter changes. Combined feature validation follows integration. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/src/apply.rs | 49 +++++++ crates/lance-graph/src/ast.rs | 11 ++ .../datafusion_planner/builder/basic_ops.rs | 19 ++- crates/lance-graph/src/logical_plan.rs | 40 +++++- crates/lance-graph/src/query.rs | 83 ++++++------ crates/lance-graph/src/semantic.rs | 28 +++- docs/cypher.md | 6 + python/python/tests/test_call_subqueries.py | 123 ++++++++++++++++++ python/python/tests/test_union.py | 21 +++ 9 files changed, 337 insertions(+), 43 deletions(-) diff --git a/crates/lance-graph/src/apply.rs b/crates/lance-graph/src/apply.rs index e93c05d..d2cf84f 100644 --- a/crates/lance-graph/src/apply.rs +++ b/crates/lance-graph/src/apply.rs @@ -412,6 +412,53 @@ mod tests { } } + #[tokio::test] + async fn call_union_branches_bind_independently_and_keep_outer_duplicates() { + let result = run("UNWIND [1, 1, 2] AS x CALL { WITH x RETURN x AS n UNION WITH x RETURN x AS n } RETURN x, n ORDER BY x").await.unwrap(); + assert_eq!(ints(&result, "x"), [1, 1, 2]); + assert_eq!(ints(&result, "n"), [1, 1, 2]); + let result = run("WITH 2 AS x, 3 AS y CALL { WITH x RETURN x AS n UNION ALL WITH y RETURN y AS n } RETURN n ORDER BY n").await.unwrap(); + assert_eq!(ints(&result, "n"), [2, 3]); + let error = run("WITH 2 AS x CALL { WITH x RETURN x AS n UNION RETURN x AS n } RETURN n") + .await + .unwrap_err() + .to_string(); + assert!(error.contains("Undefined variable"), "{error}"); + } + + #[tokio::test] + async fn call_union_node_exports_and_empty_aggregate_branches() { + let result = run("MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) RETURN f AS friend UNION WITH p MATCH (p)-[:KNOWS]->(f:Person) RETURN f AS friend } RETURN p.id AS id, friend.id AS friend ORDER BY id, friend").await.unwrap(); + assert_eq!(ints(&result, "id"), [10, 10, 20]); + assert_eq!(ints(&result, "friend"), [20, 30, 40]); + let result = run("MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) RETURN count(f) AS n UNION ALL WITH p MATCH (p)-[:KNOWS]->(f:Person) WHERE f.id < 0 RETURN count(f) AS n } RETURN p.id AS id, n ORDER BY id, n").await.unwrap(); + assert_eq!(ints(&result, "n"), [0, 3, 0, 1, 0, 0, 0, 0]); + } + + #[tokio::test] + async fn call_union_uses_coerced_schema_for_numbers_and_nulls() { + let batch = run("CALL { RETURN 1 AS n UNION ALL RETURN 2.5 AS n } RETURN n ORDER BY n") + .await + .unwrap(); + assert_eq!( + batch.schema().field(0).data_type(), + &arrow_schema::DataType::Float64 + ); + for text in [ + "RETURN null AS n UNION ALL RETURN 1 AS n", + "RETURN 1 AS n UNION ALL RETURN null AS n", + "CALL { RETURN null AS n UNION ALL RETURN 1 AS n } RETURN n", + ] { + let batch = run(text).await.unwrap(); + assert_eq!( + batch.schema().field(0).data_type(), + &arrow_schema::DataType::Int64 + ); + assert!(batch.schema().field(0).is_nullable()); + assert_eq!(batch.column(0).null_count(), 1); + } + } + #[test] fn call_bypasses_snapshot_templates_in_every_reading_stage() { let catalog = Arc::new(lance_graph_catalog::InMemoryCatalog::new()); @@ -419,6 +466,8 @@ mod tests { "CALL { RETURN 1 AS n } RETURN n", "WITH 1 AS x CALL { WITH x RETURN x AS n } RETURN n", "WITH 1 AS x WITH x AS y CALL { WITH y RETURN y AS n } RETURN n", + "RETURN 1 AS n UNION CALL { RETURN 2 AS n } RETURN n", + "RETURN 1 AS n UNION RETURN 2 AS n LIMIT $take", ] { let mut query = CypherQuery::new(text).unwrap(); assert!(!query.prepare_cached_plan(catalog.clone()).unwrap()); diff --git a/crates/lance-graph/src/ast.rs b/crates/lance-graph/src/ast.rs index d11233d..5676f95 100644 --- a/crates/lance-graph/src/ast.rs +++ b/crates/lance-graph/src/ast.rs @@ -197,6 +197,17 @@ pub enum ReadingClause { } impl CypherQuery { + /// Every outer variable needed by this CALL's independently scoped UNION branches. + pub(crate) fn call_union_imports(&self) -> crate::error::Result> { + let mut names = self.call_imports()?; + for branch in &self.union_branches { + names.extend(branch.query.call_union_imports()?); + } + names.sort(); + names.dedup(); + Ok(names) + } + /// Names explicitly imported by a leading plain WITH in a CALL body. pub(crate) fn call_imports(&self) -> crate::error::Result> { let Some(clause) = self 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 16a32bc..a9179a7 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs @@ -56,6 +56,8 @@ impl DataFusionPlanner { Self::collect_labels(input, &mut analysis.var_to_label); let mut branch = PlanningContext::new(&analysis); branch.relationship_instance_idx = ctx.relationship_instance_idx.clone(); + branch.call_input = ctx.call_input.clone(); + branch.comprehension_elements = ctx.comprehension_elements.clone(); plans.push(self.build_operator(&mut branch, input)?); // Branches consume distinct relationship instances, in analysis order. ctx.relationship_instance_idx = branch.relationship_instance_idx; @@ -121,6 +123,8 @@ impl DataFusionPlanner { Self::collect_labels(input, &mut analysis.var_to_label); let mut branch = PlanningContext::new(&analysis); branch.relationship_instance_idx = ctx.relationship_instance_idx.clone(); + branch.call_input = ctx.call_input.clone(); + branch.comprehension_elements = ctx.comprehension_elements.clone(); plans.push(self.build_operator(&mut branch, input)?); ctx.relationship_instance_idx = branch.relationship_instance_idx; } @@ -153,9 +157,18 @@ impl DataFusionPlanner { }); } } - let union = LogicalPlanBuilder::from(left) - .union(right) - .map_err(|e| self.plan_error("Failed to build query UNION", e))?; + // A CALL must know its child's final schema before physical planning per row. + // DataFusion's UNION builder deliberately leaves coercion until analysis; + // perform its standard UNION coercion now, without optimizing parameter values. + let union = datafusion::logical_expr::Union::try_new_with_loose_types(vec![ + std::sync::Arc::new(left), + std::sync::Arc::new(right), + ]) + .and_then( + datafusion::optimizer::analyzer::type_coercion::TypeCoercionRewriter::coerce_union, + ) + .map_err(|e| self.plan_error("Failed to coerce query UNION", e))?; + let union = LogicalPlanBuilder::from(union); let union = if all { union } else { diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index d962c5d..339e842 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -344,6 +344,18 @@ impl<'a> LogicalPlanner<'a> { input: Option, preserve_entities: bool, ) -> Result { + let seed = input.clone(); + let imports = ( + self.variables.clone(), + self.node_lists.clone(), + self.path_labels.clone(), + ); + if preserve_entities { + let names = query.call_imports()?; + self.variables.retain(|name, _| names.contains(name)); + self.node_lists.retain(|name, _| names.contains(name)); + self.path_labels.retain(|name, _| names.contains(name)); + } // Plan main MATCH clauses; the WHERE clause guides where traversal starts. let where_expr = query.where_clause.as_ref().map(|w| &w.expression); let mut plan = self.plan_reading_clauses(input, &query.reading_clauses, where_expr)?; @@ -421,7 +433,31 @@ impl<'a> LogicalPlanner<'a> { for branch in &query.union_branches { let mut planner = LogicalPlanner::new(self.config).with_fan_out(self.fan_out); - let right = planner.plan(&branch.query)?; + planner.calls = self.calls; + if preserve_entities { + planner.variables = imports.0.clone(); + planner.node_lists = imports.1.clone(); + planner.path_labels = imports.2.clone(); + } + let right = planner.plan_query(&branch.query, seed.clone(), preserve_entities)?; + self.calls = planner.calls; + if preserve_entities { + let node_bindings = |variables: &HashMap| { + variables + .iter() + .filter(|(_, label)| self.config.get_node_mapping(label).is_some()) + .map(|(name, label)| (name.clone(), label.clone())) + .collect::>() + }; + if node_bindings(&self.variables) != node_bindings(&planner.variables) + || self.node_lists != planner.node_lists + { + return Err(GraphError::PlanError { + message: "CALL UNION node and node-list outputs must retain the same labels in every branch".into(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + } plan = LogicalOperator::QueryUnion { left: Box::new(plan), right: Box::new(right), @@ -483,7 +519,7 @@ impl<'a> LogicalPlanner<'a> { input: Option, query: &CypherQuery, ) -> Result { - let imports = query.call_imports()?; + let imports = query.call_union_imports()?; self.calls += 1; let id = self.calls; let mut inner = Self::new(self.config).with_fan_out(self.fan_out); diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 014e551..2b5f8fc 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -235,32 +235,31 @@ impl CypherQuery { &mut self, catalog: Arc, ) -> Result { - let parameter_counts = self - .ast - .limit - .iter() - .chain(self.ast.skip.iter()) - .chain( - self.ast - .with_clause - .iter() - .filter_map(|clause| clause.limit.as_ref()), - ) - .chain( - self.ast - .further_with_stages - .iter() - .filter_map(|stage| stage.with_clause.limit.as_ref()), - ) - .any(|count| matches!(count, crate::ast::RowCount::Parameter(_))); - if parameter_counts - || self.scalar_parameters().is_none() - || self - .ast - .all_reading_clauses() - .any(|clause| matches!(clause, ReadingClause::Call(_))) - { - // CALL owns an inner template that is bound independently per outer row. + let mut shape_sensitive = false; + self.ast.for_each_query(&mut |query| { + shape_sensitive |= query + .limit + .iter() + .chain(query.skip.iter()) + .chain( + query + .with_clause + .iter() + .filter_map(|clause| clause.limit.as_ref()), + ) + .chain( + query + .further_with_stages + .iter() + .filter_map(|stage| stage.with_clause.limit.as_ref()), + ) + .any(|count| matches!(count, crate::ast::RowCount::Parameter(_))) + || query + .all_reading_clauses() + .any(|clause| matches!(clause, ReadingClause::Call(_))); + }); + if shape_sensitive || self.scalar_parameters().is_none() { + // CALL binds its own child templates; row counts alter plan shape. return Ok(false); } let (graph, datafusion) = self.build_logical_plans(catalog.clone(), true)?; @@ -658,14 +657,24 @@ impl CypherQuery { .dataframe_with_catalog_and_context(catalog, ctx) .await?; - // Get schema before collecting (in case result is empty) - let result_schema = df.schema().inner().clone(); - - // Collect results - let batches = df.collect().await.map_err(|e| GraphError::ExecutionError { - message: format!("Failed to collect query results: {}", e), - location: snafu::Location::new(file!(), line!(), column!()), - })?; + // Capture the final schema after type coercion, exactly as DataFrame::collect + // creates its physical plan. UNION batches can differ in nullability, while + // the unoptimized logical schema can still have a narrower data type. + let task_context = Arc::new(df.task_ctx()); + let physical = df + .create_physical_plan() + .await + .map_err(|e| GraphError::ExecutionError { + message: format!("Failed to plan query results: {e}"), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + let result_schema = physical.schema(); + let batches = datafusion::physical_plan::collect(physical, task_context) + .await + .map_err(|e| GraphError::ExecutionError { + message: format!("Failed to collect query results: {e}"), + location: snafu::Location::new(file!(), line!(), column!()), + })?; if batches.is_empty() { // Return empty batch with the schema from the DataFrame @@ -673,9 +682,9 @@ impl CypherQuery { return Ok(arrow::record_batch::RecordBatch::new_empty(result_schema)); } - // Combine all batches - let schema = batches[0].schema(); - concat_batches(&schema, &batches).map_err(|e| GraphError::ExecutionError { + // UNION branches may emit batches with narrower nullability. The complete + // result schema must govern concatenation regardless of which branch finishes first. + concat_batches(&result_schema, &batches).map_err(|e| GraphError::ExecutionError { message: format!("Failed to concatenate result batches: {}", e), location: snafu::Location::new(file!(), line!(), column!()), }) diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index dc8caa7..8de95c1 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -27,6 +27,7 @@ pub struct SemanticAnalyzer { current_scope: ScopeType, catalog: Option>, expression_types: HashMap, + call_body: bool, } /// Information about a variable in the query @@ -78,6 +79,7 @@ impl SemanticAnalyzer { current_scope: ScopeType::Match, catalog: None, expression_types: HashMap::new(), + call_body: false, } } @@ -96,6 +98,17 @@ impl SemanticAnalyzer { query: &CypherQuery, parameters: &HashMap, ) -> Result { + // UNION siblings receive the same available CALL imports, never variables + // created by the preceding branch. Each branch opts in with its own WITH. + let (imports, import_types) = if self.call_body { + let available = (self.variables.clone(), self.expression_types.clone()); + let names = query.call_imports()?; + self.variables.retain(|name, _| names.contains(name)); + self.expression_types.retain(|name, _| names.contains(name)); + available + } else { + (HashMap::new(), HashMap::new()) + }; // Clone the query to perform parameter substitution let mut analyzed_query = query.clone(); @@ -214,6 +227,14 @@ impl SemanticAnalyzer { } } + if self.call_body + && query.return_clause.items.iter().any(|item| { + item.alias.is_none() && !matches!(item.expression, ValueExpression::Variable(_)) + }) + { + errors.push("CALL RETURN expressions need an alias".into()); + } + // Phase 7: Schema validation self.validate_schema(&mut warnings); @@ -222,6 +243,10 @@ impl SemanticAnalyzer { for branch in &mut analyzed_query.union_branches { let mut analyzer = SemanticAnalyzer::new(self.config.clone()); + analyzer.catalog = self.catalog.clone(); + analyzer.call_body = self.call_body; + analyzer.variables = imports.clone(); + analyzer.expression_types = import_types.clone(); let analyzed = analyzer.analyze(&branch.query, parameters)?; branch.query = analyzed.ast; errors.extend( @@ -248,7 +273,8 @@ impl SemanticAnalyzer { }; let mut inner = Self::new(self.config.clone()); inner.catalog = self.catalog.clone(); - for name in query.call_imports()? { + inner.call_body = true; + for name in query.call_union_imports()? { let variable = self .variables .get(&name) diff --git a/docs/cypher.md b/docs/cypher.md index 74661d8..d6ec8fa 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -236,6 +236,12 @@ The leading importing `WITH` accepts only existing variable names, without alias `DISTINCT`, filters or pagination. Use a second `WITH` to transform or filter them. A body without an importing `WITH` has its own independent scope. Calls can nest and can appear after `WITH` or `UNWIND`. +A CALL body can combine branches with `UNION` or `UNION ALL`. Each branch needs its own importing +`WITH` and may import a different subset of the outer variables. Branches cannot see variables created +by sibling branches. Deduplication happens within each incoming row, so repeated outer rows remain +repeated. Node and node-list exports must retain the same labels in every branch; mixed-label entity +outputs are not represented by the current binding metadata. + Every body must end with `RETURN`. Give expressions an `AS` alias; bare variables keep their names. Returned names must be distinct and must not already exist outside the call. Returning a node or relationship retains its properties, so later clauses can use the returned entity. Names that the body diff --git a/python/python/tests/test_call_subqueries.py b/python/python/tests/test_call_subqueries.py index 592d297..8da6d32 100644 --- a/python/python/tests/test_call_subqueries.py +++ b/python/python/tests/test_call_subqueries.py @@ -201,3 +201,126 @@ def test_call_parameterized_child_limit_is_per_row_and_fresh(graph): ), ]: assert engine.execute(query, {"take": take}).to_pylist() == expected + + +@pytest.mark.parametrize(("union", "repeats"), [("UNION", 1), ("UNION ALL", 2)]) +def test_call_union_distinct_is_per_outer_row(graph, union, repeats): + engine, _, _ = graph + query = ( + "UNWIND [1, 1, 2] AS x CALL { WITH x RETURN x + $value AS n " + f"{union} WITH x RETURN x + $value AS n }} RETURN x, n ORDER BY x" + ) + for value in [10, 20, 10]: + expected = [{"x": x, "n": x + value} for x in [1, 1, 2] for _ in range(repeats)] + assert engine.execute(query, {"value": value}).to_pylist() == expected + assert ( + engine.execute_batches(query, {"value": value}).read_all().to_pylist() + == expected + ) + + +def test_call_union_branches_import_different_variables_without_leaking(graph): + engine, _, _ = graph + assert engine.execute( + "WITH 2 AS x, 3 AS y CALL { WITH x RETURN x AS n " + "UNION ALL WITH y RETURN y AS n } RETURN n ORDER BY n" + ).to_pylist() == [{"n": 2}, {"n": 3}] + assert engine.execute( + "WITH 3 AS y CALL { RETURN 2 AS n " + "UNION ALL WITH y RETURN y AS n } RETURN n ORDER BY n" + ).to_pylist() == [{"n": 2}, {"n": 3}] + for query in [ + "WITH 2 AS x CALL { WITH x RETURN x AS n UNION RETURN x AS n } RETURN n", + "WITH 2 AS x CALL { RETURN x AS n UNION WITH x RETURN x AS n } RETURN n", + "WITH 2 AS x CALL { WITH x RETURN x AS n UNION WITH x RETURN x + 1 } RETURN n", + ]: + with pytest.raises(ValueError): + engine.execute(query) + + +@pytest.mark.parametrize("union", ["UNION", "UNION ALL"]) +def test_call_union_nodes_and_per_branch_top_k(graph, union): + engine, _, _ = graph + query = ( + "MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) " + "RETURN f AS friend ORDER BY friend.age DESC LIMIT 1 " + f"{union} WITH p MATCH (p)-[:KNOWS]->(f:Person) " + "RETURN f AS friend ORDER BY friend.age ASC LIMIT 1 } " + "RETURN p.id AS id, friend.name AS name ORDER BY id, name" + ) + expected = [ + {"id": 10, "name": "B"}, + {"id": 10, "name": "C"}, + {"id": 20, "name": "D"}, + ] + if union == "UNION ALL": + expected.append({"id": 20, "name": "D"}) + assert engine.execute(query).to_pylist() == expected + + +def test_call_union_aggregates_and_empty_branches(graph): + engine, _, _ = graph + assert engine.execute( + "MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) " + "RETURN count(f) AS n UNION ALL WITH p MATCH (p)-[:KNOWS]->(f:Person) " + "WHERE f.id < 0 RETURN count(f) AS n } RETURN p.id AS id, n ORDER BY id, n" + ).to_pylist() == [ + {"id": person, "n": n} + for person, count in [(10, 3), (20, 1), (30, 0), (40, 0)] + for n in [0, count] + ] + assert engine.execute( + "UNWIND [1, 2] AS x CALL { WITH x UNWIND [] AS empty RETURN x AS n " + "UNION ALL WITH x RETURN x AS n } RETURN n ORDER BY n" + ).to_pylist() == [{"n": 1}, {"n": 2}] + + +def test_call_in_later_union_branch_uses_fresh_parameters(graph): + engine, _, _ = graph + query = "RETURN 0 AS n UNION ALL CALL { RETURN $value AS n } RETURN n" + for value in [1, 2, 1]: + assert sorted( + row["n"] for row in engine.execute(query, {"value": value}).to_pylist() + ) == [0, value] + + +def test_call_union_schema_validation_uses_each_branch_catalog(graph): + engine, _, _ = graph + with pytest.raises(ValueError, match="Type mismatch"): + engine.execute( + "MATCH (p:Person) CALL { WITH p RETURN p.age AS n " + "UNION ALL WITH p RETURN p.name + 1 AS n } RETURN n" + ) + + +def test_call_union_rejects_mixed_node_label_exports(): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .node_table("City", "id") + .build() + ) + engine = CypherEngine( + config, {"Person": pa.table({"id": [1]}), "City": pa.table({"id": [1]})} + ) + with pytest.raises(ValueError, match="same labels"): + engine.execute( + "CALL { MATCH (p:Person) RETURN p AS entity " + "UNION MATCH (c:City) RETURN c AS entity } RETURN entity.id" + ) + + +def test_call_union_coerces_numeric_and_null_output_schema(graph): + engine, _, _ = graph + query = ( + "UNWIND [1, 2] AS x CALL { WITH x RETURN x AS n UNION ALL " + "RETURN $value AS n } RETURN n ORDER BY n" + ) + for value, expected in [(2.5, [1.0, 2.0, 2.5, 2.5]), (None, [None, None, 1, 2])]: + table = engine.execute(query, {"value": value}) + expected_type = pa.float64() if value is not None else pa.int64() + assert table.schema.field("n").type == expected_type + assert table.column(0).to_pylist() == expected + reader = engine.execute_batches(query, {"value": value}) + assert reader.schema.field("n").type == expected_type + assert reader.read_all().column(0).to_pylist() == expected diff --git a/python/python/tests/test_union.py b/python/python/tests/test_union.py index 73ad036..acdbce9 100644 --- a/python/python/tests/test_union.py +++ b/python/python/tests/test_union.py @@ -94,3 +94,24 @@ def test_union_query_metadata(): "UNION MATCH (c:City)-[:Y]->(p:Person) RETURN p.id AS n" ) assert set(relationships.relationship_types()) == {"X", "Y"} + + +def test_union_nullability_does_not_depend_on_first_batch(engine): + for query in [ + "RETURN null AS n UNION ALL RETURN 1 AS n", + "RETURN 1 AS n UNION ALL RETURN null AS n", + "RETURN $a AS n UNION ALL RETURN $b AS n", + ]: + for _ in range(5): + table = engine.execute(query, {"a": None, "b": 1}) + assert table.schema.field("n").nullable + assert sorted(table.column(0).to_pylist(), key=lambda n: n is None) == [ + 1, + None, + ] + stream = engine.execute_batches(query, {"a": None, "b": 1}).read_all() + assert stream.schema.field("n").nullable + assert sorted(stream.column(0).to_pylist(), key=lambda n: n is None) == [ + 1, + None, + ] From 67912be24541ade550b051c3566627feed958adc Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 09:59:30 -0700 Subject: [PATCH 20/66] fix(csr): stream expansion in bounded chunks Resolve #11 by resuming neighbor traversal between output batches and storing immutable deletion occurrence offsets. Add focused traversal/snapshot tests, a reusable hub memory scenario, and documentation. Assisted by Codex. --- crates/lance-graph/src/csr_expand.rs | 439 +++++++++++++++++-------- docs/csr.md | 8 + docs/cypher.md | 2 + python/benchmarks/bounded_memory.py | 42 ++- python/python/tests/test_engine_csr.py | 35 ++ 5 files changed, 377 insertions(+), 149 deletions(-) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 010023e..0feeb78 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -92,9 +92,9 @@ use crate::error::{GraphError, Result}; /// /// An immutable CSR base, shared between snapshots via `Arc`, plus one delta applied /// by [`Adjacency::with_changes`]: appended edges per source key, and tombstones as -/// multiset counts of `(source, neighbour)` pairs. Counting pairs instead of -/// recording edge identities keeps the base at 8 bytes per edge, and is exact -/// because identical edges are indistinguishable in query results. +/// sorted offsets into each source's base-plus-appended neighbours. Offsets let +/// readers skip deleted occurrences without copying per-source deletion state. +/// Identical edges remain distinct occurrences; the base stays at 8 bytes per edge. /// /// # Design note for future changes (read before "optimizing" the delta) /// @@ -104,6 +104,8 @@ use crate::error::{GraphError, Result}; /// in-flight queries keep using the previous one, so the delta cannot be mutated in /// place. The copy is bounded: the engine rebuilds once the delta exceeds /// `CSR_REBUILD_FRACTION` (25%) of the base. +/// Deletions additionally scan each affected source's neighbours once to resolve +/// occurrence offsets; queries can then resume without allocating deletion maps. /// /// Measured at 10M edges (`python/benchmarks/csr_delta_accumulated.py`, 20 × 1% /// appends): refresh grows from about 11 ms to about 80 ms as the delta reaches 20%, @@ -120,7 +122,7 @@ use crate::error::{GraphError, Result}; pub struct Adjacency { base: Arc, added: FastMap>, - removed: FastMap>, + removed: FastMap>, added_count: u64, removed_count: u64, } @@ -173,7 +175,7 @@ impl Adjacency { /// /// Callers must only remove edges that are present. `self` is unchanged, so /// snapshots still using it keep their answers. The base is never copied. - // shortcut: copies the accumulated delta on every call (O(delta), ~80 ms at 2M delta edges); bounded by the engine's 25% rebuild threshold. Upgrade to a persistent map if refreshes get frequent on much larger graphs. + // shortcut: copies the accumulated delta on refresh; bounded by the engine's 25% rebuild threshold. Upgrade to a persistent map if refreshes get frequent on much larger graphs. pub fn with_changes(&self, added: &[(i64, i64)], removed: &[(i64, i64)]) -> Self { let mut next = Self { base: self.base.clone(), @@ -185,39 +187,79 @@ impl Adjacency { for &(from, to) in added { next.added.entry(from).or_default().push(to as u64); } + let mut removals: FastMap> = FastMap::default(); for &(from, to) in removed { - *next - .removed + *removals .entry(from) .or_default() .entry(to as u64) .or_default() += 1; } + for (from, mut counts) in removals { + let base = next.base_neighbors(from); + let added = next.added.get(&from).map_or(&[][..], Vec::as_slice); + let prior = next.removed.get(&from).map_or(&[][..], Vec::as_slice); + let mut prior = prior.iter().copied().peekable(); + let mut positions = Vec::new(); + for (position, neighbor) in base.iter().chain(added).enumerate() { + if prior.peek() == Some(&position) { + prior.next(); + continue; + } + if let Some(count) = counts.get_mut(neighbor) { + positions.push(position); + *count -= 1; + if *count == 0 { + counts.remove(neighbor); + } + if counts.is_empty() { + break; + } + } + } + debug_assert!(counts.is_empty(), "removed edges must be present"); + let tombstones = next.removed.entry(from).or_default(); + tombstones.extend(positions); + tombstones.sort_unstable(); + } next } /// Call `f` for each target endpoint value reachable from `key`, duplicates included. - pub fn for_each_neighbor(&self, key: i64, mut f: impl FnMut(u64)) { - let base = self - .base + pub fn for_each_neighbor(&self, key: i64, f: impl FnMut(u64)) { + self.neighbors_from(key, &mut 0).for_each(f); + } + + fn base_neighbors(&self, key: i64) -> &[u64] { + self.base .dense .get(&key) - .map_or(&[][..], |&v| self.base.csr.neighbors(v)); - let added = self.added.get(&key).map_or(&[][..], |v| v.as_slice()); - let all = base.iter().chain(added).copied(); - match self.removed.get(&key) { - None => all.for_each(f), - Some(removed) => { - // Skip as many occurrences of each pair as were removed. - let mut remaining = removed.clone(); - for neighbor in all { - match remaining.get_mut(&neighbor) { - Some(n) if *n > 0 => *n -= 1, - _ => f(neighbor), - } + .map_or(&[], |&v| self.base.csr.neighbors(v)) + } + + /// Resume at an occurrence offset, borrowing all neighbour and deletion storage. + fn neighbors_from<'a>( + &'a self, + key: i64, + offset: &'a mut usize, + ) -> impl Iterator + 'a { + let base = self.base_neighbors(key); + let added = self.added.get(&key).map_or(&[][..], Vec::as_slice); + let removed = self.removed.get(&key).map_or(&[][..], Vec::as_slice); + let mut removed = &removed[removed.partition_point(|&position| position < *offset)..]; + base.iter() + .chain(added) + .enumerate() + .skip(*offset) + .filter_map(move |(position, &neighbor)| { + *offset = position + 1; + if removed.first() == Some(&position) { + removed = &removed[1..]; + None + } else { + Some(neighbor) } - } - } + }) } /// Live edges: base + appended - removed. @@ -744,6 +786,13 @@ pub struct Neighbors { reverse: Option>, } +/// Occurrence offsets are sufficient to pause even a high-degree source row. +#[derive(Default)] +struct NeighborPosition { + forward: usize, + reverse: usize, +} + impl Neighbors { pub fn directed(adjacency: Arc) -> Self { Self { @@ -759,15 +808,27 @@ impl Neighbors { } } - pub fn for_each_neighbor(&self, key: i64, mut f: impl FnMut(u64)) { - self.forward.for_each_neighbor(key, &mut f); - if let Some(reverse) = &self.reverse { - reverse.for_each_neighbor(key, |n| { - if n as i64 != key { - f(n) - } - }); - } + pub fn for_each_neighbor(&self, key: i64, f: impl FnMut(u64)) { + self.neighbors_from(key, &mut NeighborPosition::default()) + .for_each(f); + } + + fn neighbors_from<'a>( + &'a self, + key: i64, + position: &'a mut NeighborPosition, + ) -> impl Iterator + 'a { + let forward = self.forward.neighbors_from(key, &mut position.forward); + let reverse = self + .reverse + .as_ref() + .map(|r| r.neighbors_from(key, &mut position.reverse)); + forward.chain( + reverse + .into_iter() + .flatten() + .filter(move |&n| n as i64 != key), + ) } /// Edges visited over all keys (self-loops count once per direction in this estimate). @@ -1035,7 +1096,7 @@ impl UserDefinedLogicalNodeCore for CsrExpandNode { } } -/// Physical operator for [`CsrExpandNode`]; streams one output batch per input batch. +/// Physical operator for [`CsrExpandNode`]; expands input rows in bounded chunks. #[derive(Debug)] pub struct CsrExpandExec { input: Arc, @@ -1079,83 +1140,78 @@ impl CsrExpandExec { } } -/// Constraints on the expanded rows, applied before the output columns are taken. -#[derive(Clone, Copy)] -struct NeighborFilter<'a> { - /// Only rows whose target satisfies this (the moved target filter). - target_filter: Option<&'a CompiledTargetFilter>, - /// Only neighbours equal to this input column's value (a closing edge). - closing_index: Option, +/// One input batch and its current row/occurrence, independent of total fan-out. +struct BatchExpansion { + batch: RecordBatch, + keys: Int64Array, + bound: Option, + row: usize, + position: NeighborPosition, } -/// Expand one batch: repeat each input row once per neighbour and append both -/// endpoints, then the target node's columns when a lookup is fused in. -/// -/// Output comes in batches of at most `batch_size` rows, as DataFusion operators -/// expect: one input row can have thousands of neighbours. (Bounded batches are also -/// what a future work-splitting, morsel-style expansion would hand out.) -async fn expand_batch( - batch: RecordBatch, - key_index: usize, - adjacency: &Neighbors, - target: Option<&TargetLookup>, - neighbor_filter: NeighborFilter<'_>, - schema: &SchemaRef, - batch_size: usize, -) -> datafusion::common::Result> { - let keys = cast(batch.column(key_index), &DataType::Int64)?; - let keys = keys.as_primitive::(); - let bound = neighbor_filter - .closing_index - .map(|i| cast(batch.column(i), &DataType::Int64)) - .transpose()?; - let bound = bound.as_ref().map(|b| b.as_primitive::()); - let mut rows = Vec::new(); - let mut key_values = Vec::new(); - let mut neighbor_values = Vec::new(); - let mut target_positions = Vec::new(); - for i in 0..keys.len() { - if keys.is_null(i) || bound.is_some_and(|b| b.is_null(i)) { - continue; - } - let key = keys.value(i); - let required = bound.map(|b| b.value(i)); - adjacency.for_each_neighbor(key, |neighbor| { - let neighbor = neighbor as i64; - if required.is_some_and(|r| r != neighbor) { - return; - } - if let Some(target) = target { - // No target row: the inner join would drop this edge too. - let Some(position) = target.rows.position(neighbor) else { - return; - }; - target_positions.push(position); - } - rows.push(i as u32); - key_values.push(key); - neighbor_values.push(neighbor); - }); - } - let mut out = Vec::new(); - for start in (0..rows.len()).step_by(batch_size.max(1)) { - let end = (start + batch_size).min(rows.len()); - let chunk = ExpandedRows { - rows: rows[start..end].to_vec(), - key_values: key_values[start..end].to_vec(), - neighbor_values: neighbor_values[start..end].to_vec(), - target_positions: target_positions - .get(start..end) - .unwrap_or_default() - .to_vec(), +impl BatchExpansion { + fn try_new( + batch: RecordBatch, + key_index: usize, + closing_index: Option, + ) -> datafusion::common::Result { + let ints = |i| { + cast(batch.column(i), &DataType::Int64).map(|a| a.as_primitive::().clone()) }; - out.push(materialize(&batch, chunk, target, neighbor_filter, schema).await?); + let keys = ints(key_index)?; + let bound = closing_index.map(ints).transpose()?; + Ok(Self { + batch, + keys, + bound, + row: 0, + position: NeighborPosition::default(), + }) + } + + /// Collect only the next chunk; target properties are fetched after this returns. + fn next_chunk( + &mut self, + adjacency: &Neighbors, + target: Option<&TargetLookup>, + batch_size: usize, + ) -> ExpandedRows { + let mut chunk = ExpandedRows::default(); + while self.row < self.keys.len() { + let i = self.row; + if self.keys.is_valid(i) && self.bound.as_ref().is_none_or(|b| b.is_valid(i)) { + let key = self.keys.value(i); + let required = self.bound.as_ref().map(|b| b.value(i)); + for neighbor in adjacency.neighbors_from(key, &mut self.position) { + let neighbor = neighbor as i64; + if required.is_some_and(|r| r != neighbor) { + continue; + } + if let Some(target) = target { + // No target row: the inner join would drop this edge too. + let Some(position) = target.rows.position(neighbor) else { + continue; + }; + chunk.target_positions.push(position); + } + chunk.rows.push(i as u32); + chunk.key_values.push(key); + chunk.neighbor_values.push(neighbor); + if chunk.rows.len() == batch_size.max(1) { + return chunk; + } + } + } + self.row += 1; + self.position = NeighborPosition::default(); + } + chunk } - Ok(out) } /// Expanded rows of one input batch: input row index, both endpoints, and the /// target row position when a lookup is fused in. +#[derive(Default)] struct ExpandedRows { rows: Vec, key_values: Vec, @@ -1168,7 +1224,7 @@ async fn materialize( batch: &RecordBatch, expanded: ExpandedRows, target: Option<&TargetLookup>, - neighbor_filter: NeighborFilter<'_>, + target_filter: Option<&CompiledTargetFilter>, schema: &SchemaRef, ) -> datafusion::common::Result { let ExpandedRows { @@ -1177,7 +1233,7 @@ async fn materialize( mut neighbor_values, mut target_positions, } = expanded; - if let (Some(target), Some(filter)) = (target, neighbor_filter.target_filter) { + if let (Some(target), Some(filter)) = (target, target_filter) { // Take only the filter's columns, evaluate, and keep passing rows before // taking anything else. Null counts as false, as in a WHERE clause. let taken = target @@ -1338,8 +1394,9 @@ impl ExecutionPlan for CsrExpandExec { let schema = self.schema.clone(); let out_schema = schema.clone(); let batch_size = context.session_config().batch_size(); - // `then` keeps batches in order and polls one take at a time per partition. - let stream = input.then(move |batch| { + // Keep one input and one materialization in flight per partition. A hub's + // remaining neighbours stay in the immutable adjacency until the next poll. + let stream = futures::stream::try_unfold((input, None), move |(mut input, mut pending)| { let (adjacency, target, target_filter, schema) = ( adjacency.clone(), target.clone(), @@ -1347,29 +1404,31 @@ impl ExecutionPlan for CsrExpandExec { schema.clone(), ); async move { - let neighbor_filter = NeighborFilter { - target_filter: target_filter.as_deref(), - closing_index, - }; - expand_batch( - batch?, - key_index, - &adjacency, - target.as_deref(), - neighbor_filter, - &schema, - batch_size, - ) - .await + loop { + if pending.is_none() { + let Some(batch) = input.next().await else { + return Ok(None); + }; + pending = Some(BatchExpansion::try_new(batch?, key_index, closing_index)?); + } + let expansion = pending.as_mut().expect("input batch just set"); + let chunk = expansion.next_chunk(&adjacency, target.as_deref(), batch_size); + if chunk.rows.is_empty() { + pending = None; + continue; + } + let batch = materialize( + &expansion.batch, + chunk, + target.as_deref(), + target_filter.as_deref(), + &schema, + ) + .await?; + return Ok(Some((batch, (input, pending)))); + } } }); - let stream = stream.flat_map(|batches| { - let batches: Vec<_> = match batches { - Ok(batches) => batches.into_iter().map(Ok).collect(), - Err(e) => vec![Err(e)], - }; - futures::stream::iter(batches) - }); Ok(Box::pin(RecordBatchStreamAdapter::new(out_schema, stream))) } } @@ -1771,7 +1830,7 @@ mod tests { } #[tokio::test] - async fn expand_batch_repeats_rows_per_neighbor() { + async fn expansion_repeats_rows_per_neighbor_and_resumes_between_chunks() { let adj = Adjacency::build( edge_stream(vec![Some(1), Some(1)], vec![Some(2), Some(3)]), "src", @@ -1798,27 +1857,117 @@ mod tests { Field::new("r__src", DataType::Int32, true), Field::new("r__dst", DataType::Int32, true), ])); - let none = NeighborFilter { - target_filter: None, - closing_index: None, - }; let adj = Neighbors::directed(Arc::new(adj)); - let out = expand_batch(input.clone(), 0, &adj, None, none, &out_schema, 8192) + let mut expansion = BatchExpansion::try_new(input.clone(), 0, None).unwrap(); + let chunk = expansion.next_chunk(&adj, None, 8192); + let out = materialize(&input, chunk, None, None, &out_schema) .await .unwrap(); - assert_eq!(out.len(), 1); - assert_eq!(out[0].num_rows(), 2); - let dst = out[0] - .column(2) - .as_primitive::(); + assert_eq!(out.num_rows(), 2); + let dst = out.column(2).as_primitive::(); assert_eq!(dst.values().to_vec(), vec![2, 3]); + assert!(expansion.next_chunk(&adj, None, 8192).rows.is_empty()); // Output is split at `batch_size` rows. - let split = expand_batch(input, 0, &adj, None, none, &out_schema, 1) - .await - .unwrap(); + let mut expansion = BatchExpansion::try_new(input, 0, None).unwrap(); + assert_eq!(expansion.next_chunk(&adj, None, 1).neighbor_values, vec![2]); + assert_eq!(expansion.row, 0); + assert_eq!(expansion.position.forward, 1); + assert_eq!(expansion.next_chunk(&adj, None, 1).neighbor_values, vec![3]); + assert!(expansion.next_chunk(&adj, None, 1).rows.is_empty()); + } + + #[tokio::test] + async fn high_degree_expansion_visits_only_the_first_chunk() { + let n = 100_000; + let base = Adjacency::build( + edge_stream(vec![Some(0); n], (1..=n as i64).map(Some).collect()), + "src", + "dst", + ) + .await + .unwrap(); + let removed: Vec<_> = (1..=n as i64).step_by(2).map(|n| (0, n)).collect(); + let adj = Neighbors::directed(Arc::new(base.with_changes(&[], &removed))); + let input = + RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(vec![0])) as ArrayRef)]) + .unwrap(); + let mut expansion = BatchExpansion::try_new(input, 0, None).unwrap(); + let chunk = expansion.next_chunk(&adj, None, 32); assert_eq!( - split.iter().map(|b| b.num_rows()).collect::>(), - vec![1, 1] + chunk.neighbor_values, + (2..=64).step_by(2).collect::>() + ); + assert_eq!(chunk.rows.len(), 32); + // The rest of the hub, including its tombstones, has not been visited. + assert_eq!(expansion.row, 0); + assert_eq!(expansion.position.forward, 64); + assert_eq!(expansion.next_chunk(&adj, None, 32).neighbor_values[0], 66); + assert_eq!(neighbors(&base, 0).len(), n); + } + + #[tokio::test] + async fn chunk_boundaries_preserve_closing_edges_and_missing_target_rows() { + let adjacency = Adjacency::build( + edge_stream(vec![Some(1); 4], vec![Some(2), Some(2), Some(3), Some(999)]), + "src", + "dst", + ) + .await + .unwrap(); + let adjacency = Neighbors::directed(Arc::new(adjacency)); + let input = RecordBatch::try_from_iter([ + ( + "key", + Arc::new(Int64Array::from(vec![Some(1), Some(1), None, Some(1)])) as ArrayRef, + ), + ( + "bound", + Arc::new(Int64Array::from(vec![Some(2), None, Some(2), Some(3)])) as ArrayRef, + ), + ]) + .unwrap(); + let nodes = + RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(vec![2])) as ArrayRef)]) + .unwrap(); + let target = TargetLookup::new( + "b".to_string(), + Arc::new(NodeRows::from_batch(nodes, "id").unwrap().unwrap()), + vec![Field::new("b__id", DataType::Int64, false)], ); + let mut expansion = BatchExpansion::try_new(input, 0, Some(1)).unwrap(); + for _ in 0..2 { + let chunk = expansion.next_chunk(&adjacency, Some(&target), 1); + assert_eq!(chunk.rows, vec![0]); + assert_eq!(chunk.neighbor_values, vec![2]); + assert_eq!(chunk.target_positions, vec![0]); + } + assert!(expansion + .next_chunk(&adjacency, Some(&target), 1) + .rows + .is_empty()); + } + + #[tokio::test] + async fn resumed_neighbors_keep_duplicate_deletions_and_self_loops_exact() { + let edges = || edge_stream(vec![Some(1); 4], vec![Some(1), Some(2), Some(2), Some(3)]); + let base = Adjacency::build(edges(), "src", "dst").await.unwrap(); + let next = base.with_changes(&[(1, 2), (1, 4)], &[(1, 2), (1, 3)]); + let later = next.with_changes(&[(1, 3)], &[(1, 2), (1, 2)]); + assert_eq!(neighbors(&later, 1), vec![1, 3, 4]); + assert_eq!(neighbors(&next, 1), vec![1, 2, 2, 4]); + assert_eq!(neighbors(&base, 1), vec![1, 2, 2, 3]); + + let reverse = Adjacency::build(edges(), "dst", "src").await.unwrap(); + let both = Neighbors::undirected(Arc::new(base), Arc::new(reverse)); + let mut position = NeighborPosition::default(); + let mut got = Vec::new(); + while let Some(n) = both.neighbors_from(1, &mut position).next() { + got.push(n); + } + assert_eq!(got, vec![1, 2, 2, 3]); + let mut position = NeighborPosition::default(); + assert_eq!(both.neighbors_from(2, &mut position).next(), Some(1)); + assert_eq!(both.neighbors_from(2, &mut position).next(), Some(1)); + assert_eq!(both.neighbors_from(2, &mut position).next(), None); } } diff --git a/docs/csr.md b/docs/csr.md index c34d0e7..f819ec6 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -91,6 +91,14 @@ Node index Expect about **70 bytes per edge** of extra peak memory, for both directions together. The index lives in memory even when the tables are read lazily from Lance. + During query execution, each CSR expansion partition keeps one input batch and collects at most one + output chunk of neighbour positions before fetching properties. It yields that chunk before visiting + the remaining neighbours, even for a single high-degree node. Deleted edge positions are read from the + shared snapshot rather than copied for each query; refresh resolves those positions by scanning the + affected sources' neighbour lists. The chunk row limit comes from DataFusion's batch + size; wide properties and multiple active partitions still increase working memory. This does not bound + the resident graph, the input data, or buffering in joins, sorts and aggregations elsewhere in the plan. + !!! warning "Experimental" The CSR is marked experimental. Results are identical with it on or off (the test suite runs each diff --git a/docs/cypher.md b/docs/cypher.md index d6ec8fa..2f04416 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -118,6 +118,8 @@ print(rows) # 4 - Batches are computed as you read them. Scans, filters and projections stream with bounded memory; joins, sorts and aggregations may still buffer their inputs. +- CSR expansion visits neighbours and fetches their properties a chunk at a time, including when one node + has many neighbours. See [CSR memory cost](csr.md#turning-it-on) for what it keeps in memory. - Planning errors (a syntax error, an unknown label) raise immediately. Execution errors raise while reading. - The reader keeps its snapshot alive, so a concurrent [`refresh()`](lance.md#refresh) doesn't affect it. - Dropping the reader cancels the remaining work. diff --git a/python/benchmarks/bounded_memory.py b/python/benchmarks/bounded_memory.py index a2b5768..4e54971 100644 --- a/python/benchmarks/bounded_memory.py +++ b/python/benchmarks/bounded_memory.py @@ -3,7 +3,7 @@ """Peak memory of streamed Cypher scans over Lance data larger than RAM limits. -``generate`` writes a deterministic Person table with a wide unused ``payload`` +``generate`` writes a deterministic Person table with a wide ``payload`` column, in batches so the generator itself stays small. ``run`` executes a query in a fresh subprocess, consumes ``execute_batches()`` incrementally (count + checksum, nothing retained), and reports peak RSS from ``getrusage``, which includes native @@ -11,10 +11,19 @@ python benchmarks/bounded_memory.py generate --out /data/large --rows 2000000 python benchmarks/bounded_memory.py run --data /data/large --query scan + +For CSR fan-out, generate a hub and return its wide properties. Compare fresh +processes on the base and changed revisions for several row counts. The pre-query +RSS includes the resident CSR and node index; peak growth is a lower bound because +getrusage also includes allocations made during engine construction. + + python benchmarks/bounded_memory.py generate --out /data/hub --rows 100000 --hub + python benchmarks/bounded_memory.py run --data /data/hub --query hub """ import argparse import json +import os import resource import subprocess import sys @@ -27,10 +36,13 @@ "scan": "MATCH (p:Person) RETURN p.id", # Selective graph join; the KNOWS side is filtered before the join. "join": "MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE a.id < 1000 RETURN b.id", + "hub": ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE a.id = 0 RETURN b.id, b.payload" + ), } -def generate(out: Path, rows: int, batch_rows: int = 16_384) -> None: +def generate(out: Path, rows: int, batch_rows: int = 16_384, hub: bool = False) -> None: import lance import numpy as np import pyarrow as pa @@ -56,7 +68,11 @@ def person_batches(): reader = pa.RecordBatchReader.from_batches(schema, person_batches()) lance.write_dataset(reader, out / "person.lance", mode="overwrite") ids = np.arange(rows, dtype=np.int64) - knows = pa.table({"src": ids, "dst": (ids + 1) % rows}) + knows = pa.table( + {"src": np.zeros(rows, dtype=np.int64), "dst": ids} + if hub + else {"src": ids, "dst": (ids + 1) % rows} + ) lance.write_dataset(knows, out / "knows.lance", mode="overwrite") logical = rows * (8 + PAYLOAD_BYTES) print(json.dumps({"rows": rows, "logical_input_bytes": logical})) @@ -80,7 +96,14 @@ def measure(data: Path, query: str) -> dict: "Person": lance.dataset(data / "person.lance"), "KNOWS": lance.dataset(data / "knows.lance"), }, + use_csr=query == "hub", + ) + peak_before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + rss_before = ( + int(subprocess.check_output(["ps", "-o", "rss=", "-p", str(os.getpid())])) + * 1024 ) + query_start = time.perf_counter() reader = engine.execute_batches(QUERIES[query]) rows = checksum = batches = 0 first_batch_s = None @@ -94,7 +117,11 @@ def measure(data: Path, query: str) -> dict: peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss # Linux reports KiB, macOS reports bytes. peak_bytes = peak if sys.platform == "darwin" else peak * 1024 + peak_before_bytes = peak_before if sys.platform == "darwin" else peak_before * 1024 input_rows = lance.dataset(data / "person.lance").count_rows() + if query in {"scan", "hub"}: + assert rows == input_rows + assert checksum == input_rows * (input_rows - 1) // 2 return { "query": query, "input_rows": input_rows, @@ -102,8 +129,14 @@ def measure(data: Path, query: str) -> dict: "checksum": checksum, "batches": batches, "time_to_first_batch_s": first_batch_s, + "query_time_to_first_batch_s": first_batch_s - (query_start - start) + if first_batch_s is not None + else None, "elapsed_s": elapsed, "peak_rss_bytes": peak_bytes, + "pre_query_rss_bytes": rss_before, + "pre_query_peak_rss_bytes": peak_before_bytes, + "peak_growth_lower_bound_bytes": max(0, peak_bytes - peak_before_bytes), "plan_scan": next( ( line.strip() @@ -121,6 +154,7 @@ def main() -> None: gen = sub.add_parser("generate") gen.add_argument("--out", type=Path, required=True) gen.add_argument("--rows", type=int, required=True) + gen.add_argument("--hub", action="store_true") run = sub.add_parser("run") run.add_argument("--data", type=Path, required=True) run.add_argument("--query", choices=QUERIES, default="scan") @@ -131,7 +165,7 @@ def main() -> None: if args.cmd == "generate": args.out.mkdir(parents=True, exist_ok=True) - generate(args.out, args.rows) + generate(args.out, args.rows, hub=args.hub) elif args.cmd == "child": print(json.dumps(measure(args.data, args.query))) else: diff --git a/python/python/tests/test_engine_csr.py b/python/python/tests/test_engine_csr.py index 7741a55..b5a4dd7 100644 --- a/python/python/tests/test_engine_csr.py +++ b/python/python/tests/test_engine_csr.py @@ -7,6 +7,8 @@ sorted results must be identical. """ +from itertools import chain + import lance import pyarrow as pa import pytest @@ -197,6 +199,39 @@ def test_execute_batches_with_csr(datasets): assert rows(streamed) == rows(csr.execute(query)) +def test_high_degree_stream_keeps_its_snapshot_across_deletions(tmp_path): + n = 20_000 + content = "wide property " * 32 + people = pa.table({"id": range(n + 1), "content": [content] * (n + 1)}) + edges = pa.table({"src": [0] * n, "dst": [1, 1, *range(2, n)]}) + knows = lance.write_dataset(edges, tmp_path / "knows.lance") + person = lance.write_dataset(people, tmp_path / "person.lance") + engine = CypherEngine(CONFIG, {"Person": person, "KNOWS": knows}, use_csr=True) + query = ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) " + "WHERE a.id = 0 RETURN b.id AS id, b.content AS content" + ) + reader = engine.execute_batches(query) + first = reader.read_next_batch() + assert 0 < first.num_rows < n + # A live reader retains the old adjacency, including both copies of 0->1. + knows.delete("dst = 1") + engine.refresh() + assert engine.csr_stats()["KNOWS"]["tombstones"] == 2 + counts = {} + for batch in chain([first], reader): + for row in batch.to_pylist(): + assert row["content"] == content + counts[row["id"]] = counts.get(row["id"], 0) + 1 + assert counts == {1: 2, **dict.fromkeys(range(2, n), 1)} + assert engine.execute(query).num_rows == n - 2 + + reader = engine.execute_batches(query) + reader.read_next_batch() + reader.close() + assert engine.execute(query).num_rows == n - 2 + + def physical_plan(engine, query): plan = engine.explain(query) return plan[plan.index("physical_plan") :] From 131741ace3e181cba1399a2d25ea01dd6625548f Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:25:46 -0700 Subject: [PATCH 21/66] perf(csr): traverse variable-length paths iteratively Resolve #13 with bounded typed traversal over CSR snapshots, retaining mixed relationship types without unrolling every length or type sequence. Implement zero-hop identity in indexed and join paths while preserving the 20-hop limit. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/src/csr_expand.rs | 103 +++++- crates/lance-graph/src/csr_traversal.rs | 328 ++++++++++++++++++ .../src/datafusion_planner/analysis.rs | 12 +- .../datafusion_planner/builder/expand_ops.rs | 217 +++++++++++- .../src/datafusion_planner/builder/helpers.rs | 7 +- .../src/datafusion_planner/join_ops.rs | 2 +- .../src/datafusion_planner/scan_ops.rs | 1 + crates/lance-graph/src/lib.rs | 1 + crates/lance-graph/src/logical_plan.rs | 86 ++++- .../lance-graph/src/pattern_alternatives.rs | 132 ++++--- crates/lance-graph/src/query.rs | 4 +- docs/csr.md | 20 +- docs/cypher.md | 9 +- .../python/tests/test_engine_alternation.py | 45 +++ python/python/tests/test_engine_csr.py | 32 +- 15 files changed, 918 insertions(+), 81 deletions(-) create mode 100644 crates/lance-graph/src/csr_traversal.rs diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 0feeb78..d3fad84 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -569,7 +569,7 @@ impl NodeRows { self.positions.len() } - fn position(&self, id: i64) -> Option { + pub(crate) fn position(&self, id: i64) -> Option { self.positions.get(&id).copied() } @@ -788,7 +788,7 @@ pub struct Neighbors { /// Occurrence offsets are sufficient to pause even a high-degree source row. #[derive(Default)] -struct NeighborPosition { +pub(crate) struct NeighborPosition { forward: usize, reverse: usize, } @@ -813,7 +813,7 @@ impl Neighbors { .for_each(f); } - fn neighbors_from<'a>( + pub(crate) fn neighbors_from<'a>( &'a self, key: i64, position: &'a mut NeighborPosition, @@ -874,6 +874,7 @@ pub struct CsrExpandNode { /// Closing edge: input column holding the already-bound target id. Only neighbours /// equal to it are kept, replacing a filter after a full expansion. closing_column: Option, + traversal: Option>, schema: DFSchemaRef, } @@ -903,6 +904,7 @@ impl CsrExpandNode { target, target_filter: None, closing_column: None, + traversal: None, schema: Arc::new(schema), }) } @@ -913,6 +915,11 @@ impl CsrExpandNode { self } + pub(crate) fn with_traversal(mut self, traversal: crate::csr_traversal::Traversal) -> Self { + self.traversal = Some(Arc::new(traversal)); + self + } + /// The same node over `input` with `target`, keeping the moved filter and closing column. fn rebuilt(&self, input: LogicalPlan, target: Option) -> Result { let mut node = Self::try_new( @@ -925,6 +932,7 @@ impl CsrExpandNode { )?; node.target_filter = self.target_filter.clone(); node.closing_column = self.closing_column.clone(); + node.traversal = self.traversal.clone(); Ok(node) } @@ -980,6 +988,11 @@ impl PartialEq for CsrExpandNode { } && self.target_filter == other.target_filter && self.closing_column == other.closing_column + && match (&self.traversal, &other.traversal) { + (Some(a), Some(b)) => Arc::ptr_eq(a, b), + (None, None) => true, + _ => false, + } } } impl Eq for CsrExpandNode {} @@ -1082,6 +1095,9 @@ impl UserDefinedLogicalNodeCore for CsrExpandNode { if let Some(column) = &self.closing_column { write!(f, ", match={}", column)?; } + if let Some(walk) = &self.traversal { + write!(f, ", iterative hops={}..{}", walk.min_hops, walk.max_hops)?; + } Ok(()) } @@ -1106,6 +1122,7 @@ pub struct CsrExpandExec { /// The moved target filter, if any. target_filter: Option, closing_column: Option, + traversal: Option>, schema: SchemaRef, properties: Arc, } @@ -1119,6 +1136,7 @@ impl CsrExpandExec { target: Option, target_filter: Option, closing_column: Option, + traversal: Option>, schema: SchemaRef, ) -> Self { let properties = Arc::new(PlanProperties::new( @@ -1134,6 +1152,7 @@ impl CsrExpandExec { target, target_filter, closing_column, + traversal, schema, properties, } @@ -1147,6 +1166,7 @@ struct BatchExpansion { bound: Option, row: usize, position: NeighborPosition, + walk: Option, } impl BatchExpansion { @@ -1166,6 +1186,7 @@ impl BatchExpansion { bound, row: 0, position: NeighborPosition::default(), + walk: None, }) } @@ -1207,6 +1228,49 @@ impl BatchExpansion { } chunk } + + async fn next_path_chunk( + &mut self, + graph: &crate::csr_traversal::Traversal, + target: Option<&TargetLookup>, + batch_size: usize, + ) -> ExpandedRows { + let mut chunk = ExpandedRows::default(); + while self.row < self.keys.len() { + let i = self.row; + if self.keys.is_valid(i) && self.bound.as_ref().is_none_or(|b| b.is_valid(i)) { + let key = self.keys.value(i); + let walk = self + .walk + .get_or_insert_with(|| crate::csr_traversal::TraversalCursor::new(graph, key)); + let mut candidates = 0usize; + while let Some(neighbor) = walk.next(graph).await { + candidates += 1; + if candidates.is_multiple_of(1024) { + cooperative_yield().await; + } + if self.bound.as_ref().is_some_and(|b| b.value(i) != neighbor) { + continue; + } + if let Some(target) = target { + let Some(position) = target.rows.position(neighbor) else { + continue; + }; + chunk.target_positions.push(position); + } + chunk.rows.push(i as u32); + chunk.key_values.push(key); + chunk.neighbor_values.push(neighbor); + if chunk.rows.len() == batch_size.max(1) { + return chunk; + } + } + } + self.row += 1; + self.walk = None; + } + chunk + } } /// Expanded rows of one input batch: input row index, both endpoints, and the @@ -1300,6 +1364,20 @@ fn retain_by_mask(values: &mut Vec, keep: &[bool]) { values.retain(|_| *keep.next().expect("one mask entry per value")); } +/// Cooperate with the stream's executor without requiring a particular async runtime. +pub(crate) async fn cooperative_yield() { + let mut yielded = false; + std::future::poll_fn(move |cx| { + if std::mem::replace(&mut yielded, true) { + std::task::Poll::Ready(()) + } else { + cx.waker().wake_by_ref(); + std::task::Poll::Pending + } + }) + .await +} + impl DisplayAs for CsrExpandExec { fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { write!( @@ -1320,6 +1398,9 @@ impl DisplayAs for CsrExpandExec { if let Some(column) = &self.closing_column { write!(f, ", match={}", column)?; } + if let Some(walk) = &self.traversal { + write!(f, ", iterative hops={}..{}", walk.min_hops, walk.max_hops)?; + } Ok(()) } } @@ -1338,6 +1419,9 @@ impl ExecutionPlan for CsrExpandExec { use datafusion::common::stats::Precision; let input = self.input.partition_statistics(partition)?; let mut stats = datafusion::common::Statistics::new_unknown(&self.schema); + if self.traversal.is_some() { + return Ok(Arc::new(stats)); + } let degree = self.adjacency.mean_degree(); stats.num_rows = match input.num_rows { Precision::Exact(n) | Precision::Inexact(n) => { @@ -1371,6 +1455,7 @@ impl ExecutionPlan for CsrExpandExec { self.target.clone(), self.target_filter.clone(), self.closing_column.clone(), + self.traversal.clone(), self.schema.clone(), ))) } @@ -1394,6 +1479,7 @@ impl ExecutionPlan for CsrExpandExec { let schema = self.schema.clone(); let out_schema = schema.clone(); let batch_size = context.session_config().batch_size(); + let traversal = self.traversal.clone(); // Keep one input and one materialization in flight per partition. A hub's // remaining neighbours stay in the immutable adjacency until the next poll. let stream = futures::stream::try_unfold((input, None), move |(mut input, mut pending)| { @@ -1403,6 +1489,7 @@ impl ExecutionPlan for CsrExpandExec { target_filter.clone(), schema.clone(), ); + let traversal = traversal.clone(); async move { loop { if pending.is_none() { @@ -1412,7 +1499,14 @@ impl ExecutionPlan for CsrExpandExec { pending = Some(BatchExpansion::try_new(batch?, key_index, closing_index)?); } let expansion = pending.as_mut().expect("input batch just set"); - let chunk = expansion.next_chunk(&adjacency, target.as_deref(), batch_size); + let chunk = match &traversal { + Some(walk) => { + expansion + .next_path_chunk(walk, target.as_deref(), batch_size) + .await + } + None => expansion.next_chunk(&adjacency, target.as_deref(), batch_size), + }; if chunk.rows.is_empty() { pending = None; continue; @@ -1477,6 +1571,7 @@ impl ExtensionPlanner for CsrExtensionPlanner { node.target.clone(), target_filter, node.closing_column.clone(), + node.traversal.clone(), Arc::new(node.schema.as_arrow().clone()), )))) } diff --git a/crates/lance-graph/src/csr_traversal.rs b/crates/lance-graph/src/csr_traversal.rs new file mode 100644 index 0000000..a19eb36 --- /dev/null +++ b/crates/lance-graph/src/csr_traversal.rs @@ -0,0 +1,328 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Resumable variable-length walks over immutable, typed CSR snapshots. + +use std::sync::Arc; + +use crate::ast::RelationshipDirection; +use crate::config::GraphConfig; +use crate::csr_expand::{ + cooperative_yield, AdjacencyCatalog, NeighborPosition, Neighbors, NodeRows, +}; + +#[derive(Debug)] +struct Step { + from: usize, + to: usize, + adjacency: Neighbors, +} + +#[derive(Debug)] +pub(crate) struct Traversal { + steps: Vec, + nodes: Vec>, + start: usize, + target: usize, + pub min_hops: u32, + pub max_hops: u32, + undirected: bool, +} + +impl Traversal { + /// Every intermediate label needs a node index: joins drop dangling edges at every hop. + #[allow(clippy::too_many_arguments)] + pub fn new( + config: &GraphConfig, + catalog: &AdjacencyCatalog, + types: &[String], + direction: &RelationshipDirection, + source_label: &str, + target_label: &str, + min_hops: u32, + max_hops: u32, + ) -> Option { + let mut labels: Vec = Vec::new(); + let mut nodes = Vec::new(); + let mut label = |name: &str| -> Option { + if let Some(index) = labels.iter().position(|l| l.eq_ignore_ascii_case(name)) { + return Some(index); + } + nodes.push(catalog.nodes(name)?.clone()); + labels.push(name.to_owned()); + Some(labels.len() - 1) + }; + let (start, target) = (label(source_label)?, label(target_label)?); + let mut steps = Vec::new(); + for name in types { + let mapping = config.get_relationship_mapping(name)?; + let (from, to) = ( + label(mapping.source_label.as_deref()?)?, + label(mapping.target_label.as_deref()?)?, + ); + match direction { + RelationshipDirection::Outgoing => steps.push(Step { + from, + to, + adjacency: Neighbors::directed(catalog.get(name, true)?.clone()), + }), + RelationshipDirection::Incoming => steps.push(Step { + from: to, + to: from, + adjacency: Neighbors::directed(catalog.get(name, false)?.clone()), + }), + RelationshipDirection::Undirected if from == to => steps.push(Step { + from, + to, + adjacency: Neighbors::undirected( + catalog.get(name, true)?.clone(), + catalog.get(name, false)?.clone(), + ), + }), + RelationshipDirection::Undirected => { + steps.push(Step { + from, + to, + adjacency: Neighbors::directed(catalog.get(name, true)?.clone()), + }); + steps.push(Step { + from: to, + to: from, + adjacency: Neighbors::directed(catalog.get(name, false)?.clone()), + }); + } + } + } + (!steps.is_empty()).then_some(Self { + steps, + nodes, + start, + target, + min_hops, + max_hops, + undirected: *direction == RelationshipDirection::Undirected, + }) + } +} + +struct Frame { + label: usize, + key: i64, + step: usize, + position: NeighborPosition, + emit: bool, +} + +impl Frame { + fn new(label: usize, key: i64) -> Self { + Self { + label, + key, + step: 0, + position: NeighborPosition::default(), + emit: true, + } + } +} + +pub(crate) struct TraversalCursor { + stack: Vec, +} + +impl TraversalCursor { + pub fn new(graph: &Traversal, key: i64) -> Self { + Self { + stack: vec![Frame::new(graph.start, key)], + } + } + + /// Keep one frame per hop, and cooperate even when candidate edges produce no rows. + pub async fn next(&mut self, graph: &Traversal) -> Option { + let mut work = 0; + while !self.stack.is_empty() { + if work == 1024 { + cooperative_yield().await; + work = 0; + } + work += 1; + let depth = self.stack.len() as u32 - 1; + let frame = self.stack.last_mut().expect("nonempty stack"); + if frame.emit { + frame.emit = false; + if depth >= graph.min_hops && frame.label == graph.target { + return Some(frame.key); + } + } + if depth == graph.max_hops || frame.step == graph.steps.len() { + self.stack.pop(); + continue; + } + let step = &graph.steps[frame.step]; + let next = (step.from == frame.label) + .then(|| { + step.adjacency + .neighbors_from(frame.key, &mut frame.position) + .next() + }) + .flatten(); + let Some(key) = next.map(|n| n as i64) else { + frame.step += 1; + frame.position = NeighborPosition::default(); + continue; + }; + if graph.nodes[step.to].position(key).is_none() { + continue; + } + // Preserve the existing undirected immediate-backtracking rule, using typed ids. + // shortcut: full relationship uniqueness needs edge row ids, as on the join path. + if graph.undirected + && self.stack.len() >= 2 + && self.stack[self.stack.len() - 2].label == step.to + && self.stack[self.stack.len() - 2].key == key + { + continue; + } + self.stack.push(Frame::new(step.to, key)); + } + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::csr_expand::Adjacency; + use arrow::array::{Int64Array, RecordBatch}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + + async fn graph() -> (GraphConfig, AdjacencyCatalog) { + let config = GraphConfig::builder() + .node_table("Tag", "id") + .node_table("Class", "id") + .rel_table("hasType", "Tag", "Class", "src", "dst") + .rel_table("isSubclass", "Class", "Class", "src", "dst") + .build() + .unwrap(); + let mut catalog = AdjacencyCatalog::default(); + for (name, ids) in [("Tag", vec![1, 2]), ("Class", vec![1, 2, 3])] { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(ids))]).unwrap(); + catalog.insert_nodes( + name, + Arc::new(NodeRows::from_batch(batch, "id").unwrap().unwrap()), + ); + } + for (name, src, dst) in [ + ("hasType", vec![1, 1, 2, 1], vec![1, 1, 2, 99]), + ("isSubclass", vec![1, 2, 3], vec![2, 3, 1]), + ] { + let schema = Arc::new(Schema::new(vec![ + Field::new("src", DataType::Int64, false), + Field::new("dst", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(src)), + Arc::new(Int64Array::from(dst)), + ], + ) + .unwrap(); + for (outgoing, from, to) in [(true, "src", "dst"), (false, "dst", "src")] { + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema.clone(), + futures::stream::iter([Ok(batch.clone())]), + )); + catalog.insert( + name, + outgoing, + Arc::new(Adjacency::build(stream, from, to).await.unwrap()), + ); + } + } + (config, catalog) + } + + #[tokio::test] + async fn typed_walk_keeps_duplicates_and_skips_missing_intermediate_nodes() { + let (config, catalog) = graph().await; + let walk = Traversal::new( + &config, + &catalog, + &["hasType".into(), "isSubclass".into()], + &RelationshipDirection::Outgoing, + "Tag", + "Class", + 0, + 3, + ) + .unwrap(); + let mut cursor = TraversalCursor::new(&walk, 1); + let mut out = Vec::new(); + while let Some(key) = cursor.next(&walk).await { + out.push(key); + assert!(cursor.stack.len() <= 4); + } + out.sort(); + // Tag 1 and Class 1 are different nodes; no extra zero-hop Class 1 result. + assert_eq!(out, vec![1, 1, 2, 2, 3, 3]); + assert!(cursor.stack.is_empty()); + } + + #[tokio::test] + async fn zero_hop_is_identity_and_depth_bound_survives_cycles() { + let (config, catalog) = graph().await; + let walk = Traversal::new( + &config, + &catalog, + &["isSubclass".into()], + &RelationshipDirection::Outgoing, + "Class", + "Class", + 0, + 20, + ) + .unwrap(); + let mut cursor = TraversalCursor::new(&walk, 1); + for depth in 0..=20 { + assert_eq!(cursor.next(&walk).await, Some(1 + depth % 3)); + assert_eq!(cursor.stack.len(), depth as usize + 1); + } + assert_eq!(cursor.next(&walk).await, None); + let zero = Traversal { + max_hops: 0, + ..walk + }; + let mut cursor = TraversalCursor::new(&zero, 2); + assert_eq!(cursor.next(&zero).await, Some(2)); + assert_eq!(cursor.next(&zero).await, None); + } + + #[tokio::test] + async fn walk_yields_while_no_candidate_can_match_the_target_label() { + let (config, mut catalog) = graph().await; + let adjacency = catalog + .get("isSubclass", true) + .unwrap() + .with_changes(&vec![(1, 1); 4096], &[]); + catalog.insert("isSubclass", true, Arc::new(adjacency)); + let walk = Traversal::new( + &config, + &catalog, + &["isSubclass".into()], + &RelationshipDirection::Outgoing, + "Class", + "Tag", + 1, + 2, + ) + .unwrap(); + let mut cursor = TraversalCursor::new(&walk, 1); + let mut next = Box::pin(cursor.next(&walk)); + assert!(futures::poll!(&mut next).is_pending()); + drop(next); + assert!(cursor.stack.len() <= 3); + } +} diff --git a/crates/lance-graph/src/datafusion_planner/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index cda0dd7..54a20df 100644 --- a/crates/lance-graph/src/datafusion_planner/analysis.rs +++ b/crates/lance-graph/src/datafusion_planner/analysis.rs @@ -146,6 +146,7 @@ fn analyze_operator( input, source_variable, target_variable, + target_label, relationship_types, direction, relationship_variable, @@ -156,13 +157,9 @@ fn analyze_operator( // Recursively analyze input first analyze_operator(input, analysis, rel_counter)?; - // Infer target variable's label from source variable - // For (a:Person)-[:KNOWS]->(b), b also gets label Person - if let Some(source_label) = analysis.var_to_label.get(source_variable).cloned() { - analysis - .var_to_label - .insert(target_variable.clone(), source_label); - } + analysis + .var_to_label + .insert(target_variable.clone(), target_label.clone()); // For variable-length paths, register multiple instances (one per hop) // We need to register instances for all possible hop counts @@ -389,6 +386,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".to_string(), target_variable: "b".to_string(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".to_string()], direction: RelationshipDirection::Outgoing, relationship_variable: None, 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 9b4ad83..26777d3 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs @@ -406,11 +406,7 @@ impl DataFusionPlanner { })) } - /// Build variable-length path expansion using unrolling + UNION strategy - /// - /// For a query like: (a)-[:KNOWS*1..3]->(b) - /// This generates: - /// 1-hop plan UNION 2-hop plan UNION 3-hop plan + /// Use one iterative CSR expansion when indexed; otherwise union the fixed-length plans. #[allow(clippy::too_many_arguments)] pub(crate) fn build_variable_length_expand( &self, @@ -424,7 +420,7 @@ impl DataFusionPlanner { max_length: Option, target_properties: &HashMap, ) -> Result { - let min_hops = min_length.unwrap_or(1).max(1); + let min_hops = min_length.unwrap_or(1); let max_hops = max_length.unwrap_or(crate::MAX_VARIABLE_LENGTH_HOPS); // Validate range @@ -465,6 +461,30 @@ impl DataFusionPlanner { .map(|f| f.name().clone()), ); + if let Some(plan) = self.try_build_csr_traversal( + ctx, + input_plan.clone(), + source_variable, + target_variable, + relationship_types, + direction, + min_hops, + max_hops, + target_properties, + )? { + let projection: Vec<_> = plan + .schema() + .fields() + .iter() + .filter(|f| expected_columns.contains(f.name().as_str())) + .map(|f| col(f.name())) + .collect(); + return LogicalPlanBuilder::from(plan) + .project(projection) + .and_then(|p| p.build()) + .map_err(|e| self.plan_error("Failed to project variable-length CSR", e)); + } + // Generate a plan for each hop count and UNION them let mut plans = Vec::new(); @@ -544,6 +564,59 @@ impl DataFusionPlanner { hop_count: u32, target_properties: &HashMap, ) -> Result { + if hop_count == 0 { + let (source_label, source_map) = self.get_target_node_mapping(ctx, source_variable)?; + let (target_label, target_map) = self.get_target_node_mapping(ctx, target_variable)?; + let source = col(qualify_column(source_variable, &source_map.id_field)); + let target = qualify_column(target_variable, &target_map.id_field); + let same_label = source_label.eq_ignore_ascii_case(&target_label); + let mut builder = LogicalPlanBuilder::from(input_plan); + if builder + .schema() + .field_with_unqualified_name(&target) + .is_ok() + { + if !target_properties.is_empty() { + return Err(crate::error::GraphError::PlanError { + message: format!("Property filters on reused variable '{target_variable}' are not yet supported"), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + builder = builder + .filter( + source + .eq(col(&target)) + .and(datafusion::logical_expr::lit(same_label)), + ) + .map_err(|e| self.plan_error("Failed to filter zero-hop identity", e))?; + } else { + let target_source = + self.get_catalog()? + .node_source(&target_label) + .ok_or_else(|| crate::error::GraphError::PlanError { + message: format!("No node source for {target_label}"), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + let scan = self.build_target_scan( + target_source, + &target_label, + target_variable, + target_properties, + )?; + builder = builder + .join_on( + scan, + datafusion::logical_expr::JoinType::Inner, + [source + .eq(col(&target)) + .and(datafusion::logical_expr::lit(same_label))], + ) + .map_err(|e| self.plan_error("Failed to join zero-hop identity", e))?; + } + return builder + .build() + .map_err(|e| self.plan_error("Failed to build zero-hop identity", e)); + } let mut current_plan = input_plan; let mut current_source = source_variable.to_string(); let mut path_nodes = vec![current_source.clone()]; @@ -604,6 +677,130 @@ impl DataFusionPlanner { .map_err(|e| self.plan_error("Failed to build undirected path", e)) } + #[allow(clippy::too_many_arguments)] + fn try_build_csr_traversal( + &self, + ctx: &mut PlanningContext, + input: LogicalPlan, + source_variable: &str, + target_variable: &str, + types: &[String], + direction: &RelationshipDirection, + min_hops: u32, + max_hops: u32, + target_properties: &HashMap, + ) -> Result> { + use crate::csr_expand::{CsrExpandNode, Neighbors}; + use datafusion::logical_expr::Extension; + if max_hops == 0 { + return Ok(None); + } + let Some(catalog) = &self.adjacency else { + return Ok(None); + }; + let (source_label, source_map) = self.get_target_node_mapping(ctx, source_variable)?; + let (target_label, target_map) = self.get_target_node_mapping(ctx, target_variable)?; + let Some(walk) = crate::csr_traversal::Traversal::new( + &self.config, + catalog, + types, + direction, + &source_label, + &target_label, + min_hops, + max_hops, + ) else { + return Ok(None); + }; + let rel_type = &types[0]; + let rel_map = self.get_relationship_mapping(rel_type)?; + let source_key = qualify_column(source_variable, &source_map.id_field); + if !input + .schema() + .field_with_unqualified_name(&source_key) + .is_ok_and(|f| crate::csr_expand::is_supported_id_type(f.data_type())) + { + return Ok(None); + } + let adjacency = match direction { + RelationshipDirection::Outgoing => Neighbors::directed( + catalog + .get(rel_type, true) + .expect("validated traversal") + .clone(), + ), + RelationshipDirection::Incoming => Neighbors::directed( + catalog + .get(rel_type, false) + .expect("validated traversal") + .clone(), + ), + RelationshipDirection::Undirected => Neighbors::undirected( + catalog + .get(rel_type, true) + .expect("validated traversal") + .clone(), + catalog + .get(rel_type, false) + .expect("validated traversal") + .clone(), + ), + }; + let cat = self.get_catalog()?; + let target = target_properties + .is_empty() + .then(|| { + self.csr_target_lookup(cat, &input, target_variable, &target_label, target_map) + }) + .flatten(); + let fused = target.is_some(); + let bound_target = qualify_column(target_variable, &target_map.id_field); + let closing = input + .schema() + .field_with_unqualified_name(&bound_target) + .is_ok(); + let alias = ctx.next_relationship_instance(rel_type)?.alias.clone(); + let output = |name| { + arrow_schema::Field::new( + qualify_column(&alias, name), + arrow_schema::DataType::Int64, + true, + ) + }; + let mut node = CsrExpandNode::try_new( + input, + source_key, + output(Self::get_source_join_key(direction, rel_map)), + output(Self::get_target_join_key(direction, rel_map)), + adjacency, + target, + )? + .with_traversal(walk); + if closing && target_properties.is_empty() { + node = node.with_closing_column(bound_target); + } + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(node), + }); + if fused || (closing && target_properties.is_empty()) { + return Ok(Some(plan)); + } + self.join_relationship_to_target( + LogicalPlanBuilder::from(plan), + cat, + &TargetJoinParams { + target_variable, + target_label: &target_label, + rel_qualifier: &alias, + node_map: target_map, + rel_map, + direction, + target_properties, + }, + ) + .map(Some) + } + /// Build a single-hop expansion on top of an existing plan (one hop of a /// variable-length path). Uses the same planning as a single hop, so a hop gets /// the CSR expansion whenever a plain `(a)-[:R]->(b)` would. @@ -712,6 +909,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: Some("r".into()), @@ -747,6 +945,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, @@ -782,6 +981,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, @@ -831,6 +1031,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, @@ -879,6 +1080,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, @@ -928,6 +1130,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, @@ -977,6 +1180,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, @@ -1026,6 +1230,7 @@ mod tests { input: Box::new(scan_a), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: crate::ast::RelationshipDirection::Outgoing, relationship_variable: None, diff --git a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs index 7044c3e..df3acd3 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs @@ -144,14 +144,12 @@ impl DataFusionPlanner { } LogicalOperator::VariableLengthExpand { input, - source_variable, target_variable, + target_label, .. } => { Self::collect_labels(input, labels); - if let Some(label) = labels.get(source_variable).cloned() { - labels.insert(target_variable.clone(), label); - } + labels.insert(target_variable.clone(), target_label.clone()); } LogicalOperator::Filter { input, .. } | LogicalOperator::Project { input, .. } @@ -319,6 +317,7 @@ mod tests { input: Box::new(scan), source_variable: "a".into(), target_variable: "b".into(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".into()], direction: RelationshipDirection::Outgoing, relationship_variable: None, diff --git a/crates/lance-graph/src/datafusion_planner/join_ops.rs b/crates/lance-graph/src/datafusion_planner/join_ops.rs index 61282b2..3fa0d8a 100644 --- a/crates/lance-graph/src/datafusion_planner/join_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/join_ops.rs @@ -43,7 +43,7 @@ pub(crate) struct TargetJoinParams<'a> { impl DataFusionPlanner { /// Build a target node scan with property filters and qualified column names - fn build_target_scan( + pub(crate) fn build_target_scan( &self, target_source: Arc, target_label: &str, diff --git a/crates/lance-graph/src/datafusion_planner/scan_ops.rs b/crates/lance-graph/src/datafusion_planner/scan_ops.rs index 5a5466f..33be41a 100644 --- a/crates/lance-graph/src/datafusion_planner/scan_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/scan_ops.rs @@ -355,6 +355,7 @@ mod tests { input: Box::new(scan), source_variable: "foo_bar".to_string(), // Will generate _temp_foo_bar_1 target_variable: "baz".to_string(), + target_label: "Person".to_string(), relationship_types: vec!["KNOWS".to_string()], direction: crate::ast::RelationshipDirection::Outgoing, min_length: Some(2), diff --git a/crates/lance-graph/src/lib.rs b/crates/lance-graph/src/lib.rs index f7746e0..4606eac 100644 --- a/crates/lance-graph/src/lib.rs +++ b/crates/lance-graph/src/lib.rs @@ -41,6 +41,7 @@ pub mod case_insensitive; pub mod config; pub mod csr_expand; pub mod csr_index; +mod csr_traversal; pub mod datafusion_planner; pub mod error; pub mod key_statistics; diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index 339e842..33dd3f0 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -80,7 +80,7 @@ pub enum LogicalOperator { /// Variable-length path expansion (*1..2 syntax) /// /// Represents multi-hop relationship traversals: (source)-[rel*min..max]->(target) - /// This is implemented by unrolling into multiple fixed-length paths and unioning them + /// Uses iterative CSR traversal when indexed, otherwise fixed-length branches and a union. VariableLengthExpand { /// The input operator (typically a node scan) input: Box, @@ -88,6 +88,8 @@ pub enum LogicalOperator { source_variable: String, /// Variable name for the target node (reachable in min..max hops) target_variable: String, + /// The label of the final node, which may differ from intermediate nodes. + target_label: String, /// Types of relationships to traverse in each hop relationship_types: Vec, /// Direction of traversal for each hop @@ -232,6 +234,7 @@ pub struct LogicalPlanner<'a> { variables: HashMap, // variable -> label config: &'a GraphConfig, fan_out: Option<&'a FanOut>, + adjacency: Option<&'a crate::csr_expand::AdjacencyCatalog>, /// 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. @@ -253,6 +256,7 @@ impl<'a> LogicalPlanner<'a> { variables: HashMap::new(), config, fan_out: None, + adjacency: None, exists_marks: 0, node_lists: HashMap::new(), comprehensions: 0, @@ -268,6 +272,73 @@ impl<'a> LogicalPlanner<'a> { self } + pub(crate) fn with_adjacency( + mut self, + adjacency: Option<&'a crate::csr_expand::AdjacencyCatalog>, + ) -> Self { + self.adjacency = adjacency; + self + } + + fn retain_variable_alternation(&self, patterns: &[GraphPattern]) -> bool { + let Some(catalog) = self.adjacency else { + return false; + }; + let labels = |n: &NodePattern| -> Vec { + n.variable + .as_ref() + .and_then(|v| self.variables.get(v)) + .map(|l| vec![l.clone()]) + .unwrap_or_else(|| n.labels.clone()) + }; + patterns.iter().all(|p| { + let GraphPattern::Path(path) = p else { + return true; + }; + let mut previous = &path.start_node; + path.segments.iter().all(|segment| { + let start = previous; + previous = &segment.end_node; + let rel = &segment.relationship; + let Some(range) = &rel.length else { + return true; + }; + if rel.types.len() <= 1 { + return true; + } + if rel.variable.is_some() + || !rel.properties.is_empty() + || rel.direction == RelationshipDirection::Undirected + || (range.min == Some(1) && range.max == Some(1)) + { + return false; + } + let (starts, ends) = (labels(start), labels(&segment.end_node)); + let (Some(start), Some(end)) = (starts.first(), ends.first()) else { + return false; + }; + if starts + .iter() + .chain(&ends) + .any(|l| catalog.nodes(l).is_none()) + { + return false; + } + crate::csr_traversal::Traversal::new( + self.config, + catalog, + &rel.types, + &rel.direction, + start, + end, + 0, + 0, + ) + .is_some() + }) + }) + } + /// For a path that continues earlier rows, the path read from the other end when that /// is the better start: always when only the far end is bound (expansion must start /// from a bound node), and, with both ends bound, when the far end's first hop has the @@ -432,7 +503,9 @@ impl<'a> LogicalPlanner<'a> { } for branch in &query.union_branches { - let mut planner = LogicalPlanner::new(self.config).with_fan_out(self.fan_out); + let mut planner = LogicalPlanner::new(self.config) + .with_fan_out(self.fan_out) + .with_adjacency(self.adjacency); planner.calls = self.calls; if preserve_entities { planner.variables = imports.0.clone(); @@ -522,7 +595,9 @@ impl<'a> LogicalPlanner<'a> { let imports = query.call_union_imports()?; self.calls += 1; let id = self.calls; - let mut inner = Self::new(self.config).with_fan_out(self.fan_out); + let mut inner = Self::new(self.config) + .with_fan_out(self.fan_out) + .with_adjacency(self.adjacency); inner.calls = self.calls; for name in &imports { if let Some(label) = self.variables.get(name) { @@ -714,11 +789,13 @@ impl<'a> LogicalPlanner<'a> { ..match_clause.clone() }; let match_clause = &inferred; - if crate::pattern_alternatives::needs_expansion(&match_clause.patterns) { + let retain_variable = self.retain_variable_alternation(&match_clause.patterns); + if crate::pattern_alternatives::needs_expansion(&match_clause.patterns, retain_variable) { let branches = crate::pattern_alternatives::expand( &match_clause.patterns, self.config, &self.variables, + retain_variable, )?; return self.plan_match_branches(base, branches, filter); } @@ -1393,6 +1470,7 @@ impl<'a> LogicalPlanner<'a> { input: Box::new(plan), source_variable: current_src.clone(), target_variable: target_variable.clone(), + target_label: target_label.clone(), relationship_types: segment.relationship.types.clone(), direction: segment.relationship.direction.clone(), relationship_variable: segment.relationship.variable.clone(), diff --git a/crates/lance-graph/src/pattern_alternatives.rs b/crates/lance-graph/src/pattern_alternatives.rs index f2d1356..64b37df 100644 --- a/crates/lance-graph/src/pattern_alternatives.rs +++ b/crates/lance-graph/src/pattern_alternatives.rs @@ -31,6 +31,12 @@ use crate::error::{GraphError, Result}; /// Branch limit, so a pattern full of alternatives can't explode the plan. const MAX_BRANCHES: usize = 64; +fn zero_hop(rel: &RelationshipPattern) -> bool { + rel.length + .as_ref() + .is_some_and(|r| r.min == Some(0) && r.max == Some(0)) +} + /// `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 @@ -149,16 +155,17 @@ pub fn infer_labels( } /// True when any relationship has several types or any node several labels. -pub fn needs_expansion(patterns: &[GraphPattern]) -> bool { +pub fn needs_expansion(patterns: &[GraphPattern], retain_variable: bool) -> bool { let nodes_multi = |n: &NodePattern| n.labels.len() > 1; patterns.iter().any(|p| match p { GraphPattern::Node(n) => nodes_multi(n), GraphPattern::Path(path) => { nodes_multi(&path.start_node) - || path - .segments - .iter() - .any(|s| s.relationship.types.len() > 1 || nodes_multi(&s.end_node)) + || path.segments.iter().any(|s| { + (s.relationship.types.len() > 1 + && !(retain_variable && s.relationship.length.is_some())) + || nodes_multi(&s.end_node) + }) } }) } @@ -178,11 +185,16 @@ pub fn expand( patterns: &[GraphPattern], config: &GraphConfig, bound: &HashMap, + retain_variable: bool, ) -> Result>> { - if let Some(variants) = unroll_variable_length_alternation(patterns, config, bound)? { + if let Some(variants) = if retain_variable { + None + } else { + unroll_variable_length_alternation(patterns, config, bound)? + } { let mut branches = Vec::new(); for variant in variants { - branches.extend(expand(&variant, config, bound)?); + branches.extend(expand(&variant, config, bound, retain_variable)?); if branches.len() > MAX_BRANCHES { return Err(too_many()); } @@ -227,12 +239,14 @@ pub fn expand( for segment in &path.segments { let rel = &segment.relationship; let right = slot(&segment.end_node, &mut allowed); - rels.push(Rel { - left, - right, - direction: rel.direction.clone(), - types: rel.types.clone(), - }); + if !(zero_hop(rel) || retain_variable && rel.length.is_some()) { + rels.push(Rel { + left, + right, + direction: rel.direction.clone(), + types: rel.types.clone(), + }); + } left = right; } } @@ -338,7 +352,13 @@ pub fn expand( } }); if consistent { - branches.push(rewrite(patterns, &occurrence_slots, &labels, &types)); + branches.push(rewrite( + patterns, + &occurrence_slots, + &labels, + &types, + retain_variable, + )); } } if branches.is_empty() { @@ -414,7 +434,7 @@ fn unroll_variable_length_alternation( labels.is_empty() || labels.iter().any(|l| l.eq_ignore_ascii_case(label)) }; let range = rel.length.as_ref().expect("a variable-length segment"); - let min = range.min.unwrap_or(1).max(1); + let min = range.min.unwrap_or(1); let max = range.max.unwrap_or(crate::MAX_VARIABLE_LENGTH_HOPS); if max > crate::MAX_VARIABLE_LENGTH_HOPS { return Err(unsupported(&format!( @@ -449,43 +469,52 @@ fn unroll_variable_length_alternation( } } } - if sequences.is_empty() { + if sequences.is_empty() && min != 0 { return Err(plan_error( "no sequence of relationship types in the variable-length pattern chains its \ endpoint labels" .into(), )); } - let variants = sequences - .into_iter() - .map(|sequence| { - let last = sequence.len() - 1; - let hops = sequence - .into_iter() - .enumerate() - .map(|(h, (t, label))| PathSegment { - relationship: RelationshipPattern { - types: vec![t], - length: None, - ..rel.clone() - }, - end_node: if h == last { - path.segments[i].end_node.clone() - } else { - NodePattern { - variable: None, - labels: vec![label], - properties: HashMap::new(), - } - }, - }); - let mut unrolled = path.clone(); - unrolled.segments.splice(i..=i, hops); - let mut variant = patterns.to_vec(); - variant[p] = GraphPattern::Path(unrolled); - variant - }) - .collect(); + let mut variants = Vec::new(); + if min == 0 { + let mut identity = path.clone(); + identity.segments[i].relationship.types = vec![rel.types[0].clone()]; + identity.segments[i].relationship.length = Some(crate::ast::LengthRange { + min: Some(0), + max: Some(0), + }); + let mut variant = patterns.to_vec(); + variant[p] = GraphPattern::Path(identity); + variants.push(variant); + } + variants.extend(sequences.into_iter().map(|sequence| { + let last = sequence.len() - 1; + let hops = sequence + .into_iter() + .enumerate() + .map(|(h, (t, label))| PathSegment { + relationship: RelationshipPattern { + types: vec![t], + length: None, + ..rel.clone() + }, + end_node: if h == last { + path.segments[i].end_node.clone() + } else { + NodePattern { + variable: None, + labels: vec![label], + properties: HashMap::new(), + } + }, + }); + let mut unrolled = path.clone(); + unrolled.segments.splice(i..=i, hops); + let mut variant = patterns.to_vec(); + variant[p] = GraphPattern::Path(unrolled); + variant + })); Ok(Some(variants)) } @@ -515,6 +544,7 @@ fn rewrite( occurrence_slots: &[usize], labels: &[Option<&str>], types: &[&str], + retain_variable: bool, ) -> Vec { let mut occurrences = occurrence_slots.iter(); let mut types = types.iter(); @@ -532,8 +562,12 @@ fn rewrite( let mut path = path.clone(); path.start_node = relabel(&path.start_node); for segment in &mut path.segments { - segment.relationship.types = - vec![types.next().expect("one type per relationship").to_string()]; + if !(zero_hop(&segment.relationship) + || retain_variable && segment.relationship.length.is_some()) + { + segment.relationship.types = + vec![types.next().expect("one type per relationship").to_string()]; + } segment.end_node = relabel(&segment.end_node); } GraphPattern::Path(path) @@ -588,7 +622,7 @@ mod tests { let ReadingClause::Match(m) = &ast.reading_clauses[0] else { panic!("not a MATCH") }; - let out = expand(&m.patterns, &ldbc_config(), &HashMap::new())?; + let out = expand(&m.patterns, &ldbc_config(), &HashMap::new(), false)?; Ok(out.iter().map(|b| describe(b)).collect()) } diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 2b5f8fc..082adb6 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -1036,7 +1036,9 @@ impl CypherQuery { } // Phase 2: Graph Logical Plan - let mut logical_planner = LogicalPlanner::new(config).with_fan_out(self.fan_out.as_ref()); + let mut logical_planner = LogicalPlanner::new(config) + .with_fan_out(self.fan_out.as_ref()) + .with_adjacency(self.adjacency.as_ref()); let ast = if parameterized { &self.ast } else { diff --git a/docs/csr.md b/docs/csr.md index f819ec6..e8d746c 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -121,7 +121,8 @@ When those hold, it covers: | --- | --- | | `(a)-[:KNOWS]->(b)`, `(a)<-[:KNOWS]-(b)` | :material-check: outgoing or incoming adjacency | | `(a)-[:KNOWS]-(b)` | :material-check: both adjacencies | -| `(a)-[:KNOWS*1..3]->(b)` | :material-check: every unrolled hop | +| `(a)-[:KNOWS*1..3]->(b)` | :material-check: one iterative expansion for the whole range | +| (t:Tag)-[:hasType|isSubclassOf*1..3]->(c:TagClass) | :material-check: typed iterative expansion when every type and label is indexed | | (a)-[:LIKES_POST|LIKES_COMMENT]->(m) | :material-check: one expansion per branch | | `OPTIONAL MATCH (a)-[:KNOWS]->(b)` | :material-check: | | `shortestPath((a)-[:KNOWS*]-(b))` | :material-check: the search walks the CSR adjacency | @@ -131,6 +132,23 @@ When those hold, it covers: Everything else in the query (filters, aggregations, sorts and any hop the CSR can't take) is planned by DataFusion as usual. You can mix: a three-hop query might run two hops on the CSR and one as a join. +### Variable-length traversal + +An eligible variable-length path keeps one cursor per hop and yields bounded output batches. Its plan +contains one `CsrExpandExec` with `iterative hops=...`, rather than a branch for every possible length or +sequence of relationship types. It preserves duplicate matches and skips missing nodes at intermediate +hops. A node's label and key travel together, so the same integer key in two labels identifies two nodes. + +The iterative path needs CSR indexes for every traversed type and direction, node indexes for every +intermediate label, and resolved labels for the endpoints. Unsupported index shapes retain the join +fallback. The 20-hop maximum still applies. `*0..` includes a zero-hop identity match when both labels and +keys agree. Undirected paths keep the existing rule that rejects immediate backtracking; neither execution +path enforces full relationship uniqueness on longer cycles. + +Traversal state grows with the hop limit, not the number of pending paths. Resident graph indexes, input +batches, fetched target properties and downstream operators still contribute to query memory. Branching +graphs can produce exponentially many matches even though traversal resumes between output batches. + !!! tip "Leave relationship variables out unless you need them" `MATCH (a)-[k:KNOWS]->(b) RETURN b.name` falls back to joins only because it names `k`. If you don't read diff --git a/docs/cypher.md b/docs/cypher.md index 2f04416..8ec7c1b 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -281,8 +281,13 @@ MATCH p = shortestPath((a:Person {id: 1})-[:KNOWS*]-(b:Person {id: 4})) RETURN length(p), nodes(p) ``` -- **Variable length** defaults to a minimum of 1 and a maximum of 20 hops; a maximum above 20 is rejected, - and zero-length paths (`*0..`) are not supported. Each length becomes one branch of a union. +- **Variable length** defaults to a minimum of 1 and a maximum of 20 hops; a maximum above 20 is rejected. + `*0..` also matches the source itself, once per input row; `*0..0` matches only that identity. The start + and end must have the same label and key for a zero-hop match. Eligible CSR paths use one iterative + expansion for the whole range; the join fallback uses one branch per length. +- **Variable-length alternation** follows each relationship type's configured endpoint labels. With + suitable CSR and node indexes, a mixed path such as `[:hasType|isSubclassOf*1..3]` stays in one operator. + Undirected or named variable-length alternation remains unsupported. - **Undirected variable-length paths** skip immediate backtracking (`a-b-a`). - **`shortestPath` and `allShortestPaths`** need exactly one relationship type, no relationship variable, and named end nodes. Their length is unbounded unless you give a maximum. `nodes(p)` returns the node keys. diff --git a/python/python/tests/test_engine_alternation.py b/python/python/tests/test_engine_alternation.py index 2d8d9fa..09e632a 100644 --- a/python/python/tests/test_engine_alternation.py +++ b/python/python/tests/test_engine_alternation.py @@ -128,6 +128,51 @@ def test_variable_length_alternation_follows_endpoint_labels(datasets, use_csr): assert rows(got) == [(20, 10), (21, 10)] +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize("minimum", [0, 1]) +def test_variable_length_mixed_labels_keep_identity_separate(use_csr, minimum): + builder = GraphConfig.builder().node_table("Tag", "id").node_table("TagClass", "id") + cfg = ( + builder.rel_table("hasType", "Tag", "TagClass", "src", "dst") + .rel_table("isSubclassOf", "TagClass", "TagClass", "src", "dst") + .build() + ) + tables = { + "Tag": pa.table({"id": [1, 2]}), + "TagClass": pa.table({"id": [1, 2, 3], "name": ["one", "two", "three"]}), + "hasType": pa.table({"src": [1, 1, 2, 1], "dst": [1, 1, 2, 99]}), + "isSubclassOf": pa.table({"src": [1, 2], "dst": [2, 3]}), + } + engine = CypherEngine(cfg, tables, use_csr=use_csr) + query = ( + f"MATCH (t:Tag {{id: 1}})-[:hasType|isSubclassOf*{minimum}..3]->(c:TagClass) " + "RETURN c.id, c.name" + ) + assert rows(engine.execute(query)) == [ + (1, "one"), + (1, "one"), + (2, "two"), + (2, "two"), + (3, "three"), + (3, "three"), + ] + if use_csr: + plan = engine.explain(query) + physical = plan[plan.index("physical_plan") :] + assert physical.count("CsrExpandExec") == 1 + assert "UnionExec" not in physical + # Equal integer keys in different labels never form a zero-hop path. + identity = ( + "MATCH (t:Tag {id: 1})-[:hasType|isSubclassOf*0..0]->(c:TagClass) RETURN c.id" + ) + assert rows(engine.execute(identity)) == [] + same_label = ( + "MATCH (t:TagClass {id: 1})-[:hasType|isSubclassOf*0..2]->(c:TagClass) " + "RETURN c.id" + ) + assert rows(engine.execute(same_label)) == [(1,), (2,), (3,)] + + def test_collect_distinct(datasets): engine = CypherEngine(config(), datasets) got = engine.execute( diff --git a/python/python/tests/test_engine_csr.py b/python/python/tests/test_engine_csr.py index b5a4dd7..b96a8b9 100644 --- a/python/python/tests/test_engine_csr.py +++ b/python/python/tests/test_engine_csr.py @@ -119,11 +119,39 @@ def test_variable_length_hops_use_csr(datasets): plan = csr.explain("MATCH (a:Person)-[:KNOWS*1..2]-(b:Person) RETURN b.id") lines = plan[plan.index("physical_plan") :].splitlines() expands = [line for line in lines if "CsrExpandExec" in line] - # One expansion for the 1-hop branch and two for the 2-hop branch. - assert len(expands) == 3 and all("undirected" in line for line in expands) + assert len(expands) == 1 and "iterative hops=1..2" in expands[0] + assert "undirected" in expands[0] + assert "UnionExec" not in "\n".join(lines) assert "HashJoinExec" not in "\n".join(lines) +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize("hops", ["0..0", "0..2"]) +def test_variable_length_zero_hop_is_identity(datasets, use_csr, hops): + engine = CypherEngine(CONFIG, datasets, use_csr=use_csr) + got = engine.execute( + f"MATCH (a:Person {{id: 1}})-[:KNOWS*{hops}]->(b:Person) " + "RETURN b.id AS id, b.name AS name" + ).to_pylist() + expected = [(1, "a")] + if hops == "0..2": + expected += [(2, "b")] * 2 + [(3, "c")] * 4 + assert sorted((r["id"], r["name"]) for r in got) == sorted(expected) + + +def test_variable_length_plan_does_not_grow_with_maximum(datasets): + engine = CypherEngine(CONFIG, datasets, use_csr=True) + for maximum in [2, 5, 20]: + plan = engine.explain( + f"MATCH (a:Person)-[:KNOWS*1..{maximum}]->(b:Person) RETURN b.id" + ) + physical = plan[plan.index("physical_plan") :] + assert physical.count("CsrExpandExec") == 1 + assert "UnionExec" not in physical + with pytest.raises(NotImplementedError, match="20"): + engine.execute("MATCH (a:Person)-[:KNOWS*1..21]->(b:Person) RETURN b.id") + + def test_source_filter_pushed_below_csr(datasets): csr = CypherEngine(CONFIG, datasets, use_csr=True) plan = csr.explain( From e307822daf802ea4b1c0ff623257ceebbff5f220 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:36:34 -0700 Subject: [PATCH 22/66] perf(csr): reuse filtered target columns within each batch Resolve #16 by retaining filter columns needed by the projection and reconstructing target IDs in predicates. Fetch only remaining properties for surviving rows, without a persistent property cache. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/src/csr_expand.rs | 182 ++++++++++++++++++++++++- docs/csr.md | 6 + python/benchmarks/csr_expand.py | 4 + python/python/tests/test_engine_csr.py | 5 + 4 files changed, 191 insertions(+), 6 deletions(-) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index d3fad84..aeb1e71 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -35,7 +35,7 @@ use std::sync::Arc; use arrow::array::{ new_empty_array, Array, ArrayRef, AsArray, Int64Array, RecordBatch, UInt32Array, UInt64Array, }; -use arrow::compute::{cast, take}; +use arrow::compute::{cast, filter as filter_array, take}; use arrow::datatypes::{DataType, Field, Int64Type, Schema, SchemaRef}; use async_trait::async_trait; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; @@ -502,6 +502,8 @@ pub struct NodeRows { schema: SchemaRef, /// Index of the id column in `schema`; its values are reconstructed, never fetched. id_index: usize, + #[cfg(test)] + fetched_cells: std::sync::atomic::AtomicUsize, } enum NodeData { @@ -544,6 +546,8 @@ impl NodeRows { data: NodeData::Lance(dataset), schema, id_index, + #[cfg(test)] + fetched_cells: std::sync::atomic::AtomicUsize::new(0), })) } @@ -562,6 +566,8 @@ impl NodeRows { schema: batch.schema(), data: NodeData::Memory(batch), id_index, + #[cfg(test)] + fetched_cells: std::sync::atomic::AtomicUsize::new(0), })) } @@ -585,6 +591,11 @@ impl NodeRows { .map(|&c| new_empty_array(self.schema.field(c).data_type())) .collect()); } + #[cfg(test)] + self.fetched_cells.fetch_add( + positions.len() * columns.len(), + std::sync::atomic::Ordering::Relaxed, + ); match &self.data { NodeData::Lance(dataset) => { let names = columns.iter().map(|&c| self.schema.field(c).name()); @@ -1297,15 +1308,40 @@ async fn materialize( mut neighbor_values, mut target_positions, } = expanded; + let mut reused = Vec::new(); if let (Some(target), Some(filter)) = (target, target_filter) { // Take only the filter's columns, evaluate, and keep passing rows before // taking anything else. Null counts as false, as in a WHERE clause. - let taken = target + let fetch: Vec<_> = filter + .columns + .iter() + .copied() + .filter(|&c| c != target.rows.id_index) + .collect(); + let mut taken = target .rows - .take(&target_positions, &filter.columns) + .take(&target_positions, &fetch) .await - .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))?; - let candidates = RecordBatch::try_new(filter.schema.clone(), taken)?; + .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))? + .into_iter(); + let arrays = filter + .columns + .iter() + .enumerate() + .map(|(i, &c)| { + if c == target.rows.id_index { + cast( + &Int64Array::from(neighbor_values.clone()), + filter.schema.field(i).data_type(), + ) + } else { + Ok(taken + .next() + .expect("one fetched array per non-id filter column")) + } + }) + .collect::, _>>()?; + let candidates = RecordBatch::try_new(filter.schema.clone(), arrays)?; let mask = filter .predicate .evaluate(&candidates)? @@ -1314,6 +1350,11 @@ async fn materialize( let keep: Vec = (0..mask.len()) .map(|i| mask.is_valid(i) && mask.value(i)) .collect(); + for (i, &column) in filter.columns.iter().enumerate() { + if column != target.rows.id_index && target.columns.contains(&column) { + reused.push((column, filter_array(candidates.column(i), mask)?)); + } + } retain_by_mask(&mut rows, &keep); retain_by_mask(&mut key_values, &keep); retain_by_mask(&mut neighbor_values, &keep); @@ -1338,7 +1379,7 @@ async fn materialize( .columns .iter() .copied() - .filter(|&c| c != id_index) + .filter(|&c| c != id_index && !reused.iter().any(|(column, _)| *column == c)) .collect(); let mut fetched = target .rows @@ -1350,6 +1391,8 @@ async fn materialize( // The target id equals the neighbour value, so it is never fetched. columns.push(if c == id_index { cast(&neighbor_values, field.data_type())? + } else if let Some((_, array)) = reused.iter().find(|(column, _)| *column == c) { + array.clone() } else { fetched.next().expect("one fetched array per non-id column") }); @@ -2065,4 +2108,131 @@ mod tests { assert_eq!(both.neighbors_from(2, &mut position).next(), Some(1)); assert_eq!(both.neighbors_from(2, &mut position).next(), None); } + + #[tokio::test] + async fn materialization_reuses_filter_columns_for_memory_and_lance() { + use arrow::array::{RecordBatchIterator, StringArray}; + use datafusion::logical_expr::{col, lit}; + use std::sync::atomic::Ordering; + + let nodes = RecordBatch::try_from_iter([ + ( + "id", + Arc::new(Int64Array::from(vec![10, 20, 30, 40])) as ArrayRef, + ), + ( + "age", + Arc::new(Int64Array::from(vec![Some(2), None, Some(5), Some(1)])) as ArrayRef, + ), + ( + "name", + Arc::new(StringArray::from(vec!["ten", "twenty", "thirty", "forty"])) as ArrayRef, + ), + ( + "unused", + Arc::new(StringArray::from(vec!["unused payload"; 4])) as ArrayRef, + ), + ]) + .unwrap(); + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("nodes.lance"); + let reader = RecordBatchIterator::new(vec![Ok(nodes.clone())], nodes.schema()); + let dataset = Dataset::write( + reader, + path.to_str().unwrap(), + None::, + ) + .await + .unwrap(); + let sources = [ + NodeRows::from_batch(nodes.clone(), "id").unwrap().unwrap(), + NodeRows::from_lance(Arc::new(dataset), "id") + .await + .unwrap() + .unwrap(), + ]; + let input = RecordBatch::try_from_iter([( + "a__id", + Arc::new(Int64Array::from(vec![999])) as ArrayRef, + )]) + .unwrap(); + let state = datafusion::execution::context::SessionContext::new().state(); + for source in sources { + let source = Arc::new(source); + for selected in [vec![0, 1, 2], vec![0, 1], vec![2], vec![]] { + let fields: Vec<_> = selected + .iter() + .map(|&i| { + let f = nodes.schema().field(i).clone(); + Field::new( + format!("b__{}", f.name()), + f.data_type().clone(), + f.is_nullable(), + ) + }) + .collect(); + let target = TargetLookup { + variable: "b".into(), + rows: source.clone(), + columns: selected.clone(), + fields: fields.clone(), + }; + let schema = Arc::new(Schema::new( + [ + Field::new("a__id", DataType::Int64, false), + Field::new("edge__src", DataType::Int64, false), + Field::new("edge__dst", DataType::Int64, false), + ] + .into_iter() + .chain(fields) + .collect::>(), + )); + for threshold in [2, 99] { + source.fetched_cells.store(0, Ordering::Relaxed); + let predicate = col("b__age") + .gt_eq(lit(threshold)) + .and(col("b__id").gt_eq(lit(10))); + let filter = target.compile_filter(&predicate, &state).unwrap(); + let neighbors = vec![30, 10, 20, 30, 40, 10]; + let expanded = ExpandedRows { + rows: vec![0; 6], + key_values: vec![999; 6], + target_positions: neighbors + .iter() + .map(|&n| source.position(n).unwrap()) + .collect(), + neighbor_values: neighbors, + }; + let out = materialize(&input, expanded, Some(&target), Some(&filter), &schema) + .await + .unwrap(); + let expected_rows: Vec = if threshold == 2 { + vec![2, 0, 2, 0] + } else { + vec![] + }; + assert_eq!(out.num_rows(), expected_rows.len()); + for (i, &column) in selected.iter().enumerate() { + let expected = take( + nodes.column(column), + &UInt32Array::from(expected_rows.clone()), + None, + ) + .unwrap(); + assert_eq!(out.column(3 + i).to_data(), expected.to_data()); + } + // Only age is read for six candidates; only name needs a later read. + let projection_cells = if selected.contains(&2) { + expected_rows.len() + } else { + 0 + }; + assert_eq!( + source.fetched_cells.load(Ordering::Relaxed), + 6 + projection_cells + ); + } + } + } + } } diff --git a/docs/csr.md b/docs/csr.md index e8d746c..4f72c39 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -169,6 +169,12 @@ print([line.strip(" |") for line in plan.splitlines() if "CsrExpandExec" in line The fused lookup needs the target's key to be a **unique integer**. Duplicate keys, or a target that is already bound earlier in the query, keep the CSR hop but join for the target's properties. +Only requested target columns are read. A pushed-down target filter reads its own columns first; remaining +properties are fetched only for neighbours that pass. When a column is used by both the filter and the +projection, its filtered values are reused within that output batch. Target keys come directly from the +neighbour IDs, including when a filter reads them. This reuse ends with the batch and keeps no property +cache across queries or dataset versions. + ### Patterns start at the most selective node The planner rewrites each `MATCH` so traversal starts at its most selective node, walking edges backwards diff --git a/python/benchmarks/csr_expand.py b/python/benchmarks/csr_expand.py index e7d591c..6ef4e92 100644 --- a/python/benchmarks/csr_expand.py +++ b/python/benchmarks/csr_expand.py @@ -34,6 +34,7 @@ # Same shapes returning a non-id target column, so the target rows are fetched. "point_name": "MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE a.id = 42 RETURN b.name", "selective_1pct_name": None, + "selective_1pct_name_filtered": None, "all_edges": "MATCH (a:Person)-[:KNOWS]->(b:Person) RETURN count(*) AS n", } @@ -85,6 +86,9 @@ def child(root: Path, mode: str, repeats: int) -> dict: f"MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE a.id < {nodes // 100} RETURN b.id" ) queries["selective_1pct_name"] = queries["selective_1pct"].replace("b.id", "b.name") + queries["selective_1pct_name_filtered"] = queries["selective_1pct_name"].replace( + "RETURN", "AND b.name STARTS WITH 'n1' RETURN" + ) def open_engine(): datasets = { diff --git a/python/python/tests/test_engine_csr.py b/python/python/tests/test_engine_csr.py index b96a8b9..3e98fe7 100644 --- a/python/python/tests/test_engine_csr.py +++ b/python/python/tests/test_engine_csr.py @@ -52,6 +52,11 @@ # Target filter moved into the CSR node (start is more selective than b). "MATCH (a:Person)-[:KNOWS]->(b:Person) " "WHERE a.id = 1 AND b.age >= 40 RETURN b.name", + # Reuse a filter column in the projection, including reconstructed target ids. + "MATCH (a:Person)-[:KNOWS]->(b:Person) " + "WHERE a.id = 1 AND b.id >= 2 AND b.age >= 30 RETURN b.id, b.age, b.name", + "MATCH (a:Person)-[:KNOWS]->(b:Person) " + "WHERE a.id = 1 AND b.age > 100 RETURN b.age, b.name", # Closing edges: cycles back to an already-bound variable (duplicates included). "MATCH (a:Person)-[:KNOWS]->(b:Person)-[:KNOWS]->(a) RETURN a.id, b.id", "MATCH (a:Person)-[:KNOWS]->(b:Person), (b)<-[:KNOWS]-(a) RETURN count(*) AS n", From 087035084a376a938526e245f3fb6f07c64ad355 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:08:40 -0700 Subject: [PATCH 23/66] perf(graph): search shortest paths from both ends Resolve #17 with complete-layer bidirectional search for single-target CSR queries, preserving all shortest paths and duplicate edges. Keep shared forward search for multiple targets. Developed by David Hughes with assistance from Codex. --- .../datafusion_planner/builder/expand_ops.rs | 16 +- crates/lance-graph/src/shortest_path.rs | 266 ++++++++++++++++-- docs/csr.md | 6 + docs/cypher.md | 6 +- python/python/tests/test_engine_paths.py | 42 +++ 5 files changed, 311 insertions(+), 25 deletions(-) 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 26777d3..ed30c57 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/expand_ops.rs @@ -358,6 +358,19 @@ impl DataFusionPlanner { .zip(cat.get(rel_type, false)) .map(|(out, inc)| Neighbors::undirected(out.clone(), inc.clone())), }); + let reverse_adjacency = self + .adjacency + .as_ref() + .filter(|_| adjacency.is_some()) + .and_then(|cat| match search { + SearchDirection::Outgoing => { + cat.get(rel_type, false).cloned().map(Neighbors::directed) + } + SearchDirection::Incoming => { + cat.get(rel_type, true).cloned().map(Neighbors::directed) + } + SearchDirection::Undirected => None, // Forward neighbours already visit both ends. + }); let source = self .get_catalog()? .relationship_source(&rel_map.relationship_type) @@ -400,7 +413,8 @@ impl DataFusionPlanner { edge_source: schema.field(src).name().clone(), edge_target: schema.field(dst).name().clone(), }; - let node = ShortestPathNode::try_new(input_plan, edges, adjacency, spec)?; + let node = ShortestPathNode::try_new(input_plan, edges, adjacency, spec)? + .with_reverse_adjacency(reverse_adjacency); Ok(LogicalPlan::Extension(Extension { node: Arc::new(node), })) diff --git a/crates/lance-graph/src/shortest_path.rs b/crates/lance-graph/src/shortest_path.rs index 2a47450..9c37dec 100644 --- a/crates/lance-graph/src/shortest_path.rs +++ b/crates/lance-graph/src/shortest_path.rs @@ -12,9 +12,9 @@ //! a second input, a scan of the relationship's two endpoint columns, and builds an adjacency //! from it once per query, shared by all partitions. //! -//! One search runs per distinct source id in a batch, level by level, until every target in -//! that batch is reached at its shortest depth (or the depth limit). Predecessors are kept -//! per node, so all shortest paths can be enumerated back from each target. +//! A source with one target searches from both ends when reverse adjacency is available, +//! expanding the smaller complete frontier. Sources with several targets share one +//! forward search. Predecessors retain every shortest path, including parallel edges. use std::cmp::Ordering; use std::collections::{HashMap, HashSet}; @@ -78,6 +78,7 @@ pub struct ShortestPathNode { /// Scan of the relationship's endpoints; `None` when `adjacency` is given. edges: Option, adjacency: Option, + reverse_adjacency: Option, spec: ShortestPathSpec, schema: DFSchemaRef, } @@ -111,10 +112,16 @@ impl ShortestPathNode { input, edges, adjacency, + reverse_adjacency: None, spec, schema: Arc::new(schema), }) } + + pub(crate) fn with_reverse_adjacency(mut self, adjacency: Option) -> Self { + self.reverse_adjacency = adjacency; + self + } } impl PartialEq for ShortestPathNode { @@ -122,6 +129,11 @@ impl PartialEq for ShortestPathNode { self.input == other.input && self.edges == other.edges && self.spec == other.spec + && match (&self.reverse_adjacency, &other.reverse_adjacency) { + (Some(a), Some(b)) => a.ptr_eq(b), + (None, None) => true, + _ => false, + } && match (&self.adjacency, &other.adjacency) { (Some(a), Some(b)) => a.ptr_eq(b), (None, None) => true, @@ -188,6 +200,7 @@ impl UserDefinedLogicalNodeCore for ShortestPathNode { let edges = (inputs.len() > 1).then(|| inputs.pop().expect("two inputs")); let input = inputs.pop().expect("an input"); Self::try_new(input, edges, self.adjacency.clone(), self.spec.clone()) + .map(|node| node.with_reverse_adjacency(self.reverse_adjacency.clone())) .map_err(|e| datafusion::common::DataFusionError::Plan(e.to_string())) } } @@ -200,6 +213,7 @@ pub struct ShortestPathExec { edges: Option>, /// The CSR, or the adjacency built from `edges` by the first batch that needs it. adjacency: Arc>>, + reverse_adjacency: Option, spec: ShortestPathSpec, schema: SchemaRef, properties: Arc, @@ -210,6 +224,7 @@ impl ShortestPathExec { input: Arc, edges: Option>, adjacency: Option, + reverse_adjacency: Option, spec: ShortestPathSpec, schema: SchemaRef, ) -> Self { @@ -223,6 +238,7 @@ impl ShortestPathExec { input, edges, adjacency: Arc::new(Mutex::new(adjacency)), + reverse_adjacency, spec, schema, properties, @@ -277,7 +293,14 @@ impl ExecutionPlan for ShortestPathExec { ) -> datafusion::common::Result> { let edges = (children.len() > 1).then(|| children.pop().expect("two children")); let input = children.pop().expect("a child"); - let mut exec = Self::new(input, edges, None, self.spec.clone(), self.schema.clone()); + let mut exec = Self::new( + input, + edges, + None, + self.reverse_adjacency.clone(), + self.spec.clone(), + self.schema.clone(), + ); exec.adjacency = self.adjacency.clone(); Ok(Arc::new(exec)) } @@ -297,7 +320,9 @@ impl ExecutionPlan for ShortestPathExec { self.schema.clone(), ); let out_schema = schema.clone(); + let reverse_adjacency = self.reverse_adjacency.clone(); let stream = input.then(move |batch| { + let reverse_adjacency = reverse_adjacency.clone(); let (cell, edges, spec, schema, context) = ( cell.clone(), edges.clone(), @@ -318,6 +343,9 @@ impl ExecutionPlan for ShortestPathExec { source_index, target_index, &adjacency, + reverse_adjacency.as_ref().or_else(|| { + (spec.direction == SearchDirection::Undirected).then_some(&adjacency) + }), &spec, &schema, ) @@ -372,22 +400,7 @@ fn search( let mut level = 0; while !frontier.is_empty() && !remaining.is_empty() && level < max_hops { level += 1; - let mut next = Vec::new(); - for &node in &frontier { - adjacency.for_each_neighbor(node, |n| { - let n = n as i64; - match depth.get(&n) { - // Another shortest way into a node first reached at this level. - Some(&d) if d == level => parents.entry(n).or_default().push(node), - Some(_) => {} - None => { - depth.insert(n, level); - parents.insert(n, vec![node]); - next.push(n); - } - } - }); - } + let next = expand_level(adjacency, &frontier, level, &mut depth, &mut parents); // Finish the level before stopping, so every shortest path is recorded. for n in &next { remaining.remove(n); @@ -397,6 +410,93 @@ fn search( (depth, parents) } +/// Finish the whole level so parallel edges and alternative shortest paths survive. +fn expand_level( + adjacency: &Neighbors, + frontier: &[i64], + level: u32, + depth: &mut FastMap, + parents: &mut FastMap>, +) -> Vec { + let mut next = Vec::new(); + for &node in frontier { + adjacency.for_each_neighbor(node, |n| { + let n = n as i64; + match depth.get(&n) { + Some(&d) if d == level => parents.entry(n).or_default().push(node), + Some(_) => {} + None => { + depth.insert(n, level); + parents.insert(n, vec![node]); + next.push(n); + } + } + }); + } + next +} + +/// Search a single pair using complete layers from the smaller frontier. At the +/// first intersection every meeting node lies on the same depth cut of a shortest path. +fn search_bidirectional( + forward: &Neighbors, + reverse: &Neighbors, + source: i64, + target: i64, + max_hops: u32, +) -> (FastMap, FastMap>) { + let mut depth = FastMap::default(); + let mut parents = FastMap::default(); + depth.insert(source, 0); + if source == target { + return (depth, parents); + } + let mut reverse_depth = FastMap::default(); + let mut reverse_parents: FastMap> = FastMap::default(); + reverse_depth.insert(target, 0); + let (mut frontier, mut reverse_frontier) = (vec![source], vec![target]); + let (mut level, mut reverse_level) = (0, 0); + while !frontier.is_empty() && !reverse_frontier.is_empty() && level + reverse_level < max_hops { + if frontier.len() <= reverse_frontier.len() { + level += 1; + frontier = expand_level(forward, &frontier, level, &mut depth, &mut parents); + } else { + reverse_level += 1; + reverse_frontier = expand_level( + reverse, + &reverse_frontier, + reverse_level, + &mut reverse_depth, + &mut reverse_parents, + ); + } + let mut pending: Vec<_> = frontier + .iter() + .copied() + .filter(|n| reverse_depth.get(n) == Some(&reverse_level)) + .collect(); + if pending.is_empty() { + continue; + } + // Reverse only the DAG beyond this cut; stitching at several depths duplicates paths. + let mut visited = HashSet::new(); + while let Some(node) = pending.pop() { + if !visited.insert(node) { + continue; + } + if let Some(next) = reverse_parents.get(&node) { + for &n in next { + parents.entry(n).or_default().push(node); + pending.push(n); + } + } + } + depth.insert(target, level + reverse_level); + break; + } + (depth, parents) +} + /// Shortest paths from the search source to `target`, as node ids from source to target: /// all of them, or the first found. fn paths_to(target: i64, parents: &FastMap>, all: bool) -> Vec> { @@ -431,6 +531,7 @@ fn search_batch( source_index: usize, target_index: usize, adjacency: &Neighbors, + reverse_adjacency: Option<&Neighbors>, spec: &ShortestPathSpec, schema: &SchemaRef, ) -> datafusion::common::Result { @@ -450,7 +551,19 @@ fn search_batch( } let searches: FastMap = by_source .iter() - .map(|(&s, ts)| (s, search(adjacency, s, ts, spec.max_hops))) + .map(|(&s, ts)| { + let result = match reverse_adjacency.filter(|_| ts.len() == 1) { + Some(reverse) => search_bidirectional( + adjacency, + reverse, + s, + *ts.iter().next().expect("one target"), + spec.max_hops, + ), + None => search(adjacency, s, ts, spec.max_hops), + }; + (s, result) + }) .collect(); let mut rows = Vec::new(); let mut lengths = Vec::new(); @@ -498,6 +611,7 @@ pub(crate) fn plan_shortest_path( physical_inputs[0].clone(), edges, node.adjacency.clone(), + node.reverse_adjacency.clone(), node.spec.clone(), Arc::new(node.schema.as_arrow().clone()), )) @@ -514,7 +628,7 @@ fn plan_error(message: &str) -> GraphError { mod tests { use super::*; - fn neighbors(edges: &[(i64, i64)]) -> Neighbors { + fn adjacencies(edges: &[(i64, i64)]) -> (Arc, Arc) { let schema = Arc::new(Schema::new(vec![ Field::new("src", DataType::Int64, false), Field::new("dst", DataType::Int64, false), @@ -547,7 +661,12 @@ mod tests { let inc = rt .block_on(Adjacency::build(stream(), "dst", "src")) .unwrap(); - Neighbors::undirected(Arc::new(out), Arc::new(inc)) + (Arc::new(out), Arc::new(inc)) + } + + fn neighbors(edges: &[(i64, i64)]) -> Neighbors { + let (out, inc) = adjacencies(edges); + Neighbors::undirected(out, inc) } #[test] @@ -564,4 +683,105 @@ mod tests { let (depth, _) = search(&adj, 1, &HashSet::from([4]), 1); assert!(!depth.contains_key(&4)); } + + #[test] + fn bidirectional_keeps_every_meeting_node_and_parallel_edge() { + // A wide first layer makes the target's side expand. Two meeting nodes, + // including parallel edges on both sides of the cut, yield six paths. + let (out, inc) = adjacencies(&[ + (1, 2), + (1, 2), + (1, 3), + (1, 9), + (2, 4), + (2, 4), + (3, 4), + (3, 4), + (4, 5), + (5, 5), + (5, 1), + ]); + let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc)); + let (depth, parents) = search_bidirectional(&forward, &reverse, 1, 5, 3); + assert_eq!(depth[&5], 3); + let mut all = paths_to(5, &parents, true); + all.sort(); + assert_eq!( + all, + vec![vec![1, 2, 4, 5]; 4] + .into_iter() + .chain(vec![vec![1, 3, 4, 5]; 2]) + .collect::>() + ); + let one = paths_to(5, &parents, false); + assert_eq!(one.len(), 1); + assert!(all.contains(&one[0])); + assert!(!search_bidirectional(&forward, &reverse, 1, 5, 2) + .0 + .contains_key(&5)); + } + + #[test] + fn bidirectional_matches_forward_search_for_small_graphs() { + // All directed graphs on three nodes, including disconnected pairs and cycles. + let edges = [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)]; + for mask in 0..(1 << edges.len()) { + let selected: Vec<_> = edges + .iter() + .enumerate() + .filter_map(|(bit, &edge)| (mask & (1 << bit) != 0).then_some(edge)) + .collect(); + let (out, inc) = adjacencies(&selected); + let directed = Neighbors::directed(out.clone()); + let reverse = Neighbors::directed(inc.clone()); + let undirected = Neighbors::undirected(out, inc); + for (forward, reverse) in [(&directed, &reverse), (&undirected, &undirected)] { + for source in 0..3 { + for target in 0..3 { + for max_hops in 0..=3 { + let (expected_depth, expected_parents) = + search(forward, source, &HashSet::from([target]), max_hops); + let (depth, parents) = + search_bidirectional(forward, reverse, source, target, max_hops); + assert_eq!( + depth.get(&target), + expected_depth.get(&target), + "graph {mask}, {source}->{target}, max {max_hops}" + ); + if depth.contains_key(&target) { + let mut expected = paths_to(target, &expected_parents, true); + let mut actual = paths_to(target, &parents, true); + expected.sort(); + actual.sort(); + assert_eq!( + actual, expected, + "graph {mask}, {source}->{target}, max {max_hops}" + ); + } + } + } + } + } + } + } + + #[test] + fn bidirectional_avoids_expanding_the_wide_middle_layer() { + let mut edges = Vec::new(); + for branch in 1..=100 { + edges.push((0, branch)); + for leaf in 0..100 { + edges.push((branch, 101 + (branch - 1) * 100 + leaf)); + } + } + let (out, inc) = adjacencies(&edges); + let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc)); + let (depth, parents) = search_bidirectional(&forward, &reverse, 0, 101, 2); + assert_eq!(paths_to(101, &parents, true), vec![vec![0, 1, 101]]); + assert_eq!(depth.len(), 102); // Root, first layer and stitched target. + assert_eq!( + search(&forward, 0, &HashSet::from([101]), 2).0.len(), + 10_101 + ); + } } diff --git a/docs/csr.md b/docs/csr.md index 4f72c39..a9677e7 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -132,6 +132,12 @@ When those hold, it covers: Everything else in the query (filters, aggregations, sorts and any hop the CSR can't take) is planned by DataFusion as usual. You can mix: a three-hop query might run two hops on the CSR and one as a join. +For `shortestPath` and `allShortestPaths`, a source with one target in an input batch searches from both +ends using the outgoing and incoming CSR. It expands the smaller frontier and finishes the whole layer +before joining paths, so it keeps all shortest alternatives. Several targets for the same source share +one forward search. Undirected searches can use both ends without a persistent CSR too; directed searches +without a CSR keep their forward search over the adjacency built for that query. + ### Variable-length traversal An eligible variable-length path keeps one cursor per hop and yields bounded output batches. Its plan diff --git a/docs/cypher.md b/docs/cypher.md index 8ec7c1b..1ab7417 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -291,7 +291,11 @@ RETURN length(p), nodes(p) - **Undirected variable-length paths** skip immediate backtracking (`a-b-a`). - **`shortestPath` and `allShortestPaths`** need exactly one relationship type, no relationship variable, and named end nodes. Their length is unbounded unless you give a maximum. `nodes(p)` returns the node keys. - With no path, `MATCH` drops the row and `OPTIONAL MATCH` returns `null`. + With no path, `MATCH` drops the row and `OPTIONAL MATCH` returns `null`. `shortestPath` chooses one path + when several have the same minimum length; `allShortestPaths` returns every shortest path, including + duplicate node sequences that use different parallel edges. A single source-target pair searches from + both ends when reverse adjacency is available. Several targets for the same source in an input batch + share one forward search. - **Named paths** (`p = ...`) are only supported with `shortestPath` and `allShortestPaths`. ### Expressions and predicates diff --git a/python/python/tests/test_engine_paths.py b/python/python/tests/test_engine_paths.py index 1206e40..922a369 100644 --- a/python/python/tests/test_engine_paths.py +++ b/python/python/tests/test_engine_paths.py @@ -95,3 +95,45 @@ def test_search_without_csr_reads_every_edge_partition(): "MATCH p = shortestPath((a)-[:KNOWS*1..30]->(b)) RETURN length(p) AS l" ) assert got.to_pylist() == [{"l": 5}] + + +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize("direction", ["outgoing", "incoming", "undirected"]) +def test_all_shortest_paths_keeps_every_meeting_node(use_csr, direction): + people = pa.table({"id": list(range(1, 10))}) + knows = pa.table( + { + "src": [1, 1, 1, 1, 2, 2, 3, 3, 4, 5, 5], + "dst": [2, 2, 3, 9, 4, 4, 4, 4, 5, 5, 6], + } + ) + engine = CypherEngine(CONFIG, {"Person": people, "KNOWS": knows}, use_csr=use_csr) + pattern = { + "outgoing": "(a)-[:KNOWS*1..3]->(b)", + "incoming": "(b)<-[:KNOWS*1..3]-(a)", + "undirected": "(a)-[:KNOWS*1..3]-(b)", + }[direction] + got = engine.execute( + "MATCH (a:Person {id: 1}), (b:Person {id: 5}) " + f"MATCH p = allShortestPaths({pattern}) RETURN length(p) AS l, nodes(p) AS ns" + ).to_pylist() + paths = [[1, 2, 4, 5]] * 4 + [[1, 3, 4, 5]] * 2 + if direction == "incoming": + paths = [p[::-1] for p in paths] + assert sorted((r["l"], r["ns"]) for r in got) == sorted((3, p) for p in paths) + + +@pytest.mark.parametrize("use_csr", [False, True]) +def test_shortest_paths_shares_a_source_across_targets(use_csr): + engine = CypherEngine(CONFIG, {"Person": PEOPLE, "KNOWS": KNOWS}, use_csr=use_csr) + got = engine.execute( + "MATCH (a:Person {id: 5}), (b:Person) " + "MATCH p = allShortestPaths((a)-[:KNOWS*1..3]->(b)) " + "RETURN b.id AS target, length(p) AS l, nodes(p) AS ns" + ).to_pylist() + assert sorted((r["target"], r["l"], r["ns"]) for r in got) == [ + (1, 1, [5, 1]), + (2, 2, [5, 1, 2]), + (2, 2, [5, 1, 2]), + (3, 2, [5, 1, 3]), + ] From 673d106b41d22cb862c129a4c3e754cbfa82ad30 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:35:59 -0700 Subject: [PATCH 24/66] feat(cypher): add map and entity keys with node labels Implement keys for maps and exposed entity property columns, preserving null map fields while omitting absent entity properties. Resolve labels from canonical configured node bindings and preserve optional nulls and WITH aliases. Reject mixed-label bindings explicitly. Validation: 344 Rust library tests, 32 focused Python cases, all 615 Python tests, workspace all-target clippy, Python lint, executable documentation, strict docs build, spelling, formatting, and whitespace checks pass. Part of #20; mixed-label labels() remains an explicit engine boundary. Submitted by David Hughes (@stugorf), with AI Collaboration. --- .../src/datafusion_planner/expression.rs | 52 ++++--- .../datafusion_planner/scalar_functions.rs | 62 +++++++- crates/lance-graph/src/logical_plan.rs | 42 ++++++ crates/lance-graph/src/semantic.rs | 32 +++- docs/cypher.md | 27 +++- python/python/tests/test_entity_functions.py | 137 ++++++++++++++++++ 6 files changed, 329 insertions(+), 23 deletions(-) create mode 100644 python/python/tests/test_entity_functions.py diff --git a/crates/lance-graph/src/datafusion_planner/expression.rs b/crates/lance-graph/src/datafusion_planner/expression.rs index 0000f36..1777625 100644 --- a/crates/lance-graph/src/datafusion_planner/expression.rs +++ b/crates/lance-graph/src/datafusion_planner/expression.rs @@ -229,7 +229,7 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { }, "abs" | "round" | "floor" | "ceil" | "sqrt" | "left" | "right" | "substring" | "replace" | "split" | "trim" | "ltrim" | "rtrim" | "size" | "reverse" - | "last" | "tail" => { + | "last" | "tail" | "keys" | "__entity_keys" | "labels" => { super::scalar_functions::call(name, args.iter().map(to_df_value_expr).collect()) } "range" => { @@ -492,6 +492,27 @@ pub(crate) fn bind_value( schema: &datafusion::common::DFSchema, ) -> ValueExpression { let has = |name: &str| schema.has_column_with_unqualified_name(name); + let entity_properties = |var: &str| { + let var = var.to_lowercase(); + let prefix = format!("{var}__"); + if has(&var) { + return Vec::new(); + } + schema + .fields() + .iter() + .filter_map(|f| f.name().strip_prefix(&prefix)) + .map(|property| { + ( + property.to_string(), + ValueExpression::Property(crate::ast::PropertyRef { + variable: var.clone(), + property: property.to_string(), + }), + ) + }) + .collect::>() + }; let mut expr = expr.clone(); crate::ast::transform_value(&mut expr, &mut |e| match e { ValueExpression::Property(p) @@ -503,28 +524,23 @@ pub(crate) fn bind_value( field: p.property.clone(), }; } + ValueExpression::ScalarFunction { name, args } if name.eq_ignore_ascii_case("keys") => { + if let [ValueExpression::Variable(var)] = args.as_slice() { + let properties = entity_properties(var); + if !properties.is_empty() { + *e = ValueExpression::ScalarFunction { + name: "__entity_keys".to_string(), + args: vec![ValueExpression::Map(properties)], + }; + } + } + } 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(); + let properties = entity_properties(var); if !properties.is_empty() { *value = ValueExpression::Map(properties); } diff --git a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs index 94cd8d7..f28e8d5 100644 --- a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs +++ b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs @@ -50,6 +50,8 @@ impl ScalarUDFImpl for CypherScalar { let valid = match self.name.as_str() { "abs" | "round" | "floor" | "ceil" | "sqrt" => t.is_integer() || t.is_floating(), "range" => t.is_integer(), + "keys" | "__entity_keys" => matches!(t, DataType::Struct(_)), + "labels" => false, "size" | "reverse" => is_string(t) || list_field(t).is_some(), "last" | "tail" => list_field(t).is_some(), "left" | "right" | "substring" if i > 0 => t.is_integer(), @@ -64,7 +66,9 @@ impl ScalarUDFImpl for CypherScalar { "abs" if types[0].is_integer() => DataType::Int64, "abs" | "round" | "floor" | "ceil" | "sqrt" => DataType::Float64, "range" => DataType::List(Arc::new(Field::new("item", DataType::Int64, true))), - "split" => DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + "split" | "keys" | "__entity_keys" | "labels" => { + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))) + } "last" => list_field(&types[0]) .map(|f| f.data_type().clone()) .unwrap_or(DataType::Null), @@ -85,6 +89,31 @@ impl ScalarUDFImpl for CypherScalar { args.number_rows, ))); } + if matches!(self.name.as_str(), "keys" | "__entity_keys") { + let input = args.args[0].to_array(args.number_rows)?; + let map = input.as_struct(); + let nulls: Vec<_> = map + .columns() + .iter() + .map(|column| column.logical_nulls()) + .collect(); + let entity = self.name == "__entity_keys"; + let mut result = ListBuilder::new(StringBuilder::new()); + for row in 0..input.len() { + let missing = |i: usize| nulls[i].as_ref().is_some_and(|nulls| nulls.is_null(row)); + if input.is_null(row) || (entity && (0..map.num_columns()).all(missing)) { + result.append(false); + continue; + } + for (i, field) in map.fields().iter().enumerate() { + if !entity || !missing(i) { + result.values().append_value(field.name()); + } + } + result.append(true); + } + return Ok(ColumnarValue::Array(Arc::new(result.finish()))); + } if self.name == "range" { for arg in &mut args.args { *arg = @@ -338,6 +367,37 @@ mod tests { .await } + #[tokio::test] + async fn cypher_entity_map_keys() { + for (expression, expected) in [ + ("keys({x: null, y: 1})", vec!["x", "y"]), + ("keys({})", vec![]), + ] { + let out = run(expression).await.unwrap(); + let keys = out.column(0).as_list::().value(0); + assert_eq!( + keys.as_string::() + .iter() + .map(Option::unwrap) + .collect::>(), + expected + ); + } + for expression in ["keys(null)", "labels(null)"] { + assert_eq!( + run(expression) + .await + .unwrap() + .column(0) + .logical_null_count(), + 1 + ); + } + for expression in ["keys(1)", "labels(1)"] { + assert!(run(expression).await.is_err()); + } + } + #[tokio::test] async fn cypher_scalar_unicode_and_slicing() { for (expression, expected) in [ diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index 33dd3f0..0e5754a 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -685,6 +685,39 @@ impl<'a> LogicalPlanner<'a> { }); } } + ValueExpression::ScalarFunction { name, args } + if name.eq_ignore_ascii_case("labels") => + { + if let ValueExpression::Variable(var) = &args[0] { + if let Some(mapping) = self + .variables + .get(var) + .and_then(|label| self.config.get_node_mapping(label)) + { + *expr = ValueExpression::Case { + whens: vec![( + BooleanExpression::IsNull(ValueExpression::Property(PropertyRef { + variable: var.clone(), + property: mapping.id_field.clone(), + })), + ValueExpression::Literal(PropertyValue::Null), + )], + otherwise: Some(Box::new(ValueExpression::ScalarFunction { + name: "__list".to_string(), + args: vec![ValueExpression::Literal(PropertyValue::String( + mapping.label.clone(), + ))], + })), + }; + return; + } + } + args.iter_mut().for_each(|arg| self.resolve_value(arg)); + } + // keys(node/relationship) binds the entity's property columns, not its ID. + ValueExpression::ScalarFunction { name, args } + if name.eq_ignore_ascii_case("keys") + && matches!(args.as_slice(), [ValueExpression::Variable(_)]) => {} // Path functions take the path variable itself. ValueExpression::ScalarFunction { name, args } if matches!(name.to_lowercase().as_str(), "length" | "nodes") => {} @@ -961,6 +994,15 @@ impl<'a> LogicalPlanner<'a> { expr: &mut ValueExpression, ) -> Result { match expr { + ValueExpression::ScalarFunction { name, args } + if name.eq_ignore_ascii_case("labels") + && matches!(args.as_slice(), [ValueExpression::Variable(var)] if self.variables.get(var).is_some_and(|label| label.contains('|'))) => + { + Err(GraphError::UnsupportedFeature { + feature: "labels() on a mixed-label binding is not supported".to_string(), + location: Location::new(file!(), line!(), column!()), + }) + } ValueExpression::Case { whens, otherwise } => { for (condition, then) in whens.iter_mut() { input = self.rewrite_exists(input, condition)?; diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index 8de95c1..358544e 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -589,7 +589,7 @@ impl SemanticAnalyzer { match function_name.as_str() { "tolower" | "lower" | "toupper" | "upper" | "abs" | "round" | "floor" | "ceil" | "sqrt" | "trim" | "ltrim" | "rtrim" | "reverse" | "last" - | "tail" => { + | "tail" | "keys" | "labels" => { if args.len() != 1 { return Err(GraphError::PlanError { message: format!( @@ -670,7 +670,7 @@ impl SemanticAnalyzer { // Unknown scalar function - reject early with helpful error return Err(GraphError::UnsupportedFeature { feature: format!( - "Cypher function '{}' is not implemented. Supported scalar functions: toLower, lower, toUpper, upper, toString, coalesce, date_part, to_epoch_ms, length, nodes, size, range, head, last, tail, reverse, left, right, substring, replace, split, trim, lTrim, rTrim, abs, round, floor, ceil, sqrt, toInteger, toFloat, date, datetime, duration.inSeconds. Supported aggregate functions: COUNT, SUM, AVG, MIN, MAX, COLLECT.", + "Cypher function '{}' is not implemented. Supported scalar functions: toLower, lower, toUpper, upper, toString, coalesce, date_part, to_epoch_ms, length, nodes, size, range, head, last, tail, keys, labels, reverse, left, right, substring, replace, split, trim, lTrim, rTrim, abs, round, floor, ceil, sqrt, toInteger, toFloat, date, datetime, duration.inSeconds. Supported aggregate functions: COUNT, SUM, AVG, MIN, MAX, COLLECT.", name ), location: snafu::Location::new(file!(), line!(), column!()), @@ -678,6 +678,34 @@ impl SemanticAnalyzer { } } + if function_name == "labels" { + if let ValueExpression::Variable(var) = &args[0] { + if let Some(info) = self.variables.get_ci(var) { + if info + .labels + .iter() + .map(|label| label.to_lowercase()) + .collect::>() + .len() + > 1 + { + return Err(GraphError::UnsupportedFeature { + feature: "labels() on a mixed-label binding is not supported" + .to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + if info.variable_type == VariableType::Relationship { + return Err(GraphError::PlanError { + message: "labels() requires a node, not a relationship" + .to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + } + } + } + // Validate arguments recursively for arg in args { self.analyze_value_expression(arg)?; diff --git a/docs/cypher.md b/docs/cypher.md index 1ab7417..9072370 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -325,6 +325,7 @@ RETURN length(p), nodes(p) | Aggregates | `count` (including `count(*)` and `count(DISTINCT x)`), `sum`, `avg`, `min`, `max`, `collect` (including `DISTINCT`; skips nulls) | | Strings | `toLower` / `lower`, `toUpper` / `upper`, `toString`, `left`, `right`, `substring`, `replace`, `split`, `trim`, `lTrim`, `rTrim`, `reverse`, `size` | | Conversion | `toInteger`, `toFloat`, `coalesce` | +| Properties | `keys(map)`, `keys(node)`, `keys(relationship)`, `labels(node)` | | Lists | `size`, `head`, `last`, `tail`, `reverse`, `range(start, end [, step])` (inclusive) | | Math | `abs`, `round`, `floor`, `ceil`, `sqrt` | | Temporal | `date`, `datetime`, `date_part('month', d)`, `to_epoch_ms`, `duration.inSeconds(a, b).seconds` | @@ -361,9 +362,31 @@ result = engine.execute(""" print(result.to_pylist()) # [{'text': '猫😀', 'numbers': [5, 3, 1], 'rounded': -1.0}] ``` +`keys` returns the field names in a map, including keys with null values. For a node or relationship, +it returns its exposed, non-null property names, including the ID or endpoint columns already exposed +by the engine. Property names use the engine's normalized column spelling. Key order is not guaranteed. +`keys({})` returns an empty list. `labels(node)` returns the canonical label configured for that node, +even if the query spells it differently. Both functions preserve aliases carried through `WITH` and +return null for a missing entity from `OPTIONAL MATCH`. + +!!! note "Labels on mixed-label bindings" + + `labels(p)` supports a binding with one configured label. A pattern such as + `(p:Person:City)` is rejected with a mixed-label error, including after `WITH p AS x`. + The engine does not retain each row's label after this union. It never guesses a label from an ID + that might occur in both tables. + +```python +result = engine.execute(""" + MATCH (p:Person {name: 'Alice'}) WITH p AS person + RETURN labels(person) AS labels, size(keys(person)) AS properties +""") +print(result.to_pylist()) # [{'labels': ['Person'], 'properties': 3}] +``` + Any other function raises `ValueError` with `Cypher function 'name' is not implemented`. This is a -bounded function set, not a claim of full Cypher conformance. Remaining functions include `keys`, -`labels`, `reduce`, the list predicates `any` / `all` / `none` / `single`, and `relationships(p)`. +bounded function set, not a claim of full Cypher conformance. Remaining functions include `reduce`, +the list predicates `any` / `all` / `none` / `single`, and `relationships(p)`. When `size()` consumes an unmapped list comprehension, such as `size([x IN values WHERE x > 0])`, the planner counts retained elements directly. Each duplicate and each retained null element still contributes one; diff --git a/python/python/tests/test_entity_functions.py b/python/python/tests/test_entity_functions.py new file mode 100644 index 0000000..e891b26 --- /dev/null +++ b/python/python/tests/test_entity_functions.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Entity functions retain graph binding information and map null semantics.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture(params=[False, True], ids=["join", "csr"]) +def engine(request): + config = ( + GraphConfig.builder() + .node_table("Person", "person_id") + .node_table("City", "person_id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .rel_table("VISITS", "Person", "City", "src", "dst") + .build() + ) + return CypherEngine( + config, + { + "Person": pa.table( + { + "person_id": [1, 2, 3], + "name": ["Alice", None, "Carol"], + "age": [30, None, 0], + "missing": pa.nulls(3), + } + ), + "City": pa.table({"person_id": [1], "name": ["London"]}), + "KNOWS": pa.table({"src": [1, 2], "dst": [2, 3], "since": [2020, None]}), + "VISITS": pa.table({"src": [1], "dst": [1]}), + }, + use_csr=request.param, + ) + + +def test_node_keys_labels_and_with_alias(engine): + result = engine.execute( + "MATCH (p:person) WITH p AS person " + "RETURN person.person_id AS id, keys(person) AS keys, labels(person) AS labels " + "ORDER BY id" + ).to_pylist() + assert [sorted(row["keys"]) for row in result] == [ + ["age", "name", "person_id"], + ["person_id"], + ["age", "name", "person_id"], + ] + assert [row["labels"] for row in result] == [["Person"]] * 3 + + +def test_optional_node_and_relationship_keys(engine): + result = engine.execute( + "MATCH (p:Person) OPTIONAL MATCH (p)-[r:KNOWS]->(q:Person) " + "RETURN p.person_id AS id, keys(q) AS node_keys, labels(q) AS labels, " + "keys(r) AS edge_keys ORDER BY id" + ).to_pylist() + for row in result: + for key in ("node_keys", "edge_keys"): + if row[key] is not None: + row[key] = sorted(row[key]) + assert result == [ + { + "id": 1, + "node_keys": ["person_id"], + "labels": ["Person"], + "edge_keys": ["dst", "since", "src"], + }, + { + "id": 2, + "node_keys": ["age", "name", "person_id"], + "labels": ["Person"], + "edge_keys": ["dst", "src"], + }, + {"id": 3, "node_keys": None, "labels": None, "edge_keys": None}, + ] + + +def test_relationship_keys_with_alias(engine): + result = engine.execute( + "MATCH (:Person)-[r:KNOWS]->(:Person) WITH r AS edge " + "RETURN edge.src AS src, keys(edge) AS keys ORDER BY src" + ).to_pylist() + assert [sorted(row["keys"]) for row in result] == [ + ["dst", "since", "src"], + ["dst", "src"], + ] + + +@pytest.mark.parametrize("value", [{}, {"MixedKey": None, "name": "Alice"}, None]) +def test_map_keys_keep_null_values(engine, value): + result = engine.execute( + "WITH $value AS map RETURN keys(map) AS keys", {"value": value} + ) + keys = result.to_pylist()[0]["keys"] + if value is None: + assert keys is None + else: + assert sorted(keys) == sorted(value) + + +def test_labels_null_and_keys_literal(engine): + result = engine.execute( + "WITH null AS value RETURN labels(value) AS labels, keys({x: null}) AS keys" + ).to_pylist() + assert result == [{"labels": None, "keys": ["x"]}] + + +@pytest.mark.parametrize( + "query", + [ + "MATCH (p:Person:City) RETURN labels(p) AS labels", + "MATCH (p:Person:City) WITH p AS x RETURN labels(x) AS labels", + "MATCH (:Person)-[:KNOWS|VISITS]->(p) RETURN labels(p) AS labels", + ], +) +def test_labels_reject_mixed_label_bindings(engine, query): + with pytest.raises((ValueError, NotImplementedError), match="mixed-label binding"): + engine.execute(query) + + +@pytest.mark.parametrize( + "query", + [ + "RETURN keys(1) AS keys", + "RETURN labels(1) AS labels", + "RETURN labels({x: 1}) AS labels", + "MATCH (:Person)-[r:KNOWS]->(:Person) RETURN labels(r) AS labels", + "RETURN keys() AS keys", + "RETURN labels(null, null) AS labels", + ], +) +def test_entity_function_invalid_arguments(engine, query): + with pytest.raises((ValueError, RuntimeError)): + engine.execute(query) From 6e535f6548d582584e3763742e236e1e8e5f4562 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:11:34 -0700 Subject: [PATCH 25/66] feat(cypher): preserve path and relationship values Retain actual edge occurrences and typed endpoint records through named paths, collect, and UNWIND. Keep length/node-only shortest queries on CSR and reuse cooperative bidirectional search for rich path values. Part of #20. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/src/ast.rs | 3 + crates/lance-graph/src/csr_expand.rs | 39 + .../src/datafusion_planner/analysis.rs | 14 + .../datafusion_planner/builder/basic_ops.rs | 2 + .../datafusion_planner/builder/call_ops.rs | 1 + .../src/datafusion_planner/builder/helpers.rs | 19 + .../src/datafusion_planner/builder/mod.rs | 2 + .../datafusion_planner/builder/path_ops.rs | 214 ++++ .../src/datafusion_planner/cypher_udf.rs | 21 +- .../src/datafusion_planner/expression.rs | 28 +- .../lance-graph/src/datafusion_planner/mod.rs | 2 +- .../datafusion_planner/scalar_functions.rs | 48 +- crates/lance-graph/src/lib.rs | 3 + crates/lance-graph/src/logical_plan.rs | 113 +- crates/lance-graph/src/parser.rs | 14 + crates/lance-graph/src/path_inventory.rs | 282 +++++ crates/lance-graph/src/path_match.rs | 967 ++++++++++++++++++ crates/lance-graph/src/path_value.rs | 273 +++++ crates/lance-graph/src/pattern_order.rs | 2 + crates/lance-graph/src/semantic.rs | 12 +- crates/lance-graph/src/shortest_path.rs | 142 +-- docs/api-reference.md | 6 + docs/csr.md | 6 + docs/cypher.md | 53 +- python/python/tests/test_path_values.py | 350 +++++++ 25 files changed, 2537 insertions(+), 79 deletions(-) create mode 100644 crates/lance-graph/src/datafusion_planner/builder/path_ops.rs create mode 100644 crates/lance-graph/src/path_inventory.rs create mode 100644 crates/lance-graph/src/path_match.rs create mode 100644 crates/lance-graph/src/path_value.rs create mode 100644 python/python/tests/test_path_values.py diff --git a/crates/lance-graph/src/ast.rs b/crates/lance-graph/src/ast.rs index 5676f95..3fac617 100644 --- a/crates/lance-graph/src/ast.rs +++ b/crates/lance-graph/src/ast.rs @@ -284,6 +284,9 @@ pub struct NodePattern { /// A path pattern connecting nodes through relationships #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct PathPattern { + /// Named ordinary path (`p = (a)-[:R]->(b)`). + #[serde(default)] + pub variable: Option, /// Starting node pub start_node: NodePattern, /// Relationships and intermediate nodes diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index aeb1e71..66b009f 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -262,6 +262,23 @@ impl Adjacency { }) } + /// Inspect one occurrence so cancelled work also stops through long tombstone runs. + fn next_candidate(&self, key: i64, offset: &mut usize) -> Option> { + let base = self.base_neighbors(key); + let position = *offset; + let neighbor = if position < base.len() { + base[position] + } else { + *self.added.get(&key)?.get(position - base.len())? + }; + *offset += 1; + let removed = self + .removed + .get(&key) + .is_some_and(|r| r.binary_search(&position).is_ok()); + Some((!removed).then_some(neighbor)) + } + /// Live edges: base + appended - removed. pub fn num_edges(&self) -> u64 { self.base.csr.num_edges() + self.added_count - self.removed_count @@ -842,6 +859,21 @@ impl Neighbors { ) } + /// Inspect one raw edge, including tombstones and suppressed reverse self-loops. + pub(crate) fn next_candidate( + &self, + key: i64, + position: &mut NeighborPosition, + ) -> Option> { + if let Some(candidate) = self.forward.next_candidate(key, &mut position.forward) { + return Some(candidate); + } + self.reverse + .as_ref()? + .next_candidate(key, &mut position.reverse) + .map(|candidate| candidate.filter(|&n| n as i64 != key)) + } + /// Edges visited over all keys (self-loops count once per direction in this estimate). pub fn num_edges(&self) -> u64 { self.forward.num_edges() + self.reverse.as_ref().map_or(0, |r| r.num_edges()) @@ -1591,6 +1623,13 @@ impl ExtensionPlanner for CsrExtensionPlanner { session_state, ))); } + if let Some(node) = node + .as_any() + .downcast_ref::() + { + return crate::path_match::plan_path_match(node, physical_inputs, session_state) + .map(Some); + } if let Some(node) = node .as_any() .downcast_ref::() diff --git a/crates/lance-graph/src/datafusion_planner/analysis.rs b/crates/lance-graph/src/datafusion_planner/analysis.rs index 54a20df..477231b 100644 --- a/crates/lance-graph/src/datafusion_planner/analysis.rs +++ b/crates/lance-graph/src/datafusion_planner/analysis.rs @@ -37,6 +37,7 @@ pub struct RelationshipInstance { /// Planning context that tracks state during plan building pub struct PlanningContext<'a> { pub analysis: &'a QueryAnalysis, + pub(crate) path_sources: crate::path_inventory::SourceBatches, pub(crate) relationship_instance_idx: HashMap, /// Element rows of the list comprehensions being built, innermost last. pub(crate) call_input: Option, @@ -47,6 +48,7 @@ impl<'a> PlanningContext<'a> { pub fn new(analysis: &'a QueryAnalysis) -> Self { Self { analysis, + path_sources: Default::default(), relationship_instance_idx: HashMap::new(), comprehension_elements: Vec::new(), call_input: None, @@ -223,6 +225,18 @@ fn analyze_operator( analysis.required_datasets.insert(label.clone()); } } + LogicalOperator::PathMatch { input, path } => { + analyze_operator(input, analysis, rel_counter)?; + for node in + std::iter::once(&path.start_node).chain(path.segments.iter().map(|s| &s.end_node)) + { + if let (Some(variable), Some(label)) = (&node.variable, node.labels.first()) { + analysis + .var_to_label + .insert(variable.clone(), label.clone()); + } + } + } LogicalOperator::ShortestPath { input, .. } => { analyze_operator(input, analysis, rel_counter)?; } 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 a9179a7..fa0944b 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs @@ -57,6 +57,7 @@ impl DataFusionPlanner { let mut branch = PlanningContext::new(&analysis); branch.relationship_instance_idx = ctx.relationship_instance_idx.clone(); branch.call_input = ctx.call_input.clone(); + branch.path_sources = ctx.path_sources.clone(); branch.comprehension_elements = ctx.comprehension_elements.clone(); plans.push(self.build_operator(&mut branch, input)?); // Branches consume distinct relationship instances, in analysis order. @@ -124,6 +125,7 @@ impl DataFusionPlanner { let mut branch = PlanningContext::new(&analysis); branch.relationship_instance_idx = ctx.relationship_instance_idx.clone(); branch.call_input = ctx.call_input.clone(); + branch.path_sources = ctx.path_sources.clone(); branch.comprehension_elements = ctx.comprehension_elements.clone(); plans.push(self.build_operator(&mut branch, input)?); ctx.relationship_instance_idx = branch.relationship_instance_idx; diff --git a/crates/lance-graph/src/datafusion_planner/builder/call_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/call_ops.rs index 15d9a77..d93e67d 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/call_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/call_ops.rs @@ -68,6 +68,7 @@ impl DataFusionPlanner { let analysis = analysis::analyze(inner)?; let mut child_ctx = PlanningContext::new(&analysis); child_ctx.call_input = Some(seed); + child_ctx.path_sources = ctx.path_sources.clone(); let child = self.build_operator(&mut child_ctx, inner)?; let node = ApplyNode::try_new(input, child, imports) .map_err(|error| self.plan_error("CALL output", error))?; diff --git a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs index df3acd3..d2d3e38 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/helpers.rs @@ -96,6 +96,15 @@ impl DataFusionPlanner { LogicalOperator::OptionalMatch { optional, .. } => { Self::collect_variables(optional, vars); } + LogicalOperator::PathMatch { input, path } => { + Self::collect_variables(input, vars); + vars.extend( + std::iter::once(&path.start_node) + .chain(path.segments.iter().map(|s| &s.end_node)) + .filter_map(|n| n.variable.clone()), + ); + vars.extend(path.variable.clone()); + } LogicalOperator::ShortestPath { input, path_variable, @@ -178,6 +187,16 @@ impl DataFusionPlanner { LogicalOperator::OptionalMatch { optional, .. } => { Self::collect_labels(optional, labels) } + LogicalOperator::PathMatch { input, path } => { + Self::collect_labels(input, labels); + for node in std::iter::once(&path.start_node) + .chain(path.segments.iter().map(|s| &s.end_node)) + { + if let (Some(variable), Some(label)) = (&node.variable, node.labels.first()) { + labels.insert(variable.clone(), label.clone()); + } + } + } LogicalOperator::ShortestPath { input, .. } => Self::collect_labels(input, labels), LogicalOperator::Union { .. } => {} LogicalOperator::ListComprehension { input, .. } => Self::collect_labels(input, labels), diff --git a/crates/lance-graph/src/datafusion_planner/builder/mod.rs b/crates/lance-graph/src/datafusion_planner/builder/mod.rs index b738755..bf4f379 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/mod.rs @@ -19,6 +19,7 @@ mod comprehension_ops; mod expand_ops; mod helpers; mod join_builder; +mod path_ops; use super::DataFusionPlanner; use crate::error::Result; @@ -153,6 +154,7 @@ impl DataFusionPlanner { optional, keys, } => self.build_optional_match(ctx, input, optional, keys), + LogicalOperator::PathMatch { input, path } => self.build_path_match(ctx, input, path), LogicalOperator::ShortestPath { input, source_variable, diff --git a/crates/lance-graph/src/datafusion_planner/builder/path_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/path_ops.rs new file mode 100644 index 0000000..7cf18c2 --- /dev/null +++ b/crates/lance-graph/src/datafusion_planner/builder/path_ops.rs @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +use crate::ast::PathPattern; +use crate::case_insensitive::qualify_column; +use crate::datafusion_planner::{analysis::PlanningContext, DataFusionPlanner}; +use crate::error::{GraphError, Result}; +use crate::logical_plan::LogicalOperator; +use crate::path_inventory::{EdgeSource, NodeSource}; +use crate::path_match::{Binding, PathMatchNode, PathSpec, Step}; +use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder}; +use std::sync::Arc; + +impl DataFusionPlanner { + pub(crate) fn build_path_match( + &self, + ctx: &mut PlanningContext, + input: &LogicalOperator, + path: &PathPattern, + ) -> Result { + let error = |message: String| GraphError::UnsupportedFeature { + feature: message, + location: snafu::Location::new(file!(), line!(), column!()), + }; + let catalog = self.get_catalog()?; + let input = self.build_operator(ctx, input)?; + let mut node_labels = Vec::::new(); + let mut edge_labels = Vec::::new(); + for segment in &path.segments { + if segment.relationship.types.is_empty() { + return Err(error("named paths need explicit relationship types".into())); + } + for kind in &segment.relationship.types { + if edge_labels.iter().any(|t| t.eq_ignore_ascii_case(kind)) { + continue; + } + let mapping = self.get_relationship_mapping(kind)?; + for label in [&mapping.source_label, &mapping.target_label] { + let label = label.as_ref().ok_or_else(|| { + error("named paths need configured endpoint labels".into()) + })?; + if !node_labels.iter().any(|l| l.eq_ignore_ascii_case(label)) { + node_labels.push(label.clone()); + } + } + edge_labels.push(mapping.relationship_type.clone()); + } + } + let mut plans = vec![input]; + let mut nodes = Vec::new(); + for label in &node_labels { + let mapping = self + .config + .get_node_mapping(label) + .ok_or_else(|| error(format!("missing node mapping {label}")))?; + let source = catalog + .node_source(label) + .ok_or_else(|| error(format!("missing node table {label}")))?; + let schema = source.schema(); + let field = schema + .fields() + .iter() + .find(|f| f.name().eq_ignore_ascii_case(&mapping.id_field)) + .ok_or_else(|| error(format!("missing node ID {}", mapping.id_field)))?; + if !crate::csr_expand::is_supported_id_type(field.data_type()) { + return Err(error(format!( + "path node ID '{}' has unsupported type {}", + field.name(), + field.data_type() + ))); + } + nodes.push(NodeSource { + label: mapping.label.clone(), + id: mapping.id_field.clone(), + schema, + }); + plans.push( + LogicalPlanBuilder::scan(label.to_lowercase(), source, None) + .and_then(|b| b.build()) + .map_err(|e| self.plan_error("Failed to scan path nodes", e))?, + ); + } + let mut edges = Vec::new(); + for label in &edge_labels { + let mapping = self.get_relationship_mapping(label)?; + let source = catalog + .relationship_source(label) + .ok_or_else(|| error(format!("missing relationship table {label}")))?; + let schema = source.schema(); + for id in [&mapping.source_id_field, &mapping.target_id_field] { + let field = schema + .fields() + .iter() + .find(|f| f.name().eq_ignore_ascii_case(id)) + .ok_or_else(|| error(format!("missing relationship ID {id}")))?; + if !crate::csr_expand::is_supported_id_type(field.data_type()) { + return Err(error(format!( + "path relationship ID '{id}' has unsupported type {}", + field.data_type() + ))); + } + } + edges.push(EdgeSource { + label: mapping.relationship_type.clone(), + source_label: node_labels + .iter() + .position(|l| Some(l) == mapping.source_label.as_ref()) + .unwrap(), + target_label: node_labels + .iter() + .position(|l| Some(l) == mapping.target_label.as_ref()) + .unwrap(), + source_id: mapping.source_id_field.clone(), + target_id: mapping.target_id_field.clone(), + schema, + }); + plans.push( + LogicalPlanBuilder::scan(label.to_lowercase(), source, None) + .and_then(|b| b.build()) + .map_err(|e| self.plan_error("Failed to scan path edges", e))?, + ); + } + let mut bindings = Vec::new(); + for node in + std::iter::once(&path.start_node).chain(path.segments.iter().map(|s| &s.end_node)) + { + let variable = node + .variable + .clone() + .ok_or_else(|| error("named path endpoint was not assigned a variable".into()))?; + let label = ctx + .analysis + .var_to_label + .get(&variable) + .or_else(|| node.labels.first()) + .ok_or_else(|| { + error(format!("path variable '{variable}' needs a resolved label")) + })?; + let label = node_labels + .iter() + .position(|l| l.eq_ignore_ascii_case(label)) + .ok_or_else(|| { + error(format!( + "path label '{label}' is incompatible with its relationship types" + )) + })?; + bindings.push(Binding { + id_column: qualify_column(&variable, &nodes[label].id), + variable, + label, + properties: node.properties.clone(), + }); + } + let shortest = path.shortest.as_ref().map(|s| s.all); + let mut steps = Vec::new(); + for segment in &path.segments { + let rel = &segment.relationship; + let min = rel.length.as_ref().and_then(|l| l.min).unwrap_or(1); + let max = rel + .length + .as_ref() + .map(|l| { + l.max.unwrap_or(if shortest.is_some() { + u32::MAX + } else { + crate::MAX_VARIABLE_LENGTH_HOPS + }) + }) + .unwrap_or(1); + if min > max || (shortest.is_none() && max > crate::MAX_VARIABLE_LENGTH_HOPS) { + return Err(error(format!("invalid named path range {min}..{max}; ordinary variable-length paths allow at most {} hops", crate::MAX_VARIABLE_LENGTH_HOPS))); + } + steps.push(Step { + types: rel + .types + .iter() + .map(|t| { + edge_labels + .iter() + .position(|l| l.eq_ignore_ascii_case(t)) + .unwrap() + }) + .collect(), + direction: rel.direction.clone(), + min, + max, + relationship: rel.variable.clone(), + properties: rel.properties.clone(), + }); + } + let variable = path + .variable + .as_ref() + .or_else(|| path.shortest.as_ref().and_then(|s| s.variable.as_ref())) + .cloned() + .ok_or_else(|| error("rich path needs a variable".into()))?; + let node = PathMatchNode::try_new( + plans, + PathSpec { + variable, + sources: ctx.path_sources.clone(), + nodes, + edges, + bindings, + steps, + shortest, + }, + ) + .map_err(|e| self.plan_error("Failed to plan path value", e))?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(node), + })) + } +} diff --git a/crates/lance-graph/src/datafusion_planner/cypher_udf.rs b/crates/lance-graph/src/datafusion_planner/cypher_udf.rs index e4f47a2..171e787 100644 --- a/crates/lance-graph/src/datafusion_planner/cypher_udf.rs +++ b/crates/lance-graph/src/datafusion_planner/cypher_udf.rs @@ -124,7 +124,14 @@ impl ScalarUDFImpl for CypherField { 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() { + let input_type = args.arg_fields[0].data_type(); + let input_type = match input_type { + DataType::Struct(fields) if crate::path_value::entity_kind(input_type).is_some() => { + fields[0].data_type() + } + other => other, + }; + let data_type = match input_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"), @@ -141,13 +148,23 @@ impl ScalarUDFImpl for CypherField { return exec_err!("cypher_field takes a field name literal"); }; let base = args.args[0].to_array(args.number_rows)?; + let parent_nulls = base.logical_nulls(); + let base = if crate::path_value::entity_kind(base.data_type()).is_some() { + base.as_struct().column(0).clone() + } else { + base + }; 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(), + arrow::buffer::NullBuffer::union( + parent_nulls.as_ref(), + base.logical_nulls().as_ref(), + ) + .as_ref(), column.logical_nulls().as_ref(), ); let data = column.to_data().into_builder().nulls(nulls).build()?; diff --git a/crates/lance-graph/src/datafusion_planner/expression.rs b/crates/lance-graph/src/datafusion_planner/expression.rs index 1777625..78ab9d0 100644 --- a/crates/lance-graph/src/datafusion_planner/expression.rs +++ b/crates/lance-graph/src/datafusion_planner/expression.rs @@ -223,13 +223,14 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { / lit(1000_i64), ]), // A path's length and node ids are columns of its variable (see shortest_path). - "length" | "nodes" => match &args[0] { + "length" | "nodes" | "relationships" | "__path_node_records" => match &args[0] { VE::Variable(path) => col(qualify_column(path, &name.to_lowercase())), _ => Expr::Literal(datafusion::scalar::ScalarValue::Null, None), }, "abs" | "round" | "floor" | "ceil" | "sqrt" | "left" | "right" | "substring" | "replace" | "split" | "trim" | "ltrim" | "rtrim" | "size" | "reverse" - | "last" | "tail" | "keys" | "__entity_keys" | "labels" => { + | "last" | "tail" | "keys" | "__entity_keys" | "labels" | "startnode" + | "endnode" => { super::scalar_functions::call(name, args.iter().map(to_df_value_expr).collect()) } "range" => { @@ -524,6 +525,29 @@ pub(crate) fn bind_value( field: p.property.clone(), }; } + ValueExpression::ScalarFunction { name, args } + if matches!( + name.to_lowercase().as_str(), + "length" | "nodes" | "relationships" | "__path_node_records" + ) => + { + if let [base] = args.as_slice() { + let value = match base { + ValueExpression::Variable(var) => has(&var.to_lowercase()), + _ => true, + }; + if value { + *e = ValueExpression::Field { + base: Box::new(base.clone()), + field: if name == "__path_node_records" { + "node_records".into() + } else { + name.to_lowercase() + }, + }; + } + } + } ValueExpression::ScalarFunction { name, args } if name.eq_ignore_ascii_case("keys") => { if let [ValueExpression::Variable(var)] = args.as_slice() { let properties = entity_properties(var); diff --git a/crates/lance-graph/src/datafusion_planner/mod.rs b/crates/lance-graph/src/datafusion_planner/mod.rs index f1f5d11..8e67b58 100644 --- a/crates/lance-graph/src/datafusion_planner/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/mod.rs @@ -18,7 +18,7 @@ pub mod analysis; mod builder; mod config_helpers; mod cypher_udf; -mod expression; +pub(crate) mod expression; mod join_ops; mod scalar_functions; mod scan_ops; diff --git a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs index f28e8d5..6494095 100644 --- a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs +++ b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs @@ -51,7 +51,10 @@ impl ScalarUDFImpl for CypherScalar { "abs" | "round" | "floor" | "ceil" | "sqrt" => t.is_integer() || t.is_floating(), "range" => t.is_integer(), "keys" | "__entity_keys" => matches!(t, DataType::Struct(_)), - "labels" => false, + "labels" => crate::path_value::entity_kind(t) == Some("node"), + "startnode" | "endnode" => { + crate::path_value::entity_kind(t) == Some("relationship") + } "size" | "reverse" => is_string(t) || list_field(t).is_some(), "last" | "tail" => list_field(t).is_some(), "left" | "right" | "substring" if i > 0 => t.is_integer(), @@ -62,6 +65,10 @@ impl ScalarUDFImpl for CypherScalar { } } Ok(match self.name.as_str() { + "startnode" | "endnode" => match &types[0] { + DataType::Struct(fields) => fields[2].data_type().clone(), + _ => DataType::Null, + }, "size" => DataType::Int64, "abs" if types[0].is_integer() => DataType::Int64, "abs" | "round" | "floor" | "ceil" | "sqrt" => DataType::Float64, @@ -89,15 +96,50 @@ impl ScalarUDFImpl for CypherScalar { args.number_rows, ))); } + if matches!(self.name.as_str(), "startnode" | "endnode") { + let input = args.args[0].to_array(args.number_rows)?; + let column = input + .as_struct() + .column(if self.name == "startnode" { 2 } else { 3 }); + let nulls = arrow::buffer::NullBuffer::union( + input.logical_nulls().as_ref(), + column.logical_nulls().as_ref(), + ); + return Ok(ColumnarValue::Array(arrow::array::make_array( + column.to_data().into_builder().nulls(nulls).build()?, + ))); + } + if self.name == "labels" { + let input = args.args[0].to_array(args.number_rows)?; + let labels = input + .as_struct() + .column(1) + .as_struct() + .column(0) + .as_string::(); + let mut result = ListBuilder::new(StringBuilder::new()); + for row in 0..input.len() { + if input.is_valid(row) { + result.values().append_value(labels.value(row)); + } + result.append(input.is_valid(row)); + } + return Ok(ColumnarValue::Array(Arc::new(result.finish()))); + } if matches!(self.name.as_str(), "keys" | "__entity_keys") { let input = args.args[0].to_array(args.number_rows)?; - let map = input.as_struct(); + let typed_entity = crate::path_value::entity_kind(input.data_type()).is_some(); + let map = if typed_entity { + input.as_struct().column(0).as_struct() + } else { + input.as_struct() + }; let nulls: Vec<_> = map .columns() .iter() .map(|column| column.logical_nulls()) .collect(); - let entity = self.name == "__entity_keys"; + let entity = self.name == "__entity_keys" || typed_entity; let mut result = ListBuilder::new(StringBuilder::new()); for row in 0..input.len() { let missing = |i: usize| nulls[i].as_ref().is_some_and(|nulls| nulls.is_null(row)); diff --git a/crates/lance-graph/src/lib.rs b/crates/lance-graph/src/lib.rs index 4606eac..4dbdeae 100644 --- a/crates/lance-graph/src/lib.rs +++ b/crates/lance-graph/src/lib.rs @@ -49,6 +49,9 @@ pub mod lance_vector_search; pub mod logical_plan; pub mod parameter_substitution; pub mod parser; +mod path_inventory; +mod path_match; +mod path_value; mod pattern_alternatives; mod pattern_order; pub mod query; diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index 0e5754a..36f1cde 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -145,6 +145,12 @@ pub enum LogicalOperator { /// `shortestPath` / `allShortestPaths` between two variables bound in `input`; adds /// `__length` and `__nodes` (node ids, source first). + /// A named path retaining actual relationship occurrences and typed node records. + PathMatch { + input: Box, + path: PathPattern, + }, + ShortestPath { input: Box, source_variable: String, @@ -244,6 +250,7 @@ pub struct LogicalPlanner<'a> { calls: usize, /// The node label of each path variable's nodes, for `nodes(path)`. path_labels: HashMap, + rich_paths: bool, /// 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. @@ -262,6 +269,7 @@ impl<'a> LogicalPlanner<'a> { comprehensions: 0, calls: 0, path_labels: HashMap::new(), + rich_paths: false, seed_source: None, } } @@ -281,6 +289,12 @@ impl<'a> LogicalPlanner<'a> { } fn retain_variable_alternation(&self, patterns: &[GraphPattern]) -> bool { + if patterns + .iter() + .all(|p| matches!(p, GraphPattern::Path(path) if path.variable.is_some())) + { + return true; + } let Some(catalog) = self.adjacency else { return false; }; @@ -346,6 +360,9 @@ impl<'a> LogicalPlanner<'a> { /// 100). `None` keeps the written order: start already right, no estimates, or a /// variable-length segment. fn cheaper_orientation(&self, path: &PathPattern) -> Option { + if path.variable.is_some() || path.shortest.is_some() { + return None; + } let bound = |n: &NodePattern| { n.variable .as_ref() @@ -401,11 +418,13 @@ impl<'a> LogicalPlanner<'a> { start_node: nodes[0].clone(), segments, shortest: path.shortest.clone(), + variable: path.variable.clone(), } } /// Convert a Cypher AST to a logical plan pub fn plan(&mut self, query: &CypherQuery) -> Result { + self.rich_paths = crate::path_value::needs_values(query); self.plan_query(query, None, false) } @@ -507,6 +526,7 @@ impl<'a> LogicalPlanner<'a> { .with_fan_out(self.fan_out) .with_adjacency(self.adjacency); planner.calls = self.calls; + planner.rich_paths = self.rich_paths; if preserve_entities { planner.variables = imports.0.clone(); planner.node_lists = imports.1.clone(); @@ -599,6 +619,7 @@ impl<'a> LogicalPlanner<'a> { .with_fan_out(self.fan_out) .with_adjacency(self.adjacency); inner.calls = self.calls; + inner.rich_paths = self.rich_paths; for name in &imports { if let Some(label) = self.variables.get(name) { inner.variables.insert(name.clone(), label.clone()); @@ -671,7 +692,7 @@ impl<'a> LogicalPlanner<'a> { match expr { ValueExpression::Variable(var) => { let property = match self.variables.get(var).map(String::as_str) { - Some("Path") => Some("length".to_string()), + Some("Path") => None, Some(label) => self .config .get_node_mapping(label) @@ -720,7 +741,10 @@ impl<'a> LogicalPlanner<'a> { && matches!(args.as_slice(), [ValueExpression::Variable(_)]) => {} // Path functions take the path variable itself. ValueExpression::ScalarFunction { name, args } - if matches!(name.to_lowercase().as_str(), "length" | "nodes") => {} + if matches!( + name.to_lowercase().as_str(), + "length" | "nodes" | "relationships" | "__path_node_records" + ) => {} ValueExpression::ScalarFunction { args, .. } | ValueExpression::AggregateFunction { args, .. } => { args.iter_mut().for_each(|a| self.resolve_value(a)) @@ -878,9 +902,23 @@ impl<'a> LogicalPlanner<'a> { } } } + GraphPattern::Path(path) if path.variable.is_some() => { + plan = Some(self.plan_named_path(plan, path)?) + } GraphPattern::Path(path) if path.shortest.is_some() => { plan = Some(self.plan_shortest_path(plan, path)?) } + GraphPattern::Path(path) + if self.rich_paths + && path + .segments + .iter() + .any(|s| s.relationship.variable.is_some()) => + { + let mut path = path.clone(); + path.variable = Some(format!("_path_value_{}", self.variables.len())); + plan = Some(self.plan_named_path(plan, &path)?); + } GraphPattern::Path(path) => plan = Some(self.plan_path(plan, path)?), } } @@ -1136,6 +1174,13 @@ impl<'a> LogicalPlanner<'a> { } => { input = self.rewrite_exists_in_value(input, list)?; self.resolve_value(list); + if self.rich_paths { + if let ValueExpression::ScalarFunction { name, .. } = list.as_mut() { + if name.eq_ignore_ascii_case("nodes") { + *name = "__path_node_records".into(); + } + } + } let label = match list.as_ref() { ValueExpression::Variable(l) => self.node_lists.get(l).cloned(), ValueExpression::ScalarFunction { name, args } @@ -1242,6 +1287,12 @@ impl<'a> LogicalPlanner<'a> { self.path_labels.insert(path_variable.clone(), label); } let length = rel.length.as_ref(); + if self.rich_paths { + return Ok(LogicalOperator::PathMatch { + input: Box::new(plan.expect("both ends planned")), + path: path.clone(), + }); + } Ok(LogicalOperator::ShortestPath { input: Box::new(plan.expect("both end nodes are planned")), source_variable: ends[0].clone(), @@ -1540,6 +1591,54 @@ impl<'a> LogicalPlanner<'a> { Ok(plan) } + fn plan_named_path( + &mut self, + base: Option, + path: &PathPattern, + ) -> Result { + let mut path = path.clone(); + if path.start_node.variable.is_none() { + path.start_node.variable = Some(format!("_path_node_{}", self.variables.len())); + } + let start = path.start_node.variable.as_ref().unwrap(); + let base = if self.variables.contains_key(start) { + base.expect("bound variables have input rows") + } else { + let scan = self.plan_node_scan(&path.start_node)?; + match base { + Some(input) => LogicalOperator::Join { + left: Box::new(input), + right: Box::new(scan), + join_type: JoinType::Cross, + }, + None => scan, + } + }; + for segment in &mut path.segments { + if segment.end_node.variable.is_none() { + segment.end_node.variable = Some(format!("_path_node_{}", self.variables.len())); + } + let variable = segment.end_node.variable.as_ref().unwrap(); + self.validate_variable_label(variable, &segment.end_node.labels)?; + let label = self + .variables + .get(variable) + .cloned() + .or_else(|| segment.end_node.labels.first().cloned()) + .ok_or_else(|| GraphError::UnsupportedFeature { + feature: "named path boundary needs a resolved node label".into(), + location: Location::new(file!(), line!(), column!()), + })?; + self.variables.insert(variable.clone(), label); + } + self.variables + .insert(path.variable.clone().unwrap(), "Path".into()); + Ok(LogicalOperator::PathMatch { + input: Box::new(base), + path, + }) + } + /// Extract the main variable from a logical plan (for chaining) #[allow(clippy::only_used_in_recursion)] fn extract_variable_from_plan(&self, plan: &LogicalOperator) -> Result { @@ -1581,6 +1680,14 @@ impl<'a> LogicalPlanner<'a> { message: "a comprehension's element rows bind no pattern variable".to_string(), location: snafu::Location::new(file!(), line!(), column!()), }), + LogicalOperator::PathMatch { path, .. } => Ok(path + .segments + .last() + .unwrap() + .end_node + .variable + .clone() + .unwrap()), LogicalOperator::ShortestPath { target_variable, .. } => Ok(target_variable.clone()), @@ -1622,7 +1729,7 @@ impl<'a> LogicalPlanner<'a> { ValueExpression::Variable(var) => { match self.variables.get(var) { // if it is a node variable, expand to all properties - Some(label) if label != "Unwound" => { + Some(label) if label != "Unwound" && label != "Path" => { let mapping = self.config.get_node_mapping(label).ok_or_else(|| { GraphError::PlanError { message: format!("Node label '{}' doesn't exist", label), diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index f9561c0..bdeacff 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -240,6 +240,7 @@ fn unwind_clause(input: &str) -> IResult<&str, UnwindClause> { fn graph_pattern(input: &str) -> IResult<&str, GraphPattern> { alt(( map(shortest_path_pattern, GraphPattern::Path), + map(named_path_pattern, GraphPattern::Path), map(path_pattern, GraphPattern::Path), map(node_pattern, GraphPattern::Node), ))(input) @@ -268,6 +269,18 @@ fn shortest_path_pattern(input: &str) -> IResult<&str, PathPattern> { Ok((input, path)) } +// Named ordinary paths retain their textual node and relationship order. +fn named_path_pattern(input: &str) -> IResult<&str, PathPattern> { + let (input, variable) = delimited( + multispace0, + identifier, + tuple((multispace0, char('='), multispace0)), + )(input)?; + let (input, mut path) = path_pattern(input)?; + path.variable = Some(variable.to_string()); + Ok((input, path)) +} + // Parse a path pattern (only if there are segments) fn path_pattern(input: &str) -> IResult<&str, PathPattern> { let (input, start_node) = node_pattern(input)?; @@ -287,6 +300,7 @@ fn path_pattern(input: &str) -> IResult<&str, PathPattern> { start_node, segments, shortest: None, + variable: None, }, )) } diff --git a/crates/lance-graph/src/path_inventory.rs b/crates/lance-graph/src/path_inventory.rs new file mode 100644 index 0000000..1741c22 --- /dev/null +++ b/crates/lance-graph/src/path_inventory.rs @@ -0,0 +1,282 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! One immutable inventory per rich-path operator. An edge's identity is its actual +//! scanned occurrence, shared by every hop, never reconstructed from endpoint pairs. + +use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; + +use arrow::array::{ + new_null_array, Array, ArrayRef, AsArray, Int64Array, RecordBatch, StringArray, StructArray, + UInt32Array, +}; +use arrow::compute::{cast, concat_batches, take}; +use arrow_schema::{DataType, Field, Fields, Schema}; +use datafusion::common::{exec_err, plan_err, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::ExecutionPlan; +use futures::{lock::Mutex, StreamExt}; + +/// One source scan per query, shared even by distinct named-path operators. +pub(crate) type SourceBatches = Arc>>; + +use crate::csr_expand::{cooperative_yield, is_supported_id_type}; +use crate::path_value::{entity_type, struct_array}; + +#[derive(Debug, Clone)] +pub(crate) struct NodeSource { + pub label: String, + pub id: String, + pub schema: Arc, +} + +#[derive(Debug, Clone)] +pub(crate) struct EdgeSource { + pub label: String, + pub source_label: usize, + pub target_label: usize, + pub source_id: String, + pub target_id: String, + pub schema: Arc, +} + +pub(crate) fn property_fields(schemas: impl Iterator>) -> Result { + let mut fields = BTreeMap::::new(); + for schema in schemas { + for field in schema.fields() { + let name = field.name().to_lowercase(); + if fields.get(&name).is_some_and(|t| t != field.data_type()) { + return plan_err!( + "path entities have conflicting Arrow types for property '{name}'" + ); + } + fields.insert(name, field.data_type().clone()); + } + } + Ok(fields + .into_iter() + .map(|(name, t)| Field::new(name, t, true)) + .collect()) +} + +pub(crate) fn id_column(batch: &RecordBatch, name: &str) -> Result { + let schema = batch.schema(); + let Some((i, field)) = schema + .fields() + .iter() + .enumerate() + .find(|(_, f)| f.name().eq_ignore_ascii_case(name)) + else { + return exec_err!("path ID column '{name}' is missing"); + }; + if !is_supported_id_type(field.data_type()) { + return exec_err!( + "path ID column '{name}' must be a signed integer or uint8/16/32, got {}", + field.data_type() + ); + } + Ok(cast(batch.column(i), &DataType::Int64)? + .as_primitive::() + .clone()) +} + +fn properties(batch: &RecordBatch, fields: &Fields) -> Result { + let columns = fields + .iter() + .map(|f| { + batch + .schema() + .fields() + .iter() + .position(|old| old.name().eq_ignore_ascii_case(f.name())) + .map(|i| batch.column(i).clone()) + .unwrap_or_else(|| new_null_array(f.data_type(), batch.num_rows())) + }) + .collect(); + Ok(RecordBatch::try_new( + Arc::new(Schema::new(fields.clone())), + columns, + )?) +} + +#[derive(Debug, Clone, Copy)] +pub(crate) struct EdgeRef { + pub edge: u32, + pub node: u32, + pub outgoing: bool, +} + +#[derive(Debug)] +pub(crate) struct Inventory { + pub nodes: StructArray, + pub edges: StructArray, + pub node_properties: RecordBatch, + pub edge_properties: RecordBatch, + pub node_labels: Vec, + pub edge_types: Vec, + pub endpoints: Vec>, + pub by_key: HashMap<(usize, i64), u32>, + pub neighbors: Vec>, +} + +impl Inventory { + pub async fn build( + node_sources: &[NodeSource], + edge_sources: &[EdgeSource], + plans: &[Arc], + context: Arc, + sources: &SourceBatches, + ) -> Result { + let mut batches = Vec::with_capacity(plans.len()); + let source_keys = node_sources + .iter() + .map(|s| (false, s.label.to_lowercase())) + .chain(edge_sources.iter().map(|s| (true, s.label.to_lowercase()))); + for (plan, key) in plans.iter().zip(source_keys) { + let mut sources = sources.lock().await; + if !sources.contains_key(&key) { + let mut stream = plan.execute(0, context.clone())?; + let mut read = Vec::new(); + while let Some(batch) = stream.next().await { + read.push(batch?); + cooperative_yield().await; + } + sources.insert(key.clone(), concat_batches(&plan.schema(), &read)?); + } + batches.push(sources[&key].clone()); + } + let node_fields = property_fields(node_sources.iter().map(|s| s.schema.clone()))?; + let edge_fields = property_fields(edge_sources.iter().map(|s| s.schema.clone()))?; + let node_type = entity_type(node_fields.clone(), "node", None); + let edge_type = entity_type(edge_fields.clone(), "relationship", Some(&node_type)); + let mut by_key = HashMap::new(); + let mut node_labels = Vec::new(); + let mut labels = Vec::new(); + let mut ids = Vec::new(); + let mut node_batches = Vec::new(); + for (label, source) in node_sources.iter().enumerate() { + let batch = &batches[label]; + let keys = id_column(batch, &source.id)?; + for row in 0..batch.num_rows() { + let dense = u32::try_from(ids.len()).map_err(|_| { + datafusion::common::DataFusionError::Execution( + "path inventory exceeds u32 row references".into(), + ) + })?; + if keys.is_valid(row) { + by_key.insert((label, keys.value(row)), dense); + } + node_labels.push(label); + labels.push(source.label.as_str()); + ids.push(keys.is_valid(row).then(|| keys.value(row))); + if row % 4096 == 4095 { + cooperative_yield().await; + } + } + node_batches.push(properties(batch, &node_fields)?); + } + let node_properties = concat_batches(&Arc::new(Schema::new(node_fields)), &node_batches)?; + let identity = StructArray::from(vec![ + ( + Arc::new(Field::new("label", DataType::Utf8, false)), + Arc::new(StringArray::from(labels)) as ArrayRef, + ), + // Null keys never enter the lookup; their records cannot appear in a path. + ( + Arc::new(Field::new("id", DataType::Int64, false)), + Arc::new(Int64Array::from_iter_values( + ids.iter().map(|v| v.unwrap_or_default()), + )) as ArrayRef, + ), + ]); + let nodes = struct_array( + &node_type, + vec![ + Arc::new(StructArray::from(node_properties.clone())), + Arc::new(identity), + ], + ); + let mut neighbors = vec![Vec::new(); nodes.len()]; + let mut edge_batches = Vec::new(); + let mut edge_types = Vec::new(); + let mut occurrence_ids = Vec::new(); + let mut endpoints = Vec::new(); + let mut labels = Vec::new(); + for (kind, source) in edge_sources.iter().enumerate() { + let batch = &batches[node_sources.len() + kind]; + let sources = id_column(batch, &source.source_id)?; + let targets = id_column(batch, &source.target_id)?; + for row in 0..batch.num_rows() { + let edge = u32::try_from(endpoints.len()).map_err(|_| { + datafusion::common::DataFusionError::Execution( + "path inventory exceeds u32 row references".into(), + ) + })?; + let ends = if sources.is_valid(row) && targets.is_valid(row) { + by_key + .get(&(source.source_label, sources.value(row))) + .zip(by_key.get(&(source.target_label, targets.value(row)))) + .map(|(&a, &b)| (a, b)) + } else { + None + }; + if let Some((a, b)) = ends { + neighbors[a as usize].push(EdgeRef { + edge, + node: b, + outgoing: true, + }); + if a != b { + neighbors[b as usize].push(EdgeRef { + edge, + node: a, + outgoing: false, + }); + } + } + endpoints.push(ends); + edge_types.push(kind); + occurrence_ids.push(row as i64); + labels.push(source.label.as_str()); + if row % 4096 == 4095 { + cooperative_yield().await; + } + } + edge_batches.push(properties(batch, &edge_fields)?); + } + let edge_properties = concat_batches(&Arc::new(Schema::new(edge_fields)), &edge_batches)?; + let identity = StructArray::from(vec![ + ( + Arc::new(Field::new("label", DataType::Utf8, false)), + Arc::new(StringArray::from(labels)) as ArrayRef, + ), + ( + Arc::new(Field::new("id", DataType::Int64, false)), + Arc::new(Int64Array::from(occurrence_ids)) as ArrayRef, + ), + ]); + let starts = UInt32Array::from_iter(endpoints.iter().map(|p| p.map(|(a, _)| a))); + let ends = UInt32Array::from_iter(endpoints.iter().map(|p| p.map(|(_, b)| b))); + let edges = struct_array( + &edge_type, + vec![ + Arc::new(StructArray::from(edge_properties.clone())), + Arc::new(identity), + take(&nodes, &starts, None)?, + take(&nodes, &ends, None)?, + ], + ); + Ok(Self { + nodes, + edges, + node_properties, + edge_properties, + node_labels, + edge_types, + endpoints, + by_key, + neighbors, + }) + } +} diff --git a/crates/lance-graph/src/path_match.rs b/crates/lance-graph/src/path_match.rs new file mode 100644 index 0000000..d63e606 --- /dev/null +++ b/crates/lance-graph/src/path_match.rs @@ -0,0 +1,967 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Rich paths use an immutable edge inventory, resumable trail enumeration, and the +//! existing bidirectional BFS on a bipartite node/edge graph for shortest paths. + +use std::cmp::Ordering; +use std::collections::{HashMap, HashSet}; +use std::fmt; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, Int64Array, RecordBatch, UInt32Array}; +use arrow::compute::take; +use arrow_schema::{Field, Schema, SchemaRef}; +use datafusion::common::{DFSchema, DFSchemaRef, Result}; +use datafusion::execution::session_state::SessionState; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::{col, Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, +}; +use futures::{lock::Mutex, StreamExt}; + +use crate::ast::{PropertyValue, RelationshipDirection}; +use crate::csr_expand::{cooperative_yield, Adjacency, Neighbors}; +use crate::path_inventory::{id_column, property_fields, EdgeSource, Inventory, NodeSource}; +use crate::path_value::{entity_type, path_type, paths_array, PathRows}; + +type Parents = HashMap, ahash::RandomState>; + +#[derive(Debug, Clone)] +pub(crate) struct Binding { + pub variable: String, + pub label: usize, + pub id_column: String, + pub properties: HashMap, +} + +#[derive(Debug, Clone)] +pub(crate) struct Step { + pub types: Vec, + pub direction: RelationshipDirection, + pub min: u32, + pub max: u32, + pub relationship: Option, + pub properties: HashMap, +} + +#[derive(Debug)] +pub(crate) struct PathSpec { + pub variable: String, + pub sources: crate::path_inventory::SourceBatches, + pub nodes: Vec, + pub edges: Vec, + pub bindings: Vec, + pub steps: Vec, + /// Some(all): exactly one shortest-path segment, with both endpoints bound. + pub shortest: Option, +} + +#[derive(Debug, Clone)] +pub(crate) struct PathMatchNode { + inputs: Vec, + spec: Arc, + schema: DFSchemaRef, + outputs: Vec, +} + +#[derive(Debug, Clone)] +enum Output { + RelationshipValue { step: usize }, + Node { boundary: usize, property: usize }, + Relationship { step: usize, property: usize }, +} + +impl PathMatchNode { + pub fn try_new(inputs: Vec, spec: PathSpec) -> Result { + let node_fields = property_fields(spec.nodes.iter().map(|s| s.schema.clone()))?; + let edge_fields = property_fields(spec.edges.iter().map(|s| s.schema.clone()))?; + let node_type = entity_type(node_fields.clone(), "node", None); + let edge_type = entity_type(edge_fields.clone(), "relationship", Some(&node_type)); + let mut fields = Vec::new(); + let mut outputs = Vec::new(); + let input = inputs[0].schema(); + for (boundary, binding) in spec.bindings.iter().enumerate().skip(1) { + if input.has_column_with_unqualified_name(&binding.id_column) { + continue; + } + for field in spec.nodes[binding.label].schema.fields() { + let name = crate::case_insensitive::qualify_column(&binding.variable, field.name()); + if fields.iter().any(|f: &Field| f.name() == &name) { + continue; + } + let property = node_fields + .iter() + .position(|f| f.name().eq_ignore_ascii_case(field.name())) + .expect("unified property"); + fields.push(Field::new(name, field.data_type().clone(), true)); + outputs.push(Output::Node { boundary, property }); + } + } + for (step, s) in spec.steps.iter().enumerate() { + if let Some(variable) = &s.relationship { + if s.min != 1 || s.max != 1 { + return datafusion::common::plan_err!("a relationship variable in a named variable-length path is not supported; use relationships(path)"); + } + fields.push(Field::new(variable.to_lowercase(), edge_type.clone(), true)); + outputs.push(Output::RelationshipValue { step }); + for field in &edge_fields { + fields.push(Field::new( + crate::case_insensitive::qualify_column(variable, field.name()), + field.data_type().clone(), + true, + )); + outputs.push(Output::Relationship { + step, + property: edge_fields + .iter() + .position(|f| f.name() == field.name()) + .unwrap(), + }); + } + } + } + fields.push(Field::new( + spec.variable.to_lowercase(), + path_type(node_type, edge_type), + true, + )); + let schema = Arc::new(input.join(&DFSchema::try_from(Schema::new(fields))?)?); + Ok(Self { + inputs, + spec: Arc::new(spec), + schema, + outputs, + }) + } +} +impl PartialEq for PathMatchNode { + fn eq(&self, other: &Self) -> bool { + self.inputs == other.inputs && Arc::ptr_eq(&self.spec, &other.spec) + } +} +impl Eq for PathMatchNode {} +impl Hash for PathMatchNode { + fn hash(&self, state: &mut H) { + self.inputs.hash(state); + Arc::as_ptr(&self.spec).hash(state); + } +} +impl PartialOrd for PathMatchNode { + fn partial_cmp(&self, other: &Self) -> Option { + Arc::as_ptr(&self.spec).partial_cmp(&Arc::as_ptr(&other.spec)) + } +} +impl UserDefinedLogicalNodeCore for PathMatchNode { + fn name(&self) -> &str { + "PathMatch" + } + fn inputs(&self) -> Vec<&LogicalPlan> { + self.inputs.iter().collect() + } + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + fn expressions(&self) -> Vec { + vec![] + } + fn prevent_predicate_push_down_columns(&self) -> HashSet { + self.schema + .fields() + .iter() + .skip(self.inputs[0].schema().fields().len()) + .map(|f| f.name().clone()) + .collect() + } + fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "PathMatch: {}, {}", + self.spec.variable, + if self.spec.shortest.is_some() { + "shortest edge inventory" + } else { + "trail edge inventory" + } + ) + } + fn with_exprs_and_inputs(&self, _: Vec, inputs: Vec) -> Result { + Ok(Self { + inputs, + ..self.clone() + }) + } +} + +#[derive(Debug)] +struct Ready { + inventory: Inventory, + nodes_allowed: Vec, + edges_allowed: Vec, + shortest: Option<(Neighbors, Neighbors)>, +} + +#[derive(Debug)] +struct PathMatchExec { + inputs: Vec>, + spec: Arc, + outputs: Vec, + schema: SchemaRef, + properties: Arc, + predicates: Vec>>, + ready: Arc>>>, +} + +fn predicate(properties: &HashMap) -> Option { + properties + .iter() + .map(|(key, value)| { + col(key.to_lowercase()).eq(crate::datafusion_planner::expression::to_df_value_expr( + &crate::ast::ValueExpression::Literal(value.clone()), + )) + }) + .reduce(Expr::and) +} + +pub(crate) fn plan_path_match( + node: &PathMatchNode, + inputs: &[Arc], + state: &SessionState, +) -> Result> { + let mut plans = vec![inputs[0].clone()]; + plans.extend( + inputs[1..] + .iter() + .map(|p| Arc::new(CoalescePartitionsExec::new(p.clone())) as Arc), + ); + let nodes = DFSchema::try_from(Schema::new(property_fields( + node.spec.nodes.iter().map(|s| s.schema.clone()), + )?))?; + let edges = DFSchema::try_from(Schema::new(property_fields( + node.spec.edges.iter().map(|s| s.schema.clone()), + )?))?; + let predicates = node + .spec + .bindings + .iter() + .map(|b| (&b.properties, &nodes)) + .chain(node.spec.steps.iter().map(|s| (&s.properties, &edges))) + .map(|(properties, schema)| { + predicate(properties) + .map(|e| state.create_physical_expr(e, schema)) + .transpose() + }) + .collect::>>()?; + let schema = Arc::new(node.schema.as_arrow().clone()); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + inputs[0].output_partitioning().clone(), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Ok(Arc::new(PathMatchExec { + inputs: plans, + spec: node.spec.clone(), + outputs: node.outputs.clone(), + schema, + properties, + predicates, + ready: Arc::new(Mutex::new(None)), + })) +} + +impl DisplayAs for PathMatchExec { + fn fmt_as(&self, _: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "PathMatchExec: {} edge inventory", self.spec.variable) + } +} +impl ExecutionPlan for PathMatchExec { + fn name(&self) -> &str { + "PathMatchExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + self.inputs.iter().collect() + } + fn with_new_children( + self: Arc, + inputs: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + inputs, + spec: self.spec.clone(), + outputs: self.outputs.clone(), + schema: self.schema.clone(), + properties: self.properties.clone(), + predicates: self.predicates.clone(), + ready: self.ready.clone(), + })) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let input = self.inputs[0].execute(partition, context.clone())?; + let (cell, plans, spec, outputs, schema, predicates) = ( + self.ready.clone(), + self.inputs[1..].to_vec(), + self.spec.clone(), + self.outputs.clone(), + self.schema.clone(), + self.predicates.clone(), + ); + let out_schema = schema.clone(); + let batch_size = context.session_config().batch_size(); + let stream = futures::stream::try_unfold( + (input, None::), + move |(mut input, mut pending)| { + let (cell, plans, spec, outputs, schema, predicates, context) = ( + cell.clone(), + plans.clone(), + spec.clone(), + outputs.clone(), + schema.clone(), + predicates.clone(), + context.clone(), + ); + async move { + loop { + if pending.is_none() { + let Some(batch) = input.next().await else { + return Ok(None); + }; + let batch = batch?; + let graph = { + let mut ready = cell.lock().await; + if ready.is_none() { + *ready = Some(Arc::new( + Ready::build(&spec, &plans, &predicates, context.clone()) + .await?, + )); + } + ready.as_ref().expect("built").clone() + }; + pending = Some(Pending::new(batch, graph, &spec)?); + } + let p = pending.as_mut().unwrap(); + let (rows, paths) = p.next(&spec, batch_size).await?; + if rows.is_empty() { + pending = None; + continue; + } + let indices = UInt32Array::from(rows); + let mut columns: Vec = p + .batch + .columns() + .iter() + .map(|c| take(c, &indices, None)) + .collect::>()?; + for output in &outputs { + let (array, positions) = match output { + Output::RelationshipValue { step } => { + let positions = UInt32Array::from_iter_values( + paths.iter().map(|p| p.edges[p.boundaries[*step]]), + ); + columns.push(take(&p.graph.inventory.edges, &positions, None)?); + continue; + } + Output::Node { boundary, property } => ( + p.graph.inventory.node_properties.column(*property), + UInt32Array::from_iter_values( + paths.iter().map(|p| p.nodes[p.boundaries[*boundary]]), + ), + ), + Output::Relationship { step, property } => ( + p.graph.inventory.edge_properties.column(*property), + UInt32Array::from_iter_values( + paths.iter().map(|p| p.edges[p.boundaries[*step]]), + ), + ), + }; + columns.push(take(array, &positions, None)?); + } + columns.push(paths_array( + &paths, + &p.graph.inventory.nodes, + &p.graph.inventory.edges, + )?); + let batch = RecordBatch::try_new(schema, columns)?; + return Ok(Some((batch, (input, pending)))); + } + } + }, + ); + Ok(Box::pin(RecordBatchStreamAdapter::new(out_schema, stream))) + } +} + +impl Ready { + async fn build( + spec: &PathSpec, + plans: &[Arc], + predicates: &[Option>], + context: Arc, + ) -> Result { + let inventory = + Inventory::build(&spec.nodes, &spec.edges, plans, context, &spec.sources).await?; + let mask = + |p: &Option>, batch: &RecordBatch| -> Result { + match p { + Some(p) => Ok(p + .evaluate(batch)? + .into_array(batch.num_rows())? + .as_boolean() + .clone()), + None => Ok(BooleanArray::from(vec![true; batch.num_rows()])), + } + }; + let nodes_allowed = predicates[..spec.bindings.len()] + .iter() + .map(|p| mask(p, &inventory.node_properties)) + .collect::>>()?; + let edges_allowed = predicates[spec.bindings.len()..] + .iter() + .map(|p| mask(p, &inventory.edge_properties)) + .collect::>>()?; + let shortest = if spec.shortest.is_some() { + let step = &spec.steps[0]; + let mut from = Vec::new(); + let mut to = Vec::new(); + for (edge, ends) in inventory.endpoints.iter().enumerate() { + if let Some((a, b)) = ends { + let vertex = inventory.nodes.len() as i64 + edge as i64; + let (a, b) = (*a as i64, *b as i64); + match step.direction { + RelationshipDirection::Outgoing => { + from.extend([a, vertex]); + to.extend([vertex, b]); + } + RelationshipDirection::Incoming => { + from.extend([b, vertex]); + to.extend([vertex, a]); + } + RelationshipDirection::Undirected => { + from.extend([a, vertex, b, vertex]); + to.extend([vertex, b, vertex, a]); + } + } + } + if edge % 1024 == 1023 { + cooperative_yield().await; + } + } + let batch = RecordBatch::try_from_iter([ + ("src", Arc::new(Int64Array::from(from)) as ArrayRef), + ("dst", Arc::new(Int64Array::from(to)) as ArrayRef), + ])?; + let stream = || { + let chunks: Vec<_> = (0..batch.num_rows()) + .step_by(1024) + .map(|start| batch.slice(start, (batch.num_rows() - start).min(1024))) + .collect(); + Box::pin(RecordBatchStreamAdapter::new( + batch.schema(), + futures::stream::iter(chunks).then(|batch| async move { + cooperative_yield().await; + Ok(batch) + }), + )) as SendableRecordBatchStream + }; + Some(( + Neighbors::directed(Arc::new( + Adjacency::build(stream(), "src", "dst") + .await + .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))?, + )), + Neighbors::directed(Arc::new( + Adjacency::build(stream(), "dst", "src") + .await + .map_err(|e| datafusion::common::DataFusionError::External(Box::new(e)))?, + )), + )) + } else { + None + }; + Ok(Self { + inventory, + nodes_allowed, + edges_allowed, + shortest, + }) + } +} + +struct Pending { + batch: RecordBatch, + graph: Arc, + keys: Vec>, + row: usize, + walk: Option, + shortest: Option, +} +impl Pending { + fn new(batch: RecordBatch, graph: Arc, spec: &PathSpec) -> Result { + let keys = spec + .bindings + .iter() + .map(|b| { + batch + .schema() + .index_of(&b.id_column) + .ok() + .map(|_| id_column(&batch, &b.id_column)) + .transpose() + }) + .collect::>>()?; + Ok(Self { + batch, + graph, + keys, + row: 0, + walk: None, + shortest: None, + }) + } + fn bound(&self, spec: &PathSpec) -> Option>> { + self.keys + .iter() + .zip(&spec.bindings) + .map(|(keys, binding)| match keys { + None => Some(None), + Some(keys) if keys.is_null(self.row) => None, + Some(keys) => self + .graph + .inventory + .by_key + .get(&(binding.label, keys.value(self.row))) + .copied() + .map(Some), + }) + .collect() + } + async fn next(&mut self, spec: &PathSpec, size: usize) -> Result<(Vec, Vec)> { + let mut rows = Vec::new(); + let mut paths = Vec::new(); + while self.row < self.batch.num_rows() && rows.len() < size { + if self.walk.is_none() && self.shortest.is_none() { + let Some(bound) = self.bound(spec) else { + self.row += 1; + continue; + }; + let source = bound[0].expect("source bound by planner"); + if !allowed(&self.graph.nodes_allowed[0], source) { + self.row += 1; + continue; + } + if let Some(all) = spec.shortest { + let target = bound[1].expect("shortest target bound"); + let (forward, reverse) = + self.graph.shortest.as_ref().expect("shortest adjacency"); + let step = &spec.steps[0]; + let max = step.max.saturating_mul(2); + let (depths, parents) = crate::shortest_path::search_bidirectional( + forward, + reverse, + source as i64, + target as i64, + max, + ) + .await; + if depths + .get(&(target as i64)) + .is_none_or(|&depth| depth / 2 < step.min) + { + self.row += 1; + continue; + } + self.shortest = Some(ShortestPaths::new( + target as i64, + parents, + all, + self.graph.inventory.nodes.len(), + )); + } else { + self.walk = Some(Walk::new(source, bound)); + } + } + let path = if let Some(search) = &mut self.shortest { + search.next().await + } else { + self.walk.as_mut().unwrap().next(spec, &self.graph).await + }; + if let Some(path) = path { + rows.push(self.row as u32); + paths.push(path); + } else { + self.walk = None; + self.shortest = None; + self.row += 1; + } + cooperative_yield().await; + } + Ok((rows, paths)) + } +} + +fn allowed(mask: &BooleanArray, row: u32) -> bool { + mask.is_valid(row as usize) && mask.value(row as usize) +} + +struct Frame { + stage: usize, + hops: u32, + node: u32, + next: usize, + advanced: bool, + entered: bool, +} +struct Walk { + frames: Vec, + path: PathRows, + bound: Vec>, +} +impl Walk { + fn new(source: u32, bound: Vec>) -> Self { + Self { + frames: vec![Frame { + stage: 0, + hops: 0, + node: source, + next: 0, + advanced: false, + entered: false, + }], + path: PathRows { + nodes: vec![source], + edges: vec![], + boundaries: vec![0], + }, + bound, + } + } + fn pop(&mut self) { + let frame = self.frames.pop().unwrap(); + if frame.entered { + self.path.nodes.pop(); + self.path.edges.pop(); + } else if frame.stage > 0 { + self.path.boundaries.pop(); + } + } + async fn next(&mut self, spec: &PathSpec, graph: &Ready) -> Option { + let mut work = 0_usize; + while let Some(frame) = self.frames.last_mut() { + work += 1; + if work.is_multiple_of(1024) { + cooperative_yield().await; + } + if frame.stage == spec.steps.len() { + if !frame.advanced { + frame.advanced = true; + return Some(self.path.clone()); + } + self.pop(); + continue; + } + let step = &spec.steps[frame.stage]; + let binding = &spec.bindings[frame.stage + 1]; + if !frame.advanced { + frame.advanced = true; + let same = spec.bindings[..frame.stage + 1] + .iter() + .position(|b| b.variable == binding.variable) + .map(|i| self.path.nodes[self.path.boundaries[i]]); + if frame.hops >= step.min + && graph.inventory.node_labels[frame.node as usize] == binding.label + && self.bound[frame.stage + 1].is_none_or(|n| n == frame.node) + && same.is_none_or(|n| n == frame.node) + && allowed(&graph.nodes_allowed[frame.stage + 1], frame.node) + { + let next = Frame { + stage: frame.stage + 1, + hops: 0, + node: frame.node, + next: 0, + advanced: false, + entered: false, + }; + self.path.boundaries.push(self.path.nodes.len() - 1); + self.frames.push(next); + continue; + } + } + if frame.hops == step.max { + self.pop(); + continue; + } + let Some(edge) = graph.inventory.neighbors[frame.node as usize].get(frame.next) else { + self.pop(); + continue; + }; + frame.next += 1; + let direction = match step.direction { + RelationshipDirection::Outgoing => edge.outgoing, + RelationshipDirection::Incoming => !edge.outgoing || edge.node == frame.node, + RelationshipDirection::Undirected => true, + }; + if !direction + || !step + .types + .contains(&graph.inventory.edge_types[edge.edge as usize]) + || self.path.edges.contains(&edge.edge) + || !allowed(&graph.edges_allowed[frame.stage], edge.edge) + { + continue; + } + let next = Frame { + stage: frame.stage, + hops: frame.hops + 1, + node: edge.node, + next: 0, + advanced: false, + entered: true, + }; + self.path.nodes.push(edge.node); + self.path.edges.push(edge.edge); + self.frames.push(next); + } + None + } +} + +/// Depth-first predecessor enumeration keeps only one in-progress shortest path. +struct ShortestPaths { + parents: Parents, + stack: Vec<(i64, usize)>, + all: bool, + done: bool, + nodes: usize, +} +impl ShortestPaths { + fn new(target: i64, parents: Parents, all: bool, nodes: usize) -> Self { + Self { + parents, + stack: vec![(target, 0)], + all, + done: false, + nodes, + } + } + async fn next(&mut self) -> Option { + if self.done { + return None; + } + let mut work = 0_usize; + while let Some((node, next)) = self.stack.last_mut() { + work += 1; + if work.is_multiple_of(1024) { + cooperative_yield().await; + } + match self.parents.get(node) { + None if *next == 0 => { + *next = 1; + let nodes: Vec<_> = self + .stack + .iter() + .rev() + .step_by(2) + .map(|(n, _)| *n as u32) + .collect(); + let edges = self + .stack + .iter() + .rev() + .skip(1) + .step_by(2) + .map(|(n, _)| (*n as usize - self.nodes) as u32) + .collect(); + let end = nodes.len() - 1; + self.done = !self.all; + return Some(PathRows { + nodes, + edges, + boundaries: vec![0, end], + }); + } + Some(parents) if *next < parents.len() => { + let parent = parents[*next]; + *next += 1; + self.stack.push((parent, 0)); + } + _ => { + self.stack.pop(); + } + } + } + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::prelude::SessionContext; + + async fn graph( + edges: &[(i64, i64)], + max: u32, + shortest: Option, + ) -> (PathSpec, Arc) { + let n = edges + .iter() + .flat_map(|&(a, b)| [a, b]) + .max() + .unwrap_or(3) + .max(3); + let nodes = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from_iter_values(1..=n)) as ArrayRef, + )]) + .unwrap(); + let edge_batch = RecordBatch::try_from_iter([ + ( + "src", + Arc::new(Int64Array::from_iter_values(edges.iter().map(|e| e.0))) as ArrayRef, + ), + ( + "dst", + Arc::new(Int64Array::from_iter_values(edges.iter().map(|e| e.1))) as ArrayRef, + ), + ( + "weight", + Arc::new(Int64Array::from_iter_values(0..edges.len() as i64)) as ArrayRef, + ), + ]) + .unwrap(); + let binding = |variable: &str| Binding { + variable: variable.into(), + label: 0, + id_column: format!("{variable}__id"), + properties: HashMap::new(), + }; + let spec = PathSpec { + variable: "p".into(), + sources: Default::default(), + nodes: vec![NodeSource { + label: "Person".into(), + id: "id".into(), + schema: nodes.schema(), + }], + edges: vec![EdgeSource { + label: "R".into(), + source_label: 0, + target_label: 0, + source_id: "src".into(), + target_id: "dst".into(), + schema: edge_batch.schema(), + }], + bindings: vec![binding("a"), binding("b")], + steps: vec![Step { + types: vec![0], + direction: RelationshipDirection::Outgoing, + min: 0, + max, + relationship: None, + properties: HashMap::new(), + }], + shortest, + }; + let ctx = SessionContext::new(); + let mut plans = Vec::new(); + for batch in [nodes, edge_batch] { + plans.push( + ctx.read_batch(batch) + .unwrap() + .create_physical_plan() + .await + .unwrap(), + ); + } + let ready = Ready::build(&spec, &plans, &[None, None, None], ctx.task_ctx()) + .await + .unwrap(); + (spec, Arc::new(ready)) + } + + #[tokio::test] + async fn named_walk_keeps_edge_occurrences_and_never_reuses_one() { + let (spec, graph) = graph(&[(1, 2), (1, 2), (2, 3), (3, 3)], 20, None).await; + let mut walk = Walk::new(0, vec![Some(0), None]); + let mut paths = Vec::new(); + while let Some(path) = walk.next(&spec, &graph).await { + paths.push(path); + } + assert_eq!(paths.len(), 7); // identity, two one-hop, two two-hop, two with self-loop + assert!(paths + .iter() + .all(|p| p.edges.iter().collect::>().len() == p.edges.len())); + let two: Vec<_> = paths + .iter() + .filter(|p| p.edges.len() == 2) + .map(|p| p.edges.clone()) + .collect(); + assert_eq!(two, vec![vec![0, 2], vec![1, 2]]); + let array = paths_array(&paths, &graph.inventory.nodes, &graph.inventory.edges).unwrap(); + let rels = array.as_struct().column(3).as_list::(); + assert_eq!(rels.value(0).len(), 0); + let first = rels.value(1); + assert_eq!( + first + .as_struct() + .column(1) + .as_struct() + .column(1) + .as_primitive::() + .value(0), + 0 + ); + } + + #[tokio::test] + async fn shortest_values_preserve_parallel_edges_and_bounded_output() { + let (spec, graph) = graph( + &[(1, 2), (1, 2), (2, 3), (2, 3), (3, 3)], + u32::MAX, + Some(true), + ) + .await; + let batch = RecordBatch::try_from_iter([ + ("a__id", Arc::new(Int64Array::from(vec![1])) as ArrayRef), + ("b__id", Arc::new(Int64Array::from(vec![3])) as ArrayRef), + ]) + .unwrap(); + let mut pending = Pending::new(batch, graph, &spec).unwrap(); + let mut edges = Vec::new(); + loop { + let (rows, paths) = pending.next(&spec, 1).await.unwrap(); + assert!(rows.len() <= 1); + if rows.is_empty() { + break; + } + assert_eq!(paths[0].nodes, vec![0, 1, 2]); + edges.push(paths[0].edges.clone()); + } + edges.sort(); + assert_eq!(edges, vec![vec![0, 2], vec![0, 3], vec![1, 2], vec![1, 3]]); + } + + #[tokio::test] + async fn rejected_hub_yields_before_scanning_all_edges() { + let (spec, graph) = graph(&(2..100_002).map(|n| (1, n)).collect::>(), 1, None).await; + let mut walk = Walk::new(0, vec![Some(0), Some(0)]); + assert!(walk.next(&spec, &graph).await.is_some()); // zero hop + let mut next = Box::pin(walk.next(&spec, &graph)); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert!(walk.frames[0].next > 0 && walk.frames[0].next <= 1024); + } +} diff --git a/crates/lance-graph/src/path_value.rs b/crates/lance-graph/src/path_value.rs new file mode 100644 index 0000000..c4d5d14 --- /dev/null +++ b/crates/lance-graph/src/path_value.rs @@ -0,0 +1,273 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Arrow values for paths and their entities. Properties live in a marked child so +//! user property names cannot collide with identity or endpoint metadata. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, AsArray, Int64Array, ListArray, StructArray, UInt32Array}; +use arrow::buffer::OffsetBuffer; +use arrow::compute::take; +use arrow_schema::{DataType, Field, Fields}; +use datafusion::common::Result; + +pub(crate) const PROPERTIES: &str = "__lance_graph_properties"; +pub(crate) const IDENTITY: &str = "__lance_graph_identity"; +pub(crate) const START: &str = "__lance_graph_start"; +pub(crate) const END: &str = "__lance_graph_end"; +const ENTITY: &str = "lance_graph.entity"; + +pub(crate) fn entity_kind(data_type: &DataType) -> Option<&str> { + let DataType::Struct(fields) = data_type else { + return None; + }; + fields.first()?.metadata().get(ENTITY).map(String::as_str) +} + +pub(crate) fn entity_type( + properties: Fields, + kind: &str, + node_type: Option<&DataType>, +) -> DataType { + let mut fields = vec![ + Field::new(PROPERTIES, DataType::Struct(properties), false) + .with_metadata(HashMap::from([(ENTITY.into(), kind.into())])), + Field::new( + IDENTITY, + DataType::Struct( + vec![ + Field::new("label", DataType::Utf8, false), + Field::new("id", DataType::Int64, false), + ] + .into(), + ), + false, + ), + ]; + if let Some(node_type) = node_type { + fields.push(Field::new(START, node_type.clone(), true)); + fields.push(Field::new(END, node_type.clone(), true)); + } + DataType::Struct(fields.into()) +} + +pub(crate) fn path_type(node_type: DataType, relationship_type: DataType) -> DataType { + let list = |t| DataType::List(Arc::new(Field::new("item", t, true))); + DataType::Struct( + vec![ + Field::new("length", DataType::Int64, false), + Field::new("nodes", list(DataType::Int64), false), + Field::new("node_records", list(node_type), false), + Field::new("relationships", list(relationship_type), false), + ] + .into(), + ) +} + +pub(crate) fn struct_array(data_type: &DataType, columns: Vec) -> StructArray { + let DataType::Struct(fields) = data_type else { + unreachable!("entity/path struct") + }; + StructArray::new(fields.clone(), columns, None) +} + +/// Ordered row references, independent of property equality or edge endpoints. +#[derive(Debug, Clone)] +pub(crate) struct PathRows { + pub nodes: Vec, + pub edges: Vec, + /// Node positions at the written segment boundaries (variable segments may have many hops). + pub boundaries: Vec, +} + +pub(crate) fn paths_array( + paths: &[PathRows], + nodes: &StructArray, + edges: &StructArray, +) -> Result { + let node_rows = + UInt32Array::from_iter_values(paths.iter().flat_map(|p| p.nodes.iter().copied())); + let edge_rows = + UInt32Array::from_iter_values(paths.iter().flat_map(|p| p.edges.iter().copied())); + let selected_nodes = take(nodes, &node_rows, None)?; + let selected_edges = take(edges, &edge_rows, None)?; + let ids = selected_nodes + .as_struct() + .column(1) + .as_struct() + .column(1) + .clone(); + let list = |values: ArrayRef, lengths: Vec| -> ArrayRef { + Arc::new(ListArray::new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + OffsetBuffer::from_lengths(lengths), + values, + None, + )) + }; + let node_lengths: Vec<_> = paths.iter().map(|p| p.nodes.len()).collect(); + let edge_lengths: Vec<_> = paths.iter().map(|p| p.edges.len()).collect(); + let data_type = path_type(nodes.data_type().clone(), edges.data_type().clone()); + Ok(Arc::new(struct_array( + &data_type, + vec![ + Arc::new(Int64Array::from_iter_values( + paths.iter().map(|p| p.edges.len() as i64), + )), + list(ids, node_lengths.clone()), + list(selected_nodes, node_lengths), + list(selected_edges, edge_lengths), + ], + ))) +} + +/// A direct length/nodes query can keep its existing CSR plan. Values carried through +/// containers and node-property comprehensions need the actual path records. +pub(crate) fn needs_values(query: &crate::ast::CypherQuery) -> bool { + use crate::ast::{GraphPattern, ReadingClause, ValueExpression as V}; + let clauses: Vec<_> = query + .reading_clauses + .iter() + .chain(&query.post_with_reading_clauses) + .chain( + query + .further_with_stages + .iter() + .flat_map(|s| &s.reading_clauses), + ) + .collect(); + let mut names = std::collections::HashSet::new(); + for clause in &clauses { + if let ReadingClause::Match(m) = clause { + for pattern in &m.patterns { + if let GraphPattern::Path(path) = pattern { + if path.variable.is_some() { + return true; + } + if let Some(name) = path.shortest.as_ref().and_then(|s| s.variable.as_ref()) { + names.insert(name.to_lowercase()); + } + } + } + } + } + let mut needed = query + .union_branches + .iter() + .any(|branch| needs_values(&branch.query)) + || clauses.iter().any(|clause| match clause { + ReadingClause::Call(inner) => { + needs_values(inner) + || inner + .call_union_imports() + .unwrap_or_default() + .iter() + .any(|name| names.contains(name)) + } + _ => false, + }); + let mut inspect = |expression: &V| { + let mut value = expression.clone(); + crate::ast::transform_value(&mut value, &mut |v| match v { + V::ScalarFunction { name, args } + if matches!(name.to_lowercase().as_str(), "length" | "nodes") => + { + args.clear() + } + V::AggregateFunction { args, .. } + if args.iter().any( + |v| matches!(v, V::Variable(name) if names.contains(&name.to_lowercase())), + ) => + { + needed = true; + } + V::ScalarFunction { name, .. } + if matches!( + name.to_lowercase().as_str(), + "relationships" | "startnode" | "endnode" + ) => + { + needed = true + } + _ => {} + }); + crate::ast::transform_value(&mut value, &mut |v| { + if matches!(v, V::Variable(name) if names.contains(&name.to_lowercase())) { + needed = true; + } + }); + }; + for clause in clauses { + match clause { + ReadingClause::Call(_) => {} + ReadingClause::Unwind(u) => inspect(&u.expression), + ReadingClause::Match(m) => { + if let Some(w) = &m.where_clause { + inspect(&V::Boolean(Box::new(w.expression.clone()))); + } + } + } + } + for with in query + .with_clause + .iter() + .chain(query.further_with_stages.iter().map(|s| &s.with_clause)) + { + for item in &with.items { + inspect(&item.expression); + } + if let Some(order) = &with.order_by { + for item in &order.items { + inspect(&item.expression); + } + } + if let Some(filter) = &with.where_clause { + inspect(&V::Boolean(Box::new(filter.expression.clone()))); + } + } + for item in &query.return_clause.items { + inspect(&item.expression); + } + if let Some(order) = &query.order_by { + for item in &order.items { + inspect(&item.expression); + } + } + for filter in query + .where_clause + .iter() + .chain(&query.post_with_where_clause) + .chain( + query + .further_with_stages + .iter() + .filter_map(|s| s.where_clause.as_ref()), + ) + { + inspect(&V::Boolean(Box::new(filter.expression.clone()))); + } + needed +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn direct_shortest_queries_keep_the_csr_route() { + let prefix = "MATCH p = shortestPath((a:Person)-[:R*]-(b:Person)) "; + for (suffix, rich) in [ + ("RETURN length(p), nodes(p)", false), + ("RETURN [n IN nodes(p) | n.id]", false), + ("RETURN p", true), + ("RETURN relationships(p)", true), + ("RETURN length(head(collect(p)))", true), + ("WITH collect(p) AS ps UNWIND ps AS p RETURN nodes(p)", true), + ] { + let query = crate::parser::parse_cypher_query(&format!("{prefix}{suffix}")).unwrap(); + assert_eq!(needs_values(&query), rich, "{suffix}"); + } + } +} diff --git a/crates/lance-graph/src/pattern_order.rs b/crates/lance-graph/src/pattern_order.rs index b5333e4..bbf4dce 100644 --- a/crates/lance-graph/src/pattern_order.rs +++ b/crates/lance-graph/src/pattern_order.rs @@ -48,6 +48,7 @@ pub(crate) fn reorder_patterns( filter: Option<&BooleanExpression>, config: &GraphConfig, ) -> Option> { + if patterns.iter().any(|p| matches!(p, GraphPattern::Path(path) if path.variable.is_some() || path.shortest.is_some())) { return None; } let graph = PatternGraph::from_patterns(patterns)?; if graph.edges.is_empty() || !graph.is_connected() { return None; @@ -415,6 +416,7 @@ impl PatternGraph { start_node, segments, shortest: None, + variable: None, })); } // The next path starts at a bound node that still has an unused edge. diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index 358544e..aff6fa4 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -386,7 +386,11 @@ impl SemanticAnalyzer { // Register end node self.register_node_variable(&segment.end_node)?; } - if let Some(path_var) = path.shortest.as_ref().and_then(|s| s.variable.as_ref()) { + if let Some(path_var) = path + .variable + .as_ref() + .or_else(|| path.shortest.as_ref().and_then(|s| s.variable.as_ref())) + { self.register_projection_alias(path_var); if let Some(info) = self.variables.get_mut(&path_var.to_lowercase()) { info.variable_type = VariableType::Path; @@ -589,7 +593,7 @@ impl SemanticAnalyzer { match function_name.as_str() { "tolower" | "lower" | "toupper" | "upper" | "abs" | "round" | "floor" | "ceil" | "sqrt" | "trim" | "ltrim" | "rtrim" | "reverse" | "last" - | "tail" | "keys" | "labels" => { + | "tail" | "keys" | "labels" | "relationships" | "startnode" | "endnode" => { if args.len() != 1 { return Err(GraphError::PlanError { message: format!( @@ -1306,6 +1310,7 @@ mod tests { end_node: end, }], shortest: None, + variable: None, }; let query = CypherQuery { @@ -1389,6 +1394,7 @@ mod tests { end_node: end, }], shortest: None, + variable: None, }; let query = CypherQuery { @@ -1516,6 +1522,7 @@ mod tests { end_node: end, }], shortest: None, + variable: None, }; let query = CypherQuery { @@ -1580,6 +1587,7 @@ mod tests { }, ], shortest: None, + variable: None, }; // Custom config that knows both relationship types to avoid warnings muddying the assertion diff --git a/crates/lance-graph/src/shortest_path.rs b/crates/lance-graph/src/shortest_path.rs index 9c37dec..024889c 100644 --- a/crates/lance-graph/src/shortest_path.rs +++ b/crates/lance-graph/src/shortest_path.rs @@ -349,6 +349,7 @@ impl ExecutionPlan for ShortestPathExec { &spec, &schema, ) + .await } }); Ok(Box::pin(RecordBatchStreamAdapter::new(out_schema, stream))) @@ -386,7 +387,7 @@ async fn build_adjacency( /// Breadth-first search from `source` until every node in `targets` is reached (or the /// frontier empties, or `max_hops`). Returns each reached node's depth and, for every node /// at depth d > 0, its predecessors at depth d - 1. -fn search( +async fn search( adjacency: &Neighbors, source: i64, targets: &HashSet, @@ -400,7 +401,7 @@ fn search( let mut level = 0; while !frontier.is_empty() && !remaining.is_empty() && level < max_hops { level += 1; - let next = expand_level(adjacency, &frontier, level, &mut depth, &mut parents); + let next = expand_level(adjacency, &frontier, level, &mut depth, &mut parents).await; // Finish the level before stopping, so every shortest path is recorded. for n in &next { remaining.remove(n); @@ -411,7 +412,7 @@ fn search( } /// Finish the whole level so parallel edges and alternative shortest paths survive. -fn expand_level( +async fn expand_level( adjacency: &Neighbors, frontier: &[i64], level: u32, @@ -419,8 +420,17 @@ fn expand_level( parents: &mut FastMap>, ) -> Vec { let mut next = Vec::new(); + let mut work = 0_usize; for &node in frontier { - adjacency.for_each_neighbor(node, |n| { + let mut position = crate::csr_expand::NeighborPosition::default(); + while let Some(candidate) = adjacency.next_candidate(node, &mut position) { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } + let Some(n) = candidate else { + continue; + }; let n = n as i64; match depth.get(&n) { Some(&d) if d == level => parents.entry(n).or_default().push(node), @@ -431,14 +441,14 @@ fn expand_level( next.push(n); } } - }); + } } next } /// Search a single pair using complete layers from the smaller frontier. At the /// first intersection every meeting node lies on the same depth cut of a shortest path. -fn search_bidirectional( +pub(crate) async fn search_bidirectional( forward: &Neighbors, reverse: &Neighbors, source: i64, @@ -459,7 +469,7 @@ fn search_bidirectional( while !frontier.is_empty() && !reverse_frontier.is_empty() && level + reverse_level < max_hops { if frontier.len() <= reverse_frontier.len() { level += 1; - frontier = expand_level(forward, &frontier, level, &mut depth, &mut parents); + frontier = expand_level(forward, &frontier, level, &mut depth, &mut parents).await; } else { reverse_level += 1; reverse_frontier = expand_level( @@ -468,7 +478,8 @@ fn search_bidirectional( reverse_level, &mut reverse_depth, &mut reverse_parents, - ); + ) + .await; } let mut pending: Vec<_> = frontier .iter() @@ -480,7 +491,12 @@ fn search_bidirectional( } // Reverse only the DAG beyond this cut; stitching at several depths duplicates paths. let mut visited = HashSet::new(); + let mut work = 0_usize; while let Some(node) = pending.pop() { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } if !visited.insert(node) { continue; } @@ -499,10 +515,15 @@ fn search_bidirectional( /// Shortest paths from the search source to `target`, as node ids from source to target: /// all of them, or the first found. -fn paths_to(target: i64, parents: &FastMap>, all: bool) -> Vec> { +async fn paths_to(target: i64, parents: &FastMap>, all: bool) -> Vec> { let mut out = Vec::new(); let mut stack = vec![vec![target]]; + let mut work = 0_usize; while let Some(path) = stack.pop() { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } let head = *path.last().expect("non-empty"); match parents.get(&head) { None => { @@ -526,7 +547,7 @@ fn paths_to(target: i64, parents: &FastMap>, all: bool) -> Vec = by_source - .iter() - .map(|(&s, ts)| { - let result = match reverse_adjacency.filter(|_| ts.len() == 1) { - Some(reverse) => search_bidirectional( + let mut searches = FastMap::default(); + for (&source, targets) in &by_source { + let result = match reverse_adjacency.filter(|_| targets.len() == 1) { + Some(reverse) => { + search_bidirectional( adjacency, reverse, - s, - *ts.iter().next().expect("one target"), + source, + *targets.iter().next().unwrap(), spec.max_hops, - ), - None => search(adjacency, s, ts, spec.max_hops), - }; - (s, result) - }) - .collect(); + ) + .await + } + None => search(adjacency, source, targets, spec.max_hops).await, + }; + searches.insert(source, result); + } let mut rows = Vec::new(); let mut lengths = Vec::new(); let mut nodes = ListBuilder::new(Int64Builder::new()); @@ -580,7 +602,7 @@ fn search_batch( if d < spec.min_hops || d > spec.max_hops { continue; } - for path in paths_to(t, parents, spec.all) { + for path in paths_to(t, parents, spec.all).await { rows.push(i as u32); lengths.push(d as i64); nodes.values().append_slice(&path); @@ -628,7 +650,7 @@ fn plan_error(message: &str) -> GraphError { mod tests { use super::*; - fn adjacencies(edges: &[(i64, i64)]) -> (Arc, Arc) { + async fn adjacencies(edges: &[(i64, i64)]) -> (Arc, Arc) { let schema = Arc::new(Schema::new(vec![ Field::new("src", DataType::Int64, false), Field::new("dst", DataType::Int64, false), @@ -652,40 +674,33 @@ mod tests { futures::stream::once(async move { Ok(batch) }), )) as SendableRecordBatchStream }; - let rt = tokio::runtime::Builder::new_current_thread() - .build() - .unwrap(); - let out = rt - .block_on(Adjacency::build(stream(), "src", "dst")) - .unwrap(); - let inc = rt - .block_on(Adjacency::build(stream(), "dst", "src")) - .unwrap(); + let out = Adjacency::build(stream(), "src", "dst").await.unwrap(); + let inc = Adjacency::build(stream(), "dst", "src").await.unwrap(); (Arc::new(out), Arc::new(inc)) } - fn neighbors(edges: &[(i64, i64)]) -> Neighbors { - let (out, inc) = adjacencies(edges); + async fn neighbors(edges: &[(i64, i64)]) -> Neighbors { + let (out, inc) = adjacencies(edges).await; Neighbors::undirected(out, inc) } - #[test] - fn finds_every_shortest_path_and_only_shortest_ones() { + #[tokio::test] + async fn finds_every_shortest_path_and_only_shortest_ones() { // Two 2-hop routes 1-2-4 and 1-3-4, and a longer 1-5-6-4. - let adj = neighbors(&[(1, 2), (2, 4), (1, 3), (3, 4), (1, 5), (5, 6), (6, 4)]); - let (depth, parents) = search(&adj, 1, &HashSet::from([4]), 10); + let adj = neighbors(&[(1, 2), (2, 4), (1, 3), (3, 4), (1, 5), (5, 6), (6, 4)]).await; + let (depth, parents) = search(&adj, 1, &HashSet::from([4]), 10).await; assert_eq!(depth[&4], 2); - let mut all = paths_to(4, &parents, true); + let mut all = paths_to(4, &parents, true).await; all.sort(); assert_eq!(all, vec![vec![1, 2, 4], vec![1, 3, 4]]); - assert_eq!(paths_to(4, &parents, false).len(), 1); + assert_eq!(paths_to(4, &parents, false).await.len(), 1); // Beyond the hop limit, nothing is reached. - let (depth, _) = search(&adj, 1, &HashSet::from([4]), 1); + let (depth, _) = search(&adj, 1, &HashSet::from([4]), 1).await; assert!(!depth.contains_key(&4)); } - #[test] - fn bidirectional_keeps_every_meeting_node_and_parallel_edge() { + #[tokio::test] + async fn bidirectional_keeps_every_meeting_node_and_parallel_edge() { // A wide first layer makes the target's side expand. Two meeting nodes, // including parallel edges on both sides of the cut, yield six paths. let (out, inc) = adjacencies(&[ @@ -700,11 +715,12 @@ mod tests { (4, 5), (5, 5), (5, 1), - ]); + ]) + .await; let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc)); - let (depth, parents) = search_bidirectional(&forward, &reverse, 1, 5, 3); + let (depth, parents) = search_bidirectional(&forward, &reverse, 1, 5, 3).await; assert_eq!(depth[&5], 3); - let mut all = paths_to(5, &parents, true); + let mut all = paths_to(5, &parents, true).await; all.sort(); assert_eq!( all, @@ -713,16 +729,17 @@ mod tests { .chain(vec![vec![1, 3, 4, 5]; 2]) .collect::>() ); - let one = paths_to(5, &parents, false); + let one = paths_to(5, &parents, false).await; assert_eq!(one.len(), 1); assert!(all.contains(&one[0])); assert!(!search_bidirectional(&forward, &reverse, 1, 5, 2) + .await .0 .contains_key(&5)); } - #[test] - fn bidirectional_matches_forward_search_for_small_graphs() { + #[tokio::test] + async fn bidirectional_matches_forward_search_for_small_graphs() { // All directed graphs on three nodes, including disconnected pairs and cycles. let edges = [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)]; for mask in 0..(1 << edges.len()) { @@ -731,7 +748,7 @@ mod tests { .enumerate() .filter_map(|(bit, &edge)| (mask & (1 << bit) != 0).then_some(edge)) .collect(); - let (out, inc) = adjacencies(&selected); + let (out, inc) = adjacencies(&selected).await; let directed = Neighbors::directed(out.clone()); let reverse = Neighbors::directed(inc.clone()); let undirected = Neighbors::undirected(out, inc); @@ -740,17 +757,18 @@ mod tests { for target in 0..3 { for max_hops in 0..=3 { let (expected_depth, expected_parents) = - search(forward, source, &HashSet::from([target]), max_hops); + search(forward, source, &HashSet::from([target]), max_hops).await; let (depth, parents) = - search_bidirectional(forward, reverse, source, target, max_hops); + search_bidirectional(forward, reverse, source, target, max_hops) + .await; assert_eq!( depth.get(&target), expected_depth.get(&target), "graph {mask}, {source}->{target}, max {max_hops}" ); if depth.contains_key(&target) { - let mut expected = paths_to(target, &expected_parents, true); - let mut actual = paths_to(target, &parents, true); + let mut expected = paths_to(target, &expected_parents, true).await; + let mut actual = paths_to(target, &parents, true).await; expected.sort(); actual.sort(); assert_eq!( @@ -765,8 +783,8 @@ mod tests { } } - #[test] - fn bidirectional_avoids_expanding_the_wide_middle_layer() { + #[tokio::test] + async fn bidirectional_avoids_expanding_the_wide_middle_layer() { let mut edges = Vec::new(); for branch in 1..=100 { edges.push((0, branch)); @@ -774,13 +792,13 @@ mod tests { edges.push((branch, 101 + (branch - 1) * 100 + leaf)); } } - let (out, inc) = adjacencies(&edges); + let (out, inc) = adjacencies(&edges).await; let (forward, reverse) = (Neighbors::directed(out), Neighbors::directed(inc)); - let (depth, parents) = search_bidirectional(&forward, &reverse, 0, 101, 2); - assert_eq!(paths_to(101, &parents, true), vec![vec![0, 1, 101]]); + let (depth, parents) = search_bidirectional(&forward, &reverse, 0, 101, 2).await; + assert_eq!(paths_to(101, &parents, true).await, vec![vec![0, 1, 101]]); assert_eq!(depth.len(), 102); // Root, first layer and stitched target. assert_eq!( - search(&forward, 0, &HashSet::from([101]), 2).0.len(), + search(&forward, 0, &HashSet::from([101]), 2).await.0.len(), 10_101 ); } diff --git a/docs/api-reference.md b/docs/api-reference.md index c1ae261..737354d 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -140,6 +140,12 @@ be opened. --- +Cypher named paths can be returned as Arrow structs containing `length`, `nodes`, `node_records` and +`relationships`. They retain their values through `collect()` and `UNWIND`; `nodes(path)` returns integer +keys, and `relationships(path)` returns typed relationship records. `startNode(r)` and `endNode(r)` return +the stored endpoint node records. Their occurrence identities are query-local. See +[path values and their memory cost](cypher.md#path-values). + ## `CypherQuery` diff --git a/docs/csr.md b/docs/csr.md index a9677e7..9db5951 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -132,6 +132,12 @@ When those hold, it covers: Everything else in the query (filters, aggregations, sorts and any hop the CSR can't take) is planned by DataFusion as usual. You can mix: a three-hop query might run two hops on the CSR and one as a join. +Named path values and `relationships(path)` use a query-local inventory of actual edge occurrences because +CSR neighbors do not carry edge identity or properties. This also preserves parallel edges when a path +passes through `collect()` and `UNWIND`. The inventory materializes the participating node and relationship +tables; bounded traversal batches do not bound that inventory. See [path values](cypher.md#path-values). +Shortest-path queries that only read length or node keys keep the CSR route below. + For `shortestPath` and `allShortestPaths`, a source with one target in an input batch searches from both ends using the outgoing and incoming CSR. It expands the smaller frontier and finishes the whole layer before joining paths, so it keeps all shortest alternatives. Several targets for the same source share diff --git a/docs/cypher.md b/docs/cypher.md index 9072370..77dfee4 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -287,7 +287,8 @@ RETURN length(p), nodes(p) expansion for the whole range; the join fallback uses one branch per length. - **Variable-length alternation** follows each relationship type's configured endpoint labels. With suitable CSR and node indexes, a mixed path such as `[:hasType|isSubclassOf*1..3]` stays in one operator. - Undirected or named variable-length alternation remains unsupported. + Named paths also support typed variable-length alternation through their edge inventory. + Unnamed undirected variable-length alternation remains unsupported. - **Undirected variable-length paths** skip immediate backtracking (`a-b-a`). - **`shortestPath` and `allShortestPaths`** need exactly one relationship type, no relationship variable, and named end nodes. Their length is unbounded unless you give a maximum. `nodes(p)` returns the node keys. @@ -296,7 +297,51 @@ RETURN length(p), nodes(p) duplicate node sequences that use different parallel edges. A single source-target pair searches from both ends when reverse adjacency is available. Several targets for the same source in an input batch share one forward search. -- **Named paths** (`p = ...`) are only supported with `shortestPath` and `allShortestPaths`. +- **Named paths** (`p = (a)-[:R]->(b)`) retain the ordered node and relationship records. Named + variable-length paths use each relationship at most once, including self-loops and parallel edges. + Use `relationships(p)` for variable-length relationship lists; binding a relationship variable inside + a named variable-length segment is not supported. + +### Path values + +A named path can pass through `WITH`, `collect()` and `UNWIND` without losing its edges. Parallel edges +with identical endpoints remain distinct, including when their properties are identical. + +```python +paths = engine.execute(""" + MATCH path = (a:Person {id: 1})-[:KNOWS]->(b:Person) + WITH collect(path) AS paths + UNWIND paths AS path + RETURN [n IN nodes(path) | n.name] AS names, + [r IN relationships(path) | endNode(r).id] AS targets + ORDER BY names +""") +print(paths.to_pylist()) +# [{'names': ['Alice', 'Bob'], 'targets': [2]}, +# {'names': ['Alice', 'Carol'], 'targets': [3]}] +``` + +`nodes(path)` still returns integer keys when returned directly. A property comprehension over it uses +stored node records when the path is a value, so equal keys in different labels remain different nodes. +`startNode(r)` and `endNode(r)` return the relationship's stored source and target, even when the path +traversed that relationship backwards. Property access, `keys()` and node `labels()` work on these values. +A zero-hop path has one node and an empty relationship list. All path functions return null for a null path. + +`RETURN path` produces an Arrow struct with `length`, `nodes`, `node_records` and `relationships` fields. +Entity records store user properties in `__lance_graph_properties` and `{label, id}` identity metadata +in `__lance_graph_identity`. Relationship records also carry the node records `__lance_graph_start` and +`__lance_graph_end`. These metadata fields are distinct from user properties, even if a property has the +same name. `keys(entity)` lists only non-null user properties; `keys(map)` includes null-valued keys. +Prefer the Cypher property and endpoint functions when writing queries. Relationship identities identify actual occurrences within this query's inventory, not persistent +IDs that can be compared across queries or refreshes. Rich paths support the same signed integer and +`uint8`/`uint16`/`uint32` keys as the CSR. Mixed labels may share compatible properties; conflicting Arrow +property types fail explicitly. Expanding a graph pattern from a mixed-label value remains unsupported. + +A path value needs a query-local inventory of the participating node and relationship tables, including +properties. It therefore materializes those tables, even with `use_csr=True`. Traversal output is emitted +in batches and named ordinary paths keep a cursor per hop, but inventory memory grows with the graph. +Queries asking only for shortest-path length or node keys, including existing homogeneous node-property +comprehensions, retain the usual CSR execution route. ### Expressions and predicates @@ -329,7 +374,7 @@ RETURN length(p), nodes(p) | Lists | `size`, `head`, `last`, `tail`, `reverse`, `range(start, end [, step])` (inclusive) | | Math | `abs`, `round`, `floor`, `ceil`, `sqrt` | | Temporal | `date`, `datetime`, `date_part('month', d)`, `to_epoch_ms`, `duration.inSeconds(a, b).seconds` | -| Paths | `length(p)`, `nodes(p)` on `shortestPath` variables | +| Paths | `length(p)`, `nodes(p)`, `relationships(p)`; `startNode(r)`, `endNode(r)` on relationship values | | Vectors | `vector_distance(a, b, metric)`, `vector_similarity(a, b, metric)`; see [Vector search](vector-search.md) | The string and list functions follow the [Cypher v9 reference](https://s3.amazonaws.com/artifacts.opencypher.org/openCypher9.pdf), @@ -386,7 +431,7 @@ print(result.to_pylist()) # [{'labels': ['Person'], 'properties': 3}] Any other function raises `ValueError` with `Cypher function 'name' is not implemented`. This is a bounded function set, not a claim of full Cypher conformance. Remaining functions include `reduce`, -the list predicates `any` / `all` / `none` / `single`, and `relationships(p)`. +the list predicates `any` / `all` / `none` / `single`. When `size()` consumes an unmapped list comprehension, such as `size([x IN values WHERE x > 0])`, the planner counts retained elements directly. Each duplicate and each retained null element still contributes one; diff --git a/python/python/tests/test_path_values.py b/python/python/tests/test_path_values.py new file mode 100644 index 0000000..fb90c00 --- /dev/null +++ b/python/python/tests/test_path_values.py @@ -0,0 +1,350 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Actual edge records survive path values, including duplicate endpoint pairs.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture(params=[False, True], ids=["scan", "csr"]) +def engine(request): + people = pa.table({"id": [1, 2, 3, 4], "name": ["a", "b", "c", "d"]}) + edges = pa.table( + {"src": [1, 1, 2, 3, 3], "dst": [2, 2, 3, 3, 99], "weight": [5, 7, 11, 13, 17]} + ) + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + return CypherEngine( + config, {"Person": people, "KNOWS": edges}, use_csr=request.param + ) + + +def test_named_path_records(engine): + rows = engine.execute( + "MATCH p = (a:Person {id: 1})-[:KNOWS]->(b:Person)" + "-[:KNOWS]->(c:Person {id: 3}) " + "RETURN length(p) AS length, nodes(p) AS ids, " + "[n IN nodes(p) | n.name] AS names, " + "[r IN relationships(p) | r.weight] AS weights" + ).to_pylist() + assert sorted(rows, key=lambda r: r["weights"]) == [ + {"length": 2, "ids": [1, 2, 3], "names": ["a", "b", "c"], "weights": [5, 11]}, + {"length": 2, "ids": [1, 2, 3], "names": ["a", "b", "c"], "weights": [7, 11]}, + ] + + +@pytest.mark.parametrize("search", ["shortestPath", "allShortestPaths"]) +def test_shortest_path_survives_collect_unwind(engine, search): + rows = engine.execute( + f"MATCH p = {search}((a:Person {{id: 1}})-[:KNOWS*0..]->(b:Person {{id: 3}})) " + "WITH collect(p) AS paths UNWIND paths AS p " + "RETURN nodes(p) AS ids, [r IN relationships(p) | r.weight] AS weights, " + "[n IN nodes(p) | n.name] AS names" + ).to_pylist() + expected = [[5, 11], [7, 11]] + assert len(rows) == (2 if search == "allShortestPaths" else 1) + assert all(r["ids"] == [1, 2, 3] and r["names"] == ["a", "b", "c"] for r in rows) + assert all(r["weights"] in expected for r in rows) + if search == "allShortestPaths": + assert sorted(r["weights"] for r in rows) == expected + + +def test_reverse_path_keeps_stored_relationship_endpoints(engine): + rows = engine.execute( + "MATCH p = (b:Person {id: 2})<-[r:KNOWS]-(a:Person {id: 1}) " + "UNWIND relationships(p) AS edge RETURN nodes(p) AS ids, " + "edge.weight AS weight, startNode(edge).id AS source, " + "endNode(edge).id AS target" + ).to_pylist() + assert sorted(rows, key=lambda r: r["weight"]) == [ + {"ids": [2, 1], "weight": 5, "source": 1, "target": 2}, + {"ids": [2, 1], "weight": 7, "source": 1, "target": 2}, + ] + + +@pytest.mark.parametrize("shortest", [False, True]) +def test_zero_path_records(engine, shortest): + pattern = "(a:Person {id: 3})-[:KNOWS*0..0]->(b:Person {id: 3})" + if shortest: + pattern = f"shortestPath({pattern})" + rows = engine.execute( + f"MATCH p = {pattern} RETURN p, length(p) AS length, " + "nodes(p) AS ids, relationships(p) AS edges" + ).to_pylist() + assert len(rows) == 1 + assert rows[0]["length"] == 0 + assert rows[0]["ids"] == [3] + assert rows[0]["edges"] == [] + assert rows[0]["p"]["node_records"][0]["__lance_graph_properties"]["name"] == "c" + + +def test_named_trails_do_not_reuse_a_relationship(engine): + rows = engine.execute( + "MATCH p = (a:Person {id: 3})-[:KNOWS*1..3]->(b:Person) " + "RETURN nodes(p) AS ids, [r IN relationships(p) | r.weight] AS weights" + ).to_pylist() + assert rows == [{"ids": [3, 3], "weights": [13]}] + + +def test_optional_path_null_values(engine): + rows = engine.execute( + "MATCH (a:Person {id: 1}), (b:Person {id: 4}) " + "OPTIONAL MATCH p = shortestPath((a)-[:KNOWS*0..]->(b)) " + "RETURN p, relationships(p) AS edges, nodes(p) AS ids" + ).to_pylist() + assert rows == [{"p": None, "edges": None, "ids": None}] + + +def test_path_entities_have_properties_and_labels(engine): + rows = engine.execute( + "MATCH p = (a:Person {id: 1})-[:KNOWS]->(b:Person) " + "UNWIND relationships(p) AS r " + "RETURN keys(r) AS properties, labels(startNode(r)) AS labels" + ).to_pylist() + assert rows == [{"properties": ["dst", "src", "weight"], "labels": ["Person"]}] * 2 + + +def test_fast_shortest_plan_retained(engine): + query = ( + "MATCH p = shortestPath((a:Person {id: 1})-[:KNOWS*0..]->(b:Person {id: 3})) " + "RETURN [n IN nodes(p) | n.id] AS ids" + ) + assert "PathMatch" not in engine.explain(query) + assert engine.execute(query).to_pylist() == [{"ids": [1, 2, 3]}] + + +def test_mixed_label_path_identity(): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .node_table("Place", "id") + .rel_table("LIVES", "Person", "Place", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table({"id": [1], "name": ["person"]}), + "Place": pa.table({"id": [1], "name": ["place"]}), + "LIVES": pa.table({"src": [1], "dst": [1]}), + }, + ) + rows = engine.execute( + "MATCH p = (a:Person)-[:LIVES]->(b:Place) RETURN nodes(p) AS ids, " + "[n IN nodes(p) | n.name] AS names, [n IN nodes(p) | labels(n)] AS labels" + ).to_pylist() + assert rows == [ + {"ids": [1, 1], "names": ["person", "place"], "labels": [["Person"], ["Place"]]} + ] + + +def test_named_variable_alternation_tracks_labels(): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .node_table("Place", "id") + .rel_table("LIVES", "Person", "Place", "src", "dst") + .rel_table("INSIDE", "Place", "Place", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table({"id": [1], "name": ["person"]}), + "Place": pa.table({"id": [1, 2], "name": ["place", "country"]}), + "LIVES": pa.table({"src": [1], "dst": [1]}), + "INSIDE": pa.table({"src": [1], "dst": [2]}), + }, + use_csr=True, + ) + query = ( + "MATCH p = (a:Person)-[:LIVES|INSIDE*0..20]->(b:Place) " + "RETURN nodes(p) AS ids, [n IN nodes(p) | n.name] AS names ORDER BY ids" + ) + assert engine.execute(query).to_pylist() == [ + {"ids": [1, 1], "names": ["person", "place"]}, + {"ids": [1, 1, 2], "names": ["person", "place", "country"]}, + ] + assert engine.explain(query).count("PathMatch:") == 1 + + +def test_identical_parallel_occurrences_are_distinct(): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table({"id": [1, 2]}), + "R": pa.table({"src": [1, 1], "dst": [2, 2], "weight": [5, 5]}), + }, + ) + rows = engine.execute( + "MATCH p = allShortestPaths((a:Person {id: 1})" + "-[:R*]->(b:Person {id: 2})) RETURN p" + ).to_pylist() + relationships = [r["p"]["relationships"][0] for r in rows] + assert len(relationships) == 2 + assert len({r["__lance_graph_identity"]["id"] for r in relationships}) == 2 + assert [r["__lance_graph_properties"]["weight"] for r in relationships] == [5, 5] + + +@pytest.mark.parametrize("bad_type", [pa.string(), pa.uint64()]) +def test_rich_path_rejects_unsupported_ids(bad_type): + values = ["1", "2"] if pa.types.is_string(bad_type) else [1, 2] + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table({"id": pa.array(values, type=bad_type)}), + "R": pa.table( + { + "src": pa.array(values[:1], type=bad_type), + "dst": pa.array(values[1:], type=bad_type), + } + ), + }, + ) + with pytest.raises(NotImplementedError, match="unsupported type"): + engine.execute("MATCH p = (a:Person)-[:R]->(b:Person) RETURN p") + + +def test_mixed_path_rejects_conflicting_property_types(): + config = ( + GraphConfig.builder() + .node_table("A", "id") + .node_table("B", "id") + .rel_table("R", "A", "B", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "A": pa.table({"id": [1], "name": ["name"]}), + "B": pa.table({"id": [1], "name": [99]}), + "R": pa.table({"src": [1], "dst": [1]}), + }, + ) + with pytest.raises(ValueError, match="conflicting Arrow types"): + engine.execute("MATCH p = (a:A)-[:R]->(b:B) RETURN p") + + +def test_path_inventory_keeps_lance_snapshot(tmp_path): + import lance + + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + ) + people_path, edges_path = ( + str(tmp_path / "people.lance"), + str(tmp_path / "edges.lance"), + ) + lance.write_dataset( + pa.table({"id": [1, 2], "name": ["old", "target"]}), people_path + ) + lance.write_dataset(pa.table({"src": [1], "dst": [2], "weight": [5]}), edges_path) + engine = CypherEngine( + config, + {"Person": lance.dataset(people_path), "R": lance.dataset(edges_path)}, + use_csr=True, + ) + query = ( + "MATCH p = (a:Person {id: 1})-[:R]->(b:Person) " + "RETURN [n IN nodes(p) | n.name] AS names, " + "[r IN relationships(p) | r.weight] AS weights" + ) + reader = engine.execute_batches(query) + lance.write_dataset( + pa.table({"src": [1], "dst": [2], "weight": [7]}), edges_path, mode="overwrite" + ) + engine.refresh() + assert reader.read_all().to_pylist() == [ + {"names": ["old", "target"], "weights": [5]} + ] + assert engine.execute(query).to_pylist() == [ + {"names": ["old", "target"], "weights": [7]} + ] + + +def test_bound_relationship_endpoints_survive_collect(engine): + rows = engine.execute( + "MATCH (a:Person {id: 1})-[r:KNOWS]->(b:Person {id: 2}) " + "WITH collect(r) AS rs UNWIND rs AS edge " + "RETURN edge.weight AS weight, keys(edge) AS properties, " + "startNode(edge).name AS source, endNode(edge).name AS target, " + "labels(endNode(edge)) AS labels ORDER BY weight" + ).to_pylist() + assert rows == [ + { + "weight": weight, + "properties": ["dst", "src", "weight"], + "source": "a", + "target": "b", + "labels": ["Person"], + } + for weight in [5, 7] + ] + + +def test_entity_metadata_does_not_shadow_properties(): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table( + {"id": [1, 2], "nullable": pa.array([None, None], pa.string())} + ), + "R": pa.table({"src": [1], "dst": [2], "__lance_graph_identity": [42]}), + }, + ) + rows = engine.execute( + "MATCH p = (a:Person)-[:R]->(b:Person) UNWIND relationships(p) AS r " + "RETURN r.__lance_graph_identity AS user_value, keys(r) AS properties, " + "keys(startNode(r)) AS node_properties, keys({nullable: null}) AS map_keys" + ).to_pylist() + assert rows == [ + { + "user_value": 42, + "properties": ["__lance_graph_identity", "dst", "src"], + "node_properties": ["id"], + "map_keys": ["nullable"], + } + ] + + +def test_multiple_path_operators_share_occurrence_identity(engine): + rows = engine.execute( + "MATCH p = (a:Person {id: 1})-[:KNOWS]->(b:Person {id: 2}) " + "MATCH q = (a)-[:KNOWS]->(b) " + "RETURN head(relationships(p)) AS first, head(relationships(q)) AS second" + ).to_pylist() + assert len(rows) == 4 + ids = {} + for row in rows: + for edge in row.values(): + weight = edge["__lance_graph_properties"]["weight"] + identity = edge["__lance_graph_identity"] + assert ids.setdefault(weight, identity) == identity + assert ids[5] != ids[7] From 8e1563f927368a3578e4cec744277888cf1dfd6a Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:01:45 -0700 Subject: [PATCH 26/66] feat(cypher): add ordered reduce and list predicates Part of #20. Keep fold variables local, preserve list order and null semantics, and correlate graph predicates with iterator bindings. Document supported forms and explicit body limits. --- crates/lance-graph/src/ast.rs | 79 +++++ .../builder/aggregate_ops.rs | 40 +-- .../datafusion_planner/builder/basic_ops.rs | 24 +- .../builder/comprehension_ops.rs | 22 +- .../builder/join_builder.rs | 4 +- .../src/datafusion_planner/expression.rs | 90 +++-- .../src/datafusion_planner/join_ops.rs | 3 +- .../src/datafusion_planner/list_folds.rs | 328 ++++++++++++++++++ .../lance-graph/src/datafusion_planner/mod.rs | 1 + .../datafusion_planner/scalar_functions.rs | 30 ++ .../src/datafusion_planner/scan_ops.rs | 8 +- crates/lance-graph/src/logical_plan.rs | 112 +++++- .../lance-graph/src/parameter_substitution.rs | 16 + crates/lance-graph/src/parser.rs | 90 +++++ crates/lance-graph/src/pattern_order.rs | 19 +- crates/lance-graph/src/semantic.rs | 56 ++- docs/cypher.md | 41 ++- python/python/tests/test_list_folds.py | 240 +++++++++++++ 18 files changed, 1117 insertions(+), 86 deletions(-) create mode 100644 crates/lance-graph/src/datafusion_planner/list_folds.rs create mode 100644 python/python/tests/test_list_folds.py diff --git a/crates/lance-graph/src/ast.rs b/crates/lance-graph/src/ast.rs index 3fac617..c3930b7 100644 --- a/crates/lance-graph/src/ast.rs +++ b/crates/lance-graph/src/ast.rs @@ -551,6 +551,23 @@ pub enum ValueExpression { filter: Option>, map: Option>, }, + /// A lexically scoped variable in a reduce or list predicate body. + LocalVariable(String), + /// `reduce(accumulator = initial, variable IN list | expression)`. + Reduce { + accumulator: String, + initial: Box, + variable: String, + list: Box, + expression: Box, + }, + /// `all/any/none/single(variable IN list WHERE predicate)`. + ListPredicate { + kind: String, + variable: String, + list: Box, + predicate: Box, + }, /// Vector literal: [0.1, 0.2, 0.3] /// Represents an inline vector for similarity search VectorLiteral(Vec), @@ -770,7 +787,24 @@ pub fn transform_value(expr: &mut ValueExpression, f: &mut dyn FnMut(&mut ValueE } transform_value(map, f); } + ValueExpression::Reduce { + initial, + list, + expression, + .. + } => { + transform_value(initial, f); + transform_value(list, f); + transform_value(expression, f); + } + ValueExpression::ListPredicate { + list, predicate, .. + } => { + transform_value(list, f); + transform_value(predicate, f); + } ValueExpression::Variable(_) + | ValueExpression::LocalVariable(_) | ValueExpression::Property(_) | ValueExpression::Literal(_) | ValueExpression::Parameter(_) @@ -779,6 +813,51 @@ pub fn transform_value(expr: &mut ValueExpression, f: &mut dyn FnMut(&mut ValueE f(expr); } +// Mark only the body's local references; the input and initializer retain outer scope. +pub(crate) fn mark_locals(expr: &mut ValueExpression, names: &[&str]) { + transform_value(expr, &mut |value| match value { + ValueExpression::Variable(name) if names.iter().any(|n| n.eq_ignore_ascii_case(name)) => { + *value = ValueExpression::LocalVariable(name.to_lowercase()); + } + ValueExpression::Property(p) + if names.iter().any(|n| n.eq_ignore_ascii_case(&p.variable)) => + { + *value = ValueExpression::Field { + base: Box::new(ValueExpression::LocalVariable(p.variable.to_lowercase())), + field: p.property.clone(), + }; + } + _ => {} + }); +} + +/// Expressions that need relational planning rather than scalar lambda execution. +pub(crate) fn needs_relational_evaluation(expr: &ValueExpression) -> bool { + fn boolean_has_pattern(expr: &BooleanExpression) -> bool { + match expr { + BooleanExpression::PatternExists { .. } => true, + BooleanExpression::And(a, b) | BooleanExpression::Or(a, b) => { + boolean_has_pattern(a) || boolean_has_pattern(b) + } + BooleanExpression::Not(b) => boolean_has_pattern(b), + _ => false, + } + } + let mut found = false; + transform_value(&mut expr.clone(), &mut |value| { + found |= match value { + ValueExpression::PatternComprehension { .. } + | ValueExpression::ListComprehension { .. } => true, + ValueExpression::Boolean(b) => boolean_has_pattern(b), + ValueExpression::Case { whens, .. } => { + whens.iter().any(|(b, _)| boolean_has_pattern(b)) + } + _ => false, + }; + }); + found +} + /// [`transform_value`] over every value in a predicate. pub fn transform_boolean_values( expr: &mut BooleanExpression, 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 35e8d00..e18d4fe 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs @@ -14,8 +14,8 @@ use datafusion::logical_expr::{col, LogicalPlan, LogicalPlanBuilder}; /// 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, .. } => { + crate::ast::transform_value(expr, &mut |expr| { + if let ValueExpression::AggregateFunction { name, .. } = expr { 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())); @@ -33,23 +33,7 @@ fn split_aggregates(expr: &mut ValueExpression, aggregates: &mut Vec<(String, Va }; } } - 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 { @@ -96,7 +80,7 @@ impl DataFusionPlanner { let first = aggregates.len(); split_aggregates(&mut rewritten, &mut aggregates); for (column, aggregate) in &aggregates[first..] { - let mut aggregate = self.value_expr(aggregate); + let mut aggregate = self.value_expr(aggregate, &schema)?; // 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 { @@ -105,10 +89,10 @@ impl DataFusionPlanner { } agg_exprs.push(aggregate.alias(column)); } - outer.push(self.value_expr(&rewritten).alias(alias)); + outer.push(self.unresolved_value_expr(&rewritten).alias(alias)); } else { // Group expressions: use raw expression for grouping, no alias - group_exprs.push(self.value_expr(&p.expression)); + group_exprs.push(self.value_expr(&p.expression, &schema)?); } } @@ -130,7 +114,7 @@ impl DataFusionPlanner { })); } else if !super::super::expression::contains_aggregate(&p.expression) { // Re-create the expression and apply alias - let expr = self.value_expr(&p.expression); + let expr = self.unresolved_value_expr(&p.expression); let aliased = if let Some(alias) = &p.alias { expr.alias(alias.to_lowercase()) } else { @@ -145,9 +129,15 @@ impl DataFusionPlanner { } } - LogicalPlanBuilder::from(input_plan) + let aggregate = LogicalPlanBuilder::from(input_plan) .aggregate(group_exprs, agg_exprs) - .map_err(|e| self.plan_error("Failed to build aggregate", e))? + .and_then(|b| b.build()) + .map_err(|e| self.plan_error("Failed to build aggregate", e))?; + let final_projection = final_projection + .into_iter() + .map(|expr| self.resolve_expr(expr, aggregate.schema())) + .collect::>>()?; + LogicalPlanBuilder::from(aggregate) .project(final_projection) .map_err(|e| self.plan_error("Failed to project after aggregate", e))? .build() 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 fa0944b..cd4b6f5 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs @@ -190,8 +190,7 @@ impl DataFusionPlanner { predicate: &crate::ast::BooleanExpression, ) -> Result { let input_plan = self.build_operator(ctx, input)?; - let predicate = super::super::expression::bind_boolean(predicate, input_plan.schema()); - let expr = self.boolean_expr(&predicate); + let expr = self.boolean_expr(predicate, input_plan.schema())?; LogicalPlanBuilder::from(input_plan) .filter(expr) .map_err(|e| self.plan_error("Failed to build filter", e))? @@ -223,7 +222,16 @@ impl DataFusionPlanner { } = input { let input_plan = self.build_operator(ctx, inner)?; - let ordering = self.sort_exprs(&input_plan, sort_items); + let ordering = self + .sort_exprs(&input_plan, sort_items) + .into_iter() + .map(|sort| { + Ok(SortExpr { + expr: self.resolve_expr(sort.expr, input_plan.schema())?, + ..sort + }) + }) + .collect::>>()?; return self.build_project_with_aggregates(input_plan, projections, &ordering); } let input_plan = self.build_operator(ctx, input)?; @@ -278,8 +286,7 @@ impl DataFusionPlanner { exprs.extend(columns); continue; } - let bound = super::super::expression::bind_value(&p.expression, &schema); - let expr = self.value_expr(&bound); + let expr = self.value_expr(&p.expression, &schema)?; // 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 { @@ -323,6 +330,8 @@ impl DataFusionPlanner { .sort(sort_exprs) .map_err(|e| self.plan_error("Failed to build sort", e))? .build() + // The sort builder first restores input columns omitted by RETURN. + .and_then(|plan| plan.resolve_lambda_variables().map(|p| p.data)) .map_err(|e| self.plan_error("Failed to build plan", e)) } @@ -338,7 +347,7 @@ impl DataFusionPlanner { let bound = super::super::expression::bind_value(&item.expression, input_plan.schema()); SortExpr { - expr: self.value_expr(&bound), + expr: self.unresolved_value_expr(&bound), asc: matches!(item.direction, crate::ast::SortDirection::Ascending), nulls_first: true, } @@ -392,8 +401,7 @@ impl DataFusionPlanner { }; // Convert expression to DataFusion Expr - let bound = super::super::expression::bind_value(expression, input_plan.schema()); - let df_expr = self.value_expr(&bound); + let df_expr = self.value_expr(expression, input_plan.schema())?; // A list of node ids: unnest under a temporary name, then join each id back to the // node table so the alias carries the node's columns like any matched node. let unnested = match label { diff --git a/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs index 45f6f4e..51c99bb 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs @@ -10,14 +10,13 @@ use datafusion::common::Column; use datafusion::functions::core::expr_fn::coalesce; use datafusion::functions_aggregate::array_agg::array_agg; use datafusion::functions_aggregate::count::count; -use datafusion::functions_nested::expr_fn::{array_element, cardinality, gen_series, make_array}; +use datafusion::functions_nested::expr_fn::{array_element, array_length, gen_series, make_array}; use datafusion::functions_window::expr_fn::row_number; use datafusion::logical_expr::{ - col, lit, Cast, Expr, ExprFunctionExt, JoinType, LogicalPlan, LogicalPlanBuilder, + col, lit, Cast, Expr, ExprFunctionExt, ExprSchemable, JoinType, LogicalPlan, LogicalPlanBuilder, }; use super::super::analysis::PlanningContext; -use super::super::expression::{bind_boolean, bind_value}; use crate::ast::{BooleanExpression, ValueExpression}; use crate::case_insensitive::qualify_column; use crate::datafusion_planner::DataFusionPlanner; @@ -56,8 +55,17 @@ impl DataFusionPlanner { .map_err(error("Failed to number comprehension rows"))? .build() .map_err(error("Failed to number comprehension rows"))?; - let list = self.value_expr(&bind_value(list, numbered.schema())); - let length = int64(coalesce(vec![cardinality(list.clone()), lit(0_u64)])); + let list = self.value_expr(list, numbered.schema())?; + let list = if list + .get_type(numbered.schema()) + .map_err(error("Failed to type comprehension list"))? + == DataType::Null + { + datafusion::logical_expr::cast(list, DataType::new_list(DataType::Null, true)) + } else { + list + }; + let length = int64(coalesce(vec![array_length(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()) @@ -132,13 +140,13 @@ impl DataFusionPlanner { // 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(self.boolean_expr(&bind_boolean(keep, elements.schema()))); + kept = kept.and(self.boolean_expr(keep, elements.schema())?); } let aggregate = if count_only { // Count rows, not the mapped value: null elements still occupy a slot. count(lit(1_i64)).filter(kept).build() } else { - let value = self.value_expr(&bind_value(map, elements.schema())); + let value = self.value_expr(map, elements.schema())?; array_agg(value) .filter(kept) .order_by(vec![col(&pos).sort(true, false)]) diff --git a/crates/lance-graph/src/datafusion_planner/builder/join_builder.rs b/crates/lance-graph/src/datafusion_planner/builder/join_builder.rs index 97be618..5086a25 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/join_builder.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/join_builder.rs @@ -50,11 +50,13 @@ impl DataFusionPlanner { .project(right_exprs) .and_then(|right| right.build()) .and_then(|right| { - LogicalPlanBuilder::from(left).join( + LogicalPlanBuilder::from(left).join_detailed( right, datafusion::logical_expr::JoinType::Left, (keys.to_vec(), right_keys), None, + // Captured scalar nulls identify the same correlation binding. + datafusion::common::NullEquality::NullEqualsNull, ) }) .and_then(|b| b.build()) diff --git a/crates/lance-graph/src/datafusion_planner/expression.rs b/crates/lance-graph/src/datafusion_planner/expression.rs index 78ab9d0..baddb70 100644 --- a/crates/lance-graph/src/datafusion_planner/expression.rs +++ b/crates/lance-graph/src/datafusion_planner/expression.rs @@ -20,12 +20,34 @@ use datafusion_functions_aggregate::min_max::min; use datafusion_functions_aggregate::sum::sum; impl super::DataFusionPlanner { - pub(crate) fn value_expr(&self, value: &ValueExpression) -> Expr { - self.type_parameters(to_df_value_expr(value)) + pub(crate) fn resolve_expr( + &self, + expr: Expr, + schema: &datafusion::common::DFSchema, + ) -> crate::error::Result { + self.type_parameters(expr) + .resolve_lambda_variables(schema) + .map(|e| e.data) + .map_err(|e| self.plan_error("Failed to resolve expression scope", e)) + } + + pub(crate) fn value_expr( + &self, + expr: &ValueExpression, + schema: &datafusion::common::DFSchema, + ) -> crate::error::Result { + self.resolve_expr(to_df_value_expr(&bind_value(expr, schema)), schema) } - pub(crate) fn boolean_expr(&self, value: &BooleanExpression) -> Expr { - self.type_parameters(to_df_boolean_expr(value)) + pub(crate) fn boolean_expr( + &self, + expr: &BooleanExpression, + schema: &datafusion::common::DFSchema, + ) -> crate::error::Result { + self.resolve_expr(to_df_boolean_expr(&bind_boolean(expr, schema)), schema) + } + pub(crate) fn unresolved_value_expr(&self, value: &ValueExpression) -> Expr { + self.type_parameters(to_df_value_expr(value)) } fn type_parameters(&self, expression: Expr) -> Expr { @@ -159,6 +181,32 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { col(qualify_column(&prop.variable, &prop.property)) } VE::Variable(v) => col(v.to_lowercase()), + VE::LocalVariable(v) => Expr::LambdaVariable( + datafusion::logical_expr::expr::LambdaVariable::new(v.to_lowercase(), None), + ), + VE::Reduce { + accumulator, + initial, + variable, + list, + expression, + } => super::list_folds::reduce( + to_df_value_expr(list), + to_df_value_expr(initial), + vec![accumulator.clone(), variable.clone()], + to_df_value_expr(expression), + ), + VE::ListPredicate { + kind, + variable, + list, + predicate, + } => super::list_folds::predicate( + kind, + to_df_value_expr(list), + variable.clone(), + to_df_value_expr(predicate), + ), VE::Literal(PV::String(s)) => lit(s.clone()), VE::Literal(PV::Integer(i)) => lit(*i), VE::Literal(PV::Float(f)) => lit(*f), @@ -230,7 +278,8 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { "abs" | "round" | "floor" | "ceil" | "sqrt" | "left" | "right" | "substring" | "replace" | "split" | "trim" | "ltrim" | "rtrim" | "size" | "reverse" | "last" | "tail" | "keys" | "__entity_keys" | "labels" | "startnode" - | "endnode" => { + | "endnode" | "__predicate_all" | "__predicate_any" | "__predicate_none" + | "__predicate_single" => { super::scalar_functions::call(name, args.iter().map(to_df_value_expr).collect()) } "range" => { @@ -587,34 +636,13 @@ pub(crate) fn bind_boolean( /// 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)) - || otherwise.as_deref().is_some_and(contains_aggregate) - } - VE::VectorDistance { left, right, .. } => { - contains_aggregate(left) || contains_aggregate(right) - } - VE::VectorSimilarity { left, right, .. } => { - contains_aggregate(left) || contains_aggregate(right) - } - _ => false, - } + let mut found = false; + crate::ast::transform_value(&mut expr.clone(), &mut |value| { + found |= matches!(value, ValueExpression::AggregateFunction { .. }); + }); + found } -/// Convert a ValueExpression to Cypher dot notation for column naming -/// -/// This generates user-friendly column names following Cypher conventions: -/// - Property references: `p.name` (variable.property) -/// - Functions: `function_name(arg)` with simplified argument representation -/// - Other expressions: Use the expression as-is -/// /// This is used when no explicit alias is provided in RETURN clauses. pub(crate) fn to_cypher_column_name(expr: &ValueExpression) -> String { use crate::ast::ValueExpression as VE; diff --git a/crates/lance-graph/src/datafusion_planner/join_ops.rs b/crates/lance-graph/src/datafusion_planner/join_ops.rs index 3fa0d8a..8d5f2dc 100644 --- a/crates/lance-graph/src/datafusion_planner/join_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/join_ops.rs @@ -59,7 +59,8 @@ impl DataFusionPlanner { // Apply target property filters (e.g., (b {age: 30})) for (k, v) in target_properties.iter() { - let lit_expr = self.value_expr(&crate::ast::ValueExpression::Literal(v.clone())); + let lit_expr = + self.unresolved_value_expr(&crate::ast::ValueExpression::Literal(v.clone())); let filter_expr = Expr::BinaryExpr(BinaryExpr { left: Box::new(col(k.to_lowercase())), op: Operator::Eq, diff --git a/crates/lance-graph/src/datafusion_planner/list_folds.rs b/crates/lance-graph/src/datafusion_planner/list_folds.rs new file mode 100644 index 0000000..9f3a8b7 --- /dev/null +++ b/crates/lance-graph/src/datafusion_planner/list_folds.rs @@ -0,0 +1,328 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Cypher ordered folds and three-valued list predicates, using DataFusion lambdas. + +use std::sync::Arc; + +use arrow::array::{ + new_empty_array, new_null_array, Array, ArrayRef, AsArray, BooleanArray, BooleanBuilder, +}; +use arrow::compute::{cast, concat, take_arrays}; +use arrow_schema::{DataType, Field, FieldRef}; +use datafusion::common::utils::{adjust_offsets_for_slice, list_values, list_values_row_number}; +use datafusion::common::{exec_err, plan_err, Result}; +use datafusion::logical_expr::expr::{HigherOrderFunction, Lambda}; +use datafusion::logical_expr::{ + ColumnarValue, Expr, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderSignature, + HigherOrderUDF, HigherOrderUDFImpl, LambdaParametersProgress, ValueOrLambda, Volatility, +}; + +pub(super) fn reduce(list: Expr, initial: Expr, params: Vec, body: Expr) -> Expr { + Expr::HigherOrderFunction(HigherOrderFunction::new( + Arc::new(HigherOrderUDF::from(CypherFold::new("reduce"))), + vec![list, initial, Expr::Lambda(Lambda::new(params, body))], + )) +} + +pub(super) fn predicate(kind: &str, list: Expr, variable: String, body: Expr) -> Expr { + Expr::HigherOrderFunction(HigherOrderFunction::new( + Arc::new(HigherOrderUDF::from(CypherFold::new(kind))), + vec![list, Expr::Lambda(Lambda::new(vec![variable], body))], + )) +} + +#[derive(Debug, PartialEq, Eq, Hash)] +struct CypherFold { + name: String, + signature: HigherOrderSignature, +} + +impl CypherFold { + fn new(name: &str) -> Self { + let mut arguments = vec![ValueOrLambda::Value(())]; + if name == "reduce" { + arguments.push(ValueOrLambda::Value(())); + } + arguments.push(ValueOrLambda::Lambda(())); + Self { + name: name.to_string(), + signature: HigherOrderSignature::exact(arguments, Volatility::Immutable), + } + } +} + +fn list_type(data_type: &DataType) -> Result { + match data_type { + DataType::List(_) | DataType::LargeList(_) => Ok(data_type.clone()), + DataType::ListView(f) | DataType::FixedSizeList(f, _) => Ok(DataType::List(f.clone())), + DataType::LargeListView(f) => Ok(DataType::LargeList(f.clone())), + DataType::Null => Ok(DataType::new_list(DataType::Null, true)), + other => plan_err!("reduce/list predicate requires a list, got {other}"), + } +} + +impl HigherOrderUDFImpl for CypherFold { + fn name(&self) -> &str { + &self.name + } + fn signature(&self) -> &HigherOrderSignature { + &self.signature + } + + fn coerce_value_types(&self, types: &[DataType]) -> Result> { + let mut result = types.to_vec(); + result[0] = list_type(&types[0])?; + Ok(result) + } + + fn lambda_parameters( + &self, + step: usize, + fields: &[ValueOrLambda>], + ) -> Result { + let ValueOrLambda::Value(list) = &fields[0] else { + return plan_err!("expected list value"); + }; + let normalized = list_type(list.data_type())?; + let (DataType::List(element) | DataType::LargeList(element)) = normalized else { + unreachable!() + }; + if self.name != "reduce" { + return Ok(LambdaParametersProgress::Complete(vec![vec![element]])); + } + let ValueOrLambda::Value(initial) = &fields[1] else { + return plan_err!("expected initial value"); + }; + let ValueOrLambda::Lambda(body) = &fields[2] else { + return plan_err!("expected reduce body"); + }; + match body { + None => Ok(LambdaParametersProgress::Partial(vec![Some(vec![ + initial.clone(), + element, + ])])), + Some(body) => { + // One rebinding permits integer-to-float promotion and a null initializer. + if step > 2 { + return plan_err!("reduce accumulator type does not converge"); + } + Ok(LambdaParametersProgress::Complete(vec![vec![ + body.clone(), + element, + ]])) + } + } + } + + fn coerce_values_for_lambdas( + &self, + fields: &[ValueOrLambda], + ) -> Result>> { + if self.name != "reduce" { + return Ok(None); + } + let [ValueOrLambda::Value(list), ValueOrLambda::Value(initial), ValueOrLambda::Lambda(body)] = + fields + else { + return plan_err!("invalid reduce arguments"); + }; + if initial != body + && *initial != DataType::Null + && !(initial.is_numeric() && body.is_numeric()) + { + return plan_err!("reduce initial value ({initial}) and accumulator ({body}) must have compatible types"); + } + Ok(Some(vec![list.clone(), body.clone()])) + } + + fn return_field_from_args(&self, args: HigherOrderReturnFieldArgs) -> Result { + let Some(ValueOrLambda::Lambda(body)) = args.arg_fields.last() else { + return plan_err!("expected lambda body"); + }; + if self.name == "reduce" { + return Ok(Arc::new(body.as_ref().clone().with_nullable(true))); + } + if !matches!(body.data_type(), DataType::Boolean | DataType::Null) { + return plan_err!( + "{} predicate must return boolean, got {}", + self.name, + body.data_type() + ); + } + Ok(Arc::new(Field::new("", DataType::Boolean, true))) + } + + fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result { + let ValueOrLambda::Value(list) = &args.args[0] else { + return exec_err!("expected list value"); + }; + let Some(ValueOrLambda::Lambda(lambda)) = args.args.last() else { + return exec_err!("expected lambda body"); + }; + let list = list.to_array(args.number_rows)?; + let ranges: Vec<(usize, usize)> = match list.data_type() { + DataType::List(_) => list + .as_list::() + .value_offsets() + .windows(2) + .map(|w| (w[0] as usize, w[1] as usize)) + .collect(), + DataType::LargeList(_) => list + .as_list::() + .value_offsets() + .windows(2) + .map(|w| (w[0] as usize, w[1] as usize)) + .collect(), + other => return exec_err!("expected list, got {other}"), + }; + if self.name == "reduce" { + let ValueOrLambda::Value(initial) = &args.args[1] else { + return exec_err!("expected initial value"); + }; + let initial = cast(&initial.to_array(args.number_rows)?, args.return_type())?; + let values = match list.data_type() { + DataType::List(_) => list.as_list::().values(), + DataType::LargeList(_) => list.as_list::().values(), + _ => unreachable!(), + }; + let mut output: Vec = Vec::with_capacity(list.len()); + for (row, (start, end)) in ranges.into_iter().enumerate() { + if list.is_null(row) { + output.push(new_null_array(args.return_type(), 1)); + continue; + } + let mut accumulator = initial.slice(row, 1); + // Each step sees the previous step's result, in list order. + for index in start..end { + let state = || Ok(accumulator.clone()); + let element = || Ok(values.slice(index, 1)); + accumulator = lambda + .evaluate(&[&state, &element], |captures| { + Ok(captures.iter().map(|a| a.slice(row, 1)).collect()) + })? + .into_array(1)?; + if accumulator.data_type() != args.return_type() { + return exec_err!( + "reduce accumulator type changed from {} to {}", + args.return_type(), + accumulator.data_type() + ); + } + } + output.push(accumulator); + } + let result = if output.is_empty() { + new_empty_array(args.return_type()) + } else { + concat(&output.iter().map(|a| a.as_ref()).collect::>())? + }; + return Ok(ColumnarValue::Array(result)); + } + let values = list_values(&list)?; + let predicate = if values.is_empty() { + new_empty_array(&DataType::Boolean) + } else { + let parameter = || Ok(values.clone()); + cast( + &lambda + .evaluate(&[¶meter], |captures| { + Ok(take_arrays( + captures, + &list_values_row_number(&list)?, + None, + )?) + })? + .into_array(values.len())?, + &DataType::Boolean, + )? + }; + let predicate = predicate.as_boolean(); + let ranges: Vec<(usize, usize)> = match list.data_type() { + DataType::List(_) => adjust_offsets_for_slice(list.as_list::()) + .windows(2) + .map(|w| (w[0] as usize, w[1] as usize)) + .collect(), + DataType::LargeList(_) => adjust_offsets_for_slice(list.as_list::()) + .windows(2) + .map(|w| (w[0] as usize, w[1] as usize)) + .collect(), + _ => unreachable!(), + }; + let mut output = BooleanBuilder::with_capacity(list.len()); + for (row, (start, end)) in ranges.into_iter().enumerate() { + output.append_option(if list.is_null(row) { + None + } else { + predicate_result(&self.name, predicate, start, end) + }); + } + Ok(ColumnarValue::Array(Arc::new(output.finish()))) + } +} + +pub(super) fn predicate_result( + kind: &str, + values: &BooleanArray, + start: usize, + end: usize, +) -> Option { + let mut trues = 0; + let mut falses = 0; + let mut nulls = 0; + for i in start..end { + if values.is_null(i) { + nulls += 1; + } else if values.value(i) { + trues += 1; + } else { + falses += 1; + } + } + match kind { + "all" if falses > 0 => Some(false), + "any" if trues > 0 => Some(true), + "none" if trues > 0 => Some(false), + "single" if trues > 1 => Some(false), + _ if nulls > 0 => None, + "all" | "none" => Some(true), + "any" => Some(false), + "single" => Some(trues == 1), + _ => unreachable!(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn list_predicates_use_three_valued_logic() { + for (values, expected) in [ + (vec![], [Some(true), Some(false), Some(true), Some(false)]), + (vec![None], [None, None, None, None]), + ( + vec![Some(true), None], + [None, Some(true), Some(false), None], + ), + (vec![Some(false), None], [Some(false), None, None, None]), + ( + vec![Some(true), Some(true), None], + [None, Some(true), Some(false), Some(false)], + ), + ( + vec![Some(true), Some(false)], + [Some(false), Some(true), Some(false), Some(true)], + ), + ] { + let values = BooleanArray::from(values); + for (kind, want) in ["all", "any", "none", "single"].iter().zip(expected) { + assert_eq!( + predicate_result(kind, &values, 0, values.len()), + want, + "{kind}: {values:?}" + ); + } + } + } +} diff --git a/crates/lance-graph/src/datafusion_planner/mod.rs b/crates/lance-graph/src/datafusion_planner/mod.rs index 8e67b58..c67009d 100644 --- a/crates/lance-graph/src/datafusion_planner/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/mod.rs @@ -20,6 +20,7 @@ mod config_helpers; mod cypher_udf; pub(crate) mod expression; mod join_ops; +mod list_folds; mod scalar_functions; mod scan_ops; mod udf; diff --git a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs index 6494095..5ff9818 100644 --- a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs +++ b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs @@ -50,6 +50,9 @@ impl ScalarUDFImpl for CypherScalar { let valid = match self.name.as_str() { "abs" | "round" | "floor" | "ceil" | "sqrt" => t.is_integer() || t.is_floating(), "range" => t.is_integer(), + "__predicate_all" | "__predicate_any" | "__predicate_none" + | "__predicate_single" => list_field(t) + .is_some_and(|f| matches!(f.data_type(), DataType::Boolean | DataType::Null)), "keys" | "__entity_keys" => matches!(t, DataType::Struct(_)), "labels" => crate::path_value::entity_kind(t) == Some("node"), "startnode" | "endnode" => { @@ -69,6 +72,9 @@ impl ScalarUDFImpl for CypherScalar { DataType::Struct(fields) => fields[2].data_type().clone(), _ => DataType::Null, }, + "__predicate_all" | "__predicate_any" | "__predicate_none" | "__predicate_single" => { + DataType::Boolean + } "size" => DataType::Int64, "abs" if types[0].is_integer() => DataType::Int64, "abs" | "round" | "floor" | "ceil" | "sqrt" => DataType::Float64, @@ -126,6 +132,30 @@ impl ScalarUDFImpl for CypherScalar { } return Ok(ColumnarValue::Array(Arc::new(result.finish()))); } + if let Some(kind) = self.name.strip_prefix("__predicate_") { + let input = args.args[0].to_array(args.number_rows)?; + let mut result = arrow::array::BooleanBuilder::with_capacity(input.len()); + for row in 0..input.len() { + if input.is_null(row) { + result.append_null(); + continue; + } + let values = match input.data_type() { + DataType::List(_) => input.as_list::().value(row), + DataType::LargeList(_) => input.as_list::().value(row), + DataType::FixedSizeList(_, _) => input.as_fixed_size_list().value(row), + _ => unreachable!("validated list predicate input"), + }; + let values = cast(&values, &DataType::Boolean)?; + result.append_option(super::list_folds::predicate_result( + kind, + values.as_boolean(), + 0, + values.len(), + )); + } + return Ok(ColumnarValue::Array(Arc::new(result.finish()))); + } if matches!(self.name.as_str(), "keys" | "__entity_keys") { let input = args.args[0].to_array(args.number_rows)?; let typed_entity = crate::path_value::entity_kind(input.data_type()).is_some(); diff --git a/crates/lance-graph/src/datafusion_planner/scan_ops.rs b/crates/lance-graph/src/datafusion_planner/scan_ops.rs index 33be41a..b0d7c40 100644 --- a/crates/lance-graph/src/datafusion_planner/scan_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/scan_ops.rs @@ -44,8 +44,9 @@ impl DataFusionPlanner { let filter_exprs: Vec = properties .iter() .map(|(k, v)| { - let lit_expr = - self.value_expr(&crate::ast::ValueExpression::Literal(v.clone())); + let lit_expr = self.unresolved_value_expr( + &crate::ast::ValueExpression::Literal(v.clone()), + ); Expr::BinaryExpr(BinaryExpr { left: Box::new(col(k)), op: Operator::Eq, @@ -139,7 +140,8 @@ impl DataFusionPlanner { // Apply relationship property filters (e.g., -[r {since: 2020}]->) for (k, v) in relationship_properties.iter() { - let lit_expr = self.value_expr(&crate::ast::ValueExpression::Literal(v.clone())); + let lit_expr = + self.unresolved_value_expr(&crate::ast::ValueExpression::Literal(v.clone())); let filter_expr = Expr::BinaryExpr(BinaryExpr { left: Box::new(col(k)), op: Operator::Eq, diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index 36f1cde..d77ed26 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -762,6 +762,22 @@ impl<'a> LogicalPlanner<'a> { self.resolve_value(e); } } + ValueExpression::Reduce { + initial, + list, + expression, + .. + } => { + self.resolve_value(initial); + self.resolve_value(list); + self.resolve_value(expression); + } + ValueExpression::ListPredicate { + list, predicate, .. + } => { + self.resolve_value(list); + self.resolve_value(predicate); + } 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 @@ -1080,6 +1096,80 @@ impl<'a> LogicalPlanner<'a> { input = self.rewrite_exists_in_value(input, left)?; self.rewrite_exists_in_value(input, right) } + ValueExpression::Reduce { + initial, + list, + expression, + .. + } => { + input = self.rewrite_exists_in_value(input, initial)?; + input = self.rewrite_exists_in_value(input, list)?; + self.rewrite_exists_in_value(input, expression) + } + ValueExpression::ListPredicate { + kind, + variable, + list, + predicate, + } if crate::ast::needs_relational_evaluation(predicate) => { + input = self.rewrite_exists_in_value(input, list)?; + let original_list = (**list).clone(); + let mut map = (**predicate).clone(); + // The relational comprehension supplies this iterator as a column. Restore + // inner lambda bindings after replacing references to the outer iterator. + crate::ast::transform_value(&mut map, &mut |value| match value { + ValueExpression::LocalVariable(name) if name == variable => { + *value = ValueExpression::Variable(name.clone()); + } + ValueExpression::Field { base, field } if matches!(base.as_ref(), ValueExpression::Variable(name) if name == variable) => + { + *value = ValueExpression::Property(PropertyRef::new( + variable.clone(), + field.clone(), + )); + } + ValueExpression::Reduce { + accumulator, + variable, + expression, + .. + } => { + crate::ast::mark_locals(expression, &[accumulator, variable]); + } + ValueExpression::ListPredicate { + variable, + predicate, + .. + } => { + crate::ast::mark_locals(predicate, &[variable]); + } + _ => {} + }); + let mut comprehension = ValueExpression::ListComprehension { + variable: variable.clone(), + list: list.clone(), + filter: None, + map: Some(Box::new(map)), + }; + input = self.rewrite_exists_in_value(input, &mut comprehension)?; + *expr = ValueExpression::Case { + whens: vec![( + BooleanExpression::IsNull(original_list), + ValueExpression::Literal(PropertyValue::Null), + )], + otherwise: Some(Box::new(ValueExpression::ScalarFunction { + name: format!("__predicate_{kind}"), + args: vec![comprehension], + })), + }; + Ok(input) + } + ValueExpression::ListPredicate { + list, predicate, .. + } => { + input = self.rewrite_exists_in_value(input, list)?; + self.rewrite_exists_in_value(input, predicate) + } ValueExpression::Boolean(b) => self.rewrite_exists(input, b), ValueExpression::Map(entries) => { for (_, v) in entries.iter_mut() { @@ -1336,9 +1426,6 @@ impl<'a> LogicalPlanner<'a> { } } } - if shared.is_empty() { - return Err(unsupported("without a variable bound by an earlier clause")); - } let mut keys = Vec::with_capacity(shared.len()); for var in &shared { let label = &self.variables[var]; @@ -1355,12 +1442,29 @@ impl<'a> LogicalPlanner<'a> { 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)) + .filter(|n| !local.contains(*n) && !shared.iter().any(|v| v == *n)) .map(str::to_string) .collect(); scalars.sort(); + for variable in &scalars { + // Binding variables distinguish iterator contexts. Untracked aggregate + // captures can be re-evaluated in a different order by the seed plan. + if !self.variables.contains_key(variable) && !shared.is_empty() { + continue; + } + keys.push( + self.variables + .get(variable) + .and_then(|label| self.config.get_node_mapping(label)) + .map(|mapping| qualify_column(variable, &mapping.id_field)) + .unwrap_or_else(|| variable.to_lowercase()), + ); + } carried.extend(scalars); } + if carried.is_empty() { + return Err(unsupported("without a variable bound by an earlier clause")); + } let seed = LogicalOperator::Distinct { input: Box::new(LogicalOperator::Project { input: Box::new(seed_source.unwrap_or_else(|| base.clone())), diff --git a/crates/lance-graph/src/parameter_substitution.rs b/crates/lance-graph/src/parameter_substitution.rs index 3808a58..e4676d6 100644 --- a/crates/lance-graph/src/parameter_substitution.rs +++ b/crates/lance-graph/src/parameter_substitution.rs @@ -303,6 +303,22 @@ fn substitute_in_value_expression( } substitute_in_value_expression(map, parameters)?; } + ValueExpression::Reduce { + initial, + list, + expression, + .. + } => { + substitute_in_value_expression(initial, parameters)?; + substitute_in_value_expression(list, parameters)?; + substitute_in_value_expression(expression, parameters)?; + } + ValueExpression::ListPredicate { + list, predicate, .. + } => { + substitute_in_value_expression(list, parameters)?; + substitute_in_value_expression(predicate, parameters)?; + } ValueExpression::ListComprehension { list, filter, map, .. } => { diff --git a/crates/lance-graph/src/parser.rs b/crates/lance-graph/src/parser.rs index bdeacff..716b351 100644 --- a/crates/lance-graph/src/parser.rs +++ b/crates/lance-graph/src/parser.rs @@ -542,6 +542,11 @@ fn boolean_primary_expression(input: &str) -> IResult<&str, BooleanExpression> { where_clause: None, }), comparison_expression, + map(list_predicate, |left| BooleanExpression::Comparison { + left, + operator: ComparisonOperator::Equal, + right: ValueExpression::Literal(PropertyValue::Boolean(true)), + }), ))(input) } @@ -719,6 +724,8 @@ fn basic_value_expression(input: &str) -> IResult<&str, ValueExpression> { list_comprehension, list_literal, // Then a list of any expressions [a.x, 'y'] (or []) map_literal, + reduce_expression, + list_predicate, parse_parameter, // Try $parameter function_call, // Regular function calls map(property_value, ValueExpression::Literal), // Try literals BEFORE property references @@ -783,6 +790,62 @@ fn list_comprehension(input: &str) -> IResult<&str, ValueExpression> { )) } +fn reduce_expression(input: &str) -> IResult<&str, ValueExpression> { + let (input, _) = tuple((tag_no_case("reduce"), multispace0, char('('), multispace0))(input)?; + let (input, accumulator) = identifier(input)?; + let (input, initial) = + preceded(tuple((multispace0, char('='), multispace0)), expression)(input)?; + let (input, variable) = preceded(comma_ws, identifier)(input)?; + let (input, list) = preceded( + tuple((multispace1, tag_no_case("IN"), multispace1)), + value_expression, + )(input)?; + let (input, mut body) = + preceded(tuple((multispace0, char('|'), multispace0)), expression)(input)?; + let (input, _) = tuple((multispace0, char(')')))(input)?; + mark_locals(&mut body, &[accumulator, variable]); + Ok(( + input, + ValueExpression::Reduce { + accumulator: accumulator.to_lowercase(), + initial: Box::new(initial), + variable: variable.to_lowercase(), + list: Box::new(list), + expression: Box::new(body), + }, + )) +} + +fn list_predicate(input: &str) -> IResult<&str, ValueExpression> { + let (input, kind) = alt(( + tag_no_case("all"), + tag_no_case("any"), + tag_no_case("none"), + tag_no_case("single"), + ))(input)?; + let (input, variable) = + preceded(tuple((multispace0, char('('), multispace0)), identifier)(input)?; + let (input, list) = preceded( + tuple((multispace1, tag_no_case("IN"), multispace1)), + value_expression, + )(input)?; + let (input, mut predicate) = preceded( + tuple((multispace1, tag_no_case("WHERE"), multispace1)), + expression, + )(input)?; + let (input, _) = tuple((multispace0, char(')')))(input)?; + mark_locals(&mut predicate, &[variable]); + Ok(( + input, + ValueExpression::ListPredicate { + kind: kind.to_lowercase(), + variable: variable.to_lowercase(), + list: Box::new(list), + predicate: Box::new(predicate), + }, + )) +} + // `{key: value, ...}` fn map_literal(input: &str) -> IResult<&str, ValueExpression> { let entry = tuple(( @@ -2591,4 +2654,31 @@ mod tests { ValueExpression::Variable("nullCount".into()) ); } + #[test] + fn reduce_keeps_inputs_in_outer_scope_and_marks_local_properties() { + let q = + parse_cypher_query("WITH {n:1} AS x RETURN reduce(s=x.n, x IN [x] | s+x.n)").unwrap(); + let ValueExpression::Reduce { + initial, + list, + expression, + .. + } = &q.return_clause.items[0].expression + else { + panic!("expected reduce"); + }; + assert!(matches!(initial.as_ref(), ValueExpression::Property(p) if p.variable == "x")); + let mut input_locals = 0; + transform_value(&mut list.as_ref().clone(), &mut |v| { + input_locals += usize::from(matches!(v, ValueExpression::LocalVariable(_))) + }); + assert_eq!(input_locals, 0); + let mut body_locals = Vec::new(); + transform_value(&mut expression.as_ref().clone(), &mut |v| { + if let ValueExpression::LocalVariable(n) = v { + body_locals.push(n.clone()); + } + }); + assert_eq!(body_locals, ["s", "x"]); + } } diff --git a/crates/lance-graph/src/pattern_order.rs b/crates/lance-graph/src/pattern_order.rs index bbf4dce..818b9dd 100644 --- a/crates/lance-graph/src/pattern_order.rs +++ b/crates/lance-graph/src/pattern_order.rs @@ -259,7 +259,24 @@ fn value_variables<'a>(expr: &'a ValueExpression, out: &mut HashSet<&'a str>) { value_variables(m, out); } } - ValueExpression::Literal(_) + ValueExpression::Reduce { + initial, + list, + expression, + .. + } => { + value_variables(initial, out); + value_variables(list, out); + value_variables(expression, out); + } + ValueExpression::ListPredicate { + list, predicate, .. + } => { + value_variables(list, out); + value_variables(predicate, out); + } + ValueExpression::LocalVariable(_) + | ValueExpression::Literal(_) | ValueExpression::Parameter(_) | ValueExpression::VectorLiteral(_) => {} } diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index aff6fa4..ea7a9eb 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -578,7 +578,7 @@ impl SemanticAnalyzer { ValueExpression::Literal(_) => { // Literals are always valid } - ValueExpression::Variable(var) => { + ValueExpression::Variable(var) | ValueExpression::LocalVariable(var) => { // Use case-insensitive lookup if !self.variables.contains_key_ci(var) { return Err(GraphError::PlanError { @@ -902,6 +902,32 @@ impl SemanticAnalyzer { self.expression_types = before_types; result?; } + ValueExpression::Reduce { + accumulator, + initial, + variable, + list, + expression, + } => { + self.analyze_value_expression(initial)?; + self.analyze_value_expression(list)?; + if accumulator.eq_ignore_ascii_case(variable) { + return Err(GraphError::PlanError { + message: "reduce accumulator and element variables must differ".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + self.analyze_lambda(&[accumulator, variable], expression, false)?; + } + ValueExpression::ListPredicate { + variable, + list, + predicate, + .. + } => { + self.analyze_value_expression(list)?; + self.analyze_lambda(&[variable], predicate, true)?; + } ValueExpression::ListComprehension { variable, list, @@ -933,6 +959,34 @@ impl SemanticAnalyzer { Ok(()) } + fn analyze_lambda( + &mut self, + names: &[&String], + expression: &ValueExpression, + allow_patterns: bool, + ) -> Result<()> { + let before = self.variables.clone(); + for name in names { + self.variables.remove(&name.to_lowercase()); + self.register_projection_alias(name); + } + let mut unsupported = + !allow_patterns && crate::ast::needs_relational_evaluation(expression); + crate::ast::transform_value(&mut expression.clone(), &mut |value| { + unsupported |= matches!(value, ValueExpression::AggregateFunction { .. }); + }); + let result = if unsupported { + Err(GraphError::UnsupportedFeature { + feature: "graph expressions/comprehensions inside a reduce body, or aggregates inside a lambda body; compute the list before the body".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }) + } else { + self.analyze_value_expression(expression) + }; + self.variables = before; + result + } + fn register_projection_alias(&mut self, alias: &str) { // Use case-insensitive lookup and store normalized key if self.variables.contains_key_ci(alias) { diff --git a/docs/cypher.md b/docs/cypher.md index 77dfee4..e363fca 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -181,7 +181,8 @@ implemented; `ExecutionStrategy.LanceNative` is reserved and raises an error. The Graph Data Council (GDC), formerly known as the Linked Data Benchmark Council (LDBC), maintains the [GDC SNB Interactive](https://ldbcouncil.org/benchmarks/snb/) complex queries. lance-csr implements a large -read-only subset of openCypher, enough to run all 14 (IC14 without `reduce()`). +read-only subset of openCypher. The tables below specify its supported surface; relationship-path +operations needed by the unchanged IC14 query are not yet supported. ### Clauses @@ -371,7 +372,8 @@ comprehensions, retain the usual CSR execution route. | Strings | `toLower` / `lower`, `toUpper` / `upper`, `toString`, `left`, `right`, `substring`, `replace`, `split`, `trim`, `lTrim`, `rTrim`, `reverse`, `size` | | Conversion | `toInteger`, `toFloat`, `coalesce` | | Properties | `keys(map)`, `keys(node)`, `keys(relationship)`, `labels(node)` | -| Lists | `size`, `head`, `last`, `tail`, `reverse`, `range(start, end [, step])` (inclusive) | +| Lists | `size`, `head`, `last`, `tail`, `reverse`, `range(start, end [, step])` (inclusive), `reduce` | +| List predicates | `all`, `any`, `none`, `single` | | Math | `abs`, `round`, `floor`, `ceil`, `sqrt` | | Temporal | `date`, `datetime`, `date_part('month', d)`, `to_epoch_ms`, `duration.inSeconds(a, b).seconds` | | Paths | `length(p)`, `nodes(p)`, `relationships(p)`; `startNode(r)`, `endNode(r)` on relationship values | @@ -430,8 +432,7 @@ print(result.to_pylist()) # [{'labels': ['Person'], 'properties': 3}] ``` Any other function raises `ValueError` with `Cypher function 'name' is not implemented`. This is a -bounded function set, not a claim of full Cypher conformance. Remaining functions include `reduce`, -the list predicates `any` / `all` / `none` / `single`. +bounded function set, not a claim of full Cypher conformance. When `size()` consumes an unmapped list comprehension, such as `size([x IN values WHERE x > 0])`, the planner counts retained elements directly. Each duplicate and each retained null element still contributes one; @@ -440,6 +441,38 @@ an empty or null input list produces zero, as for the existing comprehension. A unnests and groups rows; `explain()` shows a filtered count in place of ordered list aggregation. A direct count output is marked non-nullable in its Arrow schema; the previous list-size expression was marked nullable. +### Ordered reduce and list predicates + +`reduce(total = initial, item IN list | expression)` evaluates the expression once per element in +list order. Each step sees the previous accumulator value. An empty list returns the initial value; +a null list returns null. Null elements are passed to the expression, so use `coalesce` if you want +to replace them. The accumulator and iterator names are local to the body: the initial value and +list still refer to outer variables. The accumulator must keep a compatible Arrow type; integer +initial values can promote to floating point when the expression requires it. + +```python +result = engine.execute(""" + RETURN reduce(total = 0, item IN [1, 2, 3] | total * 10 + item) AS ordered, + all(item IN [1, 2, 3] WHERE item > 0) AS positive, + single(item IN [1, 2, 3] WHERE item = 2) AS exactly_one +""") +print(result.to_pylist()) # [{'ordered': 123, 'positive': True, 'exactly_one': True}] +``` + +The list predicates use `function(item IN list WHERE predicate)` syntax. `all` requires every +predicate result to be true, `any` requires at least one, `none` requires no true results, and +`single` requires exactly one. For an empty list, `all` and `none` are true; `any` and `single` are +false. A null list returns null. A null predicate result makes the answer null unless a known result +already decides it: false decides `all`, true decides `any` and `none`, and two true results decide +`single`. Predicates can inspect null elements explicitly with `item IS NULL`. + +Both forms can capture outer values and nest scalar folds or predicates. List predicates also +support correlated graph predicates such as `EXISTS { MATCH ... }`. A reduce input can be a list +comprehension, pattern comprehension, or `collect(...)`; its body supports scalar expressions but +rejects graph patterns, comprehensions, and aggregation. Compute such a list before the body. +These semantics follow the pinned [Cypher 4.4 list reference](https://neo4j.com/docs/cypher-manual/4.4/functions/list/#functions-reduce) +and [predicate reference](https://neo4j.com/docs/cypher-manual/4.4/functions/predicate/). + ### Example: a little of everything ```python diff --git a/python/python/tests/test_list_folds.py b/python/python/tests/test_list_folds.py new file mode 100644 index 0000000..906b196 --- /dev/null +++ b/python/python/tests/test_list_folds.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Ordered reduce, lexical scope, and Cypher's three-valued list predicates.""" + +import itertools + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture(params=[False, True], ids=["join", "csr"]) +def engine(request): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + return CypherEngine( + config, + { + "Person": pa.table( + { + "id": [1, 2, 3], + "items": pa.array([[1, None, 3], None, []], pa.list_(pa.int64())), + "large": pa.array([[1, 2], None, []], pa.large_list(pa.int64())), + "pair": pa.array([[1, 2], None, [3, 4]], pa.list_(pa.int64(), 2)), + } + ), + "KNOWS": pa.table({"src": [1, 1, 2], "dst": [2, 3, 3]}), + }, + use_csr=request.param, + ) + + +@pytest.mark.parametrize( + ("expression", "expected"), + [ + ("reduce(s=0, x IN [1,2,3] | s*10+x)", 123), + ("reduce(s=9, x IN [] | s+x)", 9), + ("reduce(s=9, x IN null | s+x)", None), + ("reduce(s=0, x IN [1,null,3] | s+x)", None), + ("reduce(s=0, x IN [1,null,3] | s+coalesce(x,0))", 4), + ("reduce(s=null, x IN [1,2,3] | coalesce(s,0)+x)", 6), + ("reduce(s=0, x IN [1.5,2.5] | s+x)", 4.0), + ("reduce(s=1, x IN [2,3] | s/x)", 0), + ("reduce(s='猫😀猫', x IN ['猫','😀'] | replace(s,x,'a'))", "aaa"), + ("reduce(s={n:0}, x IN [1,2] | {n:s.n+x}).n", 3), + ("reduce(s=0, x IN [{n:1},{n:2}] | s+x.n)", 3), + ("reduce(s=0, xs IN [[1,2],[3]] | s+reduce(t=0, x IN xs | t+x))", 6), + ("[xs IN [[1,2],[3]] | reduce(s=0,x IN xs | s+x)]", [3, 3]), + ("reduce(s=0, x IN [1,2] | s+reduce(s=10, y IN [x] | s+y))", 23), + ("reduce(s=0, x IN [1,2] | s+reduce(t=s, y IN [x] | t+y))", 4), + ( + "reduce(s=0, x IN [1,2] | s+CASE WHEN any(y IN [x] WHERE y=2) " + "THEN 10 ELSE x END)", + 11, + ), + ("any(x IN [null] WHERE x IS NULL)", True), + ("all(x IN [null] WHERE x IS NULL)", True), + ("single(x IN [null] WHERE x IS NULL)", True), + ("none(x IN [null] WHERE x IS NULL)", False), + ("any(x IN [[1,2],[3]] WHERE size(x)=2)", True), + ("all(x IN [{n:1},{n:2}] WHERE x.n>0)", True), + ("all(x IN [1,2] WHERE any(x IN [3,4] WHERE x>2))", True), + ], +) +def test_fold_expression(engine, expression, expected): + assert engine.execute(f"RETURN {expression} AS result").to_pylist() == [ + {"result": expected} + ] + + +def test_predicate_truth_tables(engine): + """Exhaust true/false/null combinations through length three.""" + lists = [None] + [ + list(values) + for length in range(4) + for values in itertools.product([True, False, None], repeat=length) + ] + expected = [] + for values in lists: + if values is None: + expected.append(dict.fromkeys(["a", "b", "c", "d"])) + continue + yes = values.count(True) + no = values.count(False) + unknown = None in values + expected.append( + { + "a": False if no else None if unknown else True, + "b": True if yes else None if unknown else False, + "c": False if yes else None if unknown else True, + "d": False if yes > 1 else None if unknown else yes == 1, + } + ) + for values, want in zip(lists, expected, strict=True): + result = engine.execute( + "RETURN all(x IN $xs WHERE x) AS a, any(x IN $xs WHERE x) AS b, " + "none(x IN $xs WHERE x) AS c, single(x IN $xs WHERE x) AS d", + {"xs": values}, + ).to_pylist() + assert result == [want], values + + +def test_captured_properties_scope_and_arrow_lists(engine): + result = engine.execute( + "MATCH (p:Person) RETURN p.id AS id, " + "reduce(p=p.id, x IN p.items | p+coalesce(x,0)) AS reduced, " + "all(x IN p.items WHERE x(q:Person)|q.id] | s+x) AS result" + ).to_pylist() == [{"result": 5}] + assert engine.execute( + "MATCH (p:Person {id:1}) RETURN " + "[x IN [2,3] | reduce(w=0.0, v IN [(p)-[:KNOWS]->(q:Person) " + "WHERE q.id=x | 0.5] | w+v)] AS result" + ).to_pylist() == [{"result": [0.5, 0.5]}] + + +def test_graph_predicates_reuse_correlated_comprehensions(engine): + result = engine.execute( + "MATCH (p:Person) RETURN p.id AS id, " + "any(x IN [2,3] WHERE EXISTS {MATCH (p)-[:KNOWS]->(q:Person) " + "WHERE q.id=x}) AS connected " + "ORDER BY id" + ).to_pylist() + assert result == [ + {"id": 1, "connected": True}, + {"id": 2, "connected": True}, + {"id": 3, "connected": False}, + ] + for value, expected in [(None, None), ([], True)]: + assert engine.execute( + "MATCH (p:Person {id:1}) RETURN " + "all(x IN $xs WHERE EXISTS {MATCH (p)-[:KNOWS]->(q:Person) " + "WHERE q.id=x}) AS result", + {"xs": value}, + ).to_pylist() == [{"result": expected}] + + +def test_graph_predicate_null_and_duplicate_capture_bindings(engine): + result = engine.execute( + "MATCH (p:Person {id:1}) RETURN " + "[x IN [2,2,null,99] | any(y IN [x] WHERE EXISTS {" + "MATCH (p)-[:KNOWS]->(q:Person) WHERE q.id=y OR y IS NULL})] AS result" + ).to_pylist() + assert result == [{"result": [True, True, True, False]}] + # No shared graph variable: correlation consists only of the scalar iterator. + assert engine.execute( + "RETURN all(x IN [1,2,3] WHERE EXISTS {MATCH (p:Person) " + "WHERE p.id=x}) AS result" + ).to_pylist() == [{"result": True}] + + +def test_pattern_reduce_correlates_captured_struct_and_duplicate_rows(engine): + assert engine.execute( + "WITH [{id:2},{id:2},{id:3}] AS values UNWIND values AS captured " + "RETURN captured.id AS id, reduce(s=0, x IN " + "[(a:Person)-[:KNOWS]->(b:Person) WHERE b.id=captured.id | 1] | s+x) AS n " + "ORDER BY id" + ).to_pylist() == [{"id": 2, "n": 1}, {"id": 2, "n": 1}, {"id": 3, "n": 2}] + + +@pytest.mark.parametrize( + "query", + [ + "RETURN reduce(s=0, s IN [1] | s)", + "RETURN reduce(s=0, x IN 1 | s+x)", + "RETURN any(x IN 1 WHERE true)", + "RETURN all(x IN [1] WHERE x)", + "RETURN reduce(s=0, x IN [1] | s+undefined)", + "RETURN reduce(s=0, x IN [1] | s+x) AS result, x", + "RETURN any(x IN [1] WHERE x=1) AS result, x", + "RETURN reduce(s=0, x IN [1] | toString(x))", + "RETURN reduce(s=0, x IN [1] | size([y IN [x] | y]))", + "MATCH (p:Person) RETURN reduce(s=0,x IN [1] | " + "CASE WHEN EXISTS {MATCH (p)-[:KNOWS]->(:Person)} THEN s+x ELSE s END)", + ], +) +def test_invalid_fold_scope_and_types(engine, query): + with pytest.raises((ValueError, RuntimeError, NotImplementedError)): + engine.execute(query) + + +@pytest.mark.parametrize("use_csr", [False, True], ids=["join", "csr"]) +def test_sliced_lists_and_unreachable_null_storage(use_csr): + config = GraphConfig.builder().node_table("Person", "id").build() + items = pa.ListArray.from_arrays( + pa.array([0, 2, 4], type=pa.int32()), + pa.array([0, 0, 1, 2]), + mask=pa.array([True, False]), + ) + people = pa.table({"id": [1, 2], "items": items}) + query = ( + "MATCH (p:Person) RETURN p.id AS id, " + "reduce(s=12,x IN p.items | s/x) AS value, " + "all(x IN p.items WHERE 12/x>0) AS predicate ORDER BY id" + ) + expected = [ + {"id": 1, "value": None, "predicate": None}, + {"id": 2, "value": 6, "predicate": True}, + ] + for table, want in [(people, expected), (people.slice(1), expected[1:])]: + local = CypherEngine(config, {"Person": table}, use_csr=use_csr) + assert local.execute(query).to_pylist() == want + + +def test_fold_in_sort_and_unwind(engine): + assert engine.execute( + "MATCH (p:Person) RETURN p.id AS id ORDER BY reduce(s=0,x IN [p.id] | s-x)" + ).to_pylist() == [{"id": 3}, {"id": 2}, {"id": 1}] + assert engine.execute( + "UNWIND [reduce(s=0,x IN [1,2] | s+x)] AS value RETURN value" + ).to_pylist() == [{"value": 3}] From ae1218f0cbb8f9d3f5d822506e37475fa89d7a4b Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 09:56:11 -0700 Subject: [PATCH 27/66] feat(python): ship public API type stubs Add PEP 561 stubs matching the extension graph module, strict consumer checks, and runtime signature parity tests. Document the typing workflow and check wheel inclusion. Runtime and benchmark validation remain pending the coordinated build. --- ci/check_wheel.py | 3 + docs/api-reference.md | 29 ++++ python/DEVELOPMENT.md | 24 ++- python/pyproject.toml | 13 +- python/python/lance_csr/__init__.py | 2 +- .../python/lance_csr/_internal/__init__.pyi | 3 + python/python/lance_csr/_internal/graph.pyi | 155 ++++++++++++++++++ python/python/lance_csr/py.typed | 0 python/python/tests/test_typing.py | 149 +++++++++++++++++ python/python/tests/typing/public_api.py | 95 +++++++++++ python/uv.lock | 14 ++ 11 files changed, 482 insertions(+), 5 deletions(-) create mode 100644 python/python/lance_csr/_internal/__init__.pyi create mode 100644 python/python/lance_csr/_internal/graph.pyi create mode 100644 python/python/lance_csr/py.typed create mode 100644 python/python/tests/test_typing.py create mode 100644 python/python/tests/typing/public_api.py diff --git a/ci/check_wheel.py b/ci/check_wheel.py index 3d3e31c..4a5999e 100644 --- a/ci/check_wheel.py +++ b/ci/check_wheel.py @@ -36,6 +36,9 @@ def main(): # outside sys.prefix, so check for site-packages itself). extension = Path(_internal.__file__).resolve() assert "site-packages" in extension.parts, extension + package = extension.parent + for relative in ("py.typed", "_internal/__init__.pyi", "_internal/graph.pyi"): + assert (package / relative).is_file(), f"Missing typing file: {relative}" config = ( GraphConfig.builder() diff --git a/docs/api-reference.md b/docs/api-reference.md index 737354d..0162102 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -16,6 +16,35 @@ code. The guides link here for exact signatures and link back for examples. `lance_csr.__version__` is the package version. +## Type checking + +The package ships type stubs for every public class, including the compiled extension, and a `py.typed` +marker. Editors and type checkers can follow builder methods, parameters, query results and statistics. +The stubs describe the current API; they do not change runtime validation or exceptions. + +PyArrow distributes its type information separately. For strict checks of Arrow inputs and results, add +`pyarrow-stubs` to your development dependencies: + +```bash +uv add --dev pyright pyarrow-stubs +uv run pyright your_script.py +``` + +```python +from typing import assert_type + +import pyarrow as pa +from lance_csr import CypherEngine, GraphConfig + +config = GraphConfig.builder().node_table("Person", "id").build() +engine = CypherEngine(config, {"Person": pa.table({"id": [1, 2]})}) +result = engine.execute("MATCH (p:Person) RETURN p.id ORDER BY p.id") +assert_type(result, pa.Table) +print(result.to_pylist()) # [{'p.id': 1}, {'p.id': 2}] +``` + +This reference remains maintained by hand alongside the stubs and Rust bindings. + ## Datasets Methods that take `datasets` take a `dict[str, object]` mapping each node label or relationship type (or SQL diff --git a/python/DEVELOPMENT.md b/python/DEVELOPMENT.md index 885a34c..87d7401 100644 --- a/python/DEVELOPMENT.md +++ b/python/DEVELOPMENT.md @@ -20,8 +20,8 @@ uv sync --locked --no-install-project # locked dependencies and the dev group uv run --no-sync maturin develop # debug build of the extension, installed in editable mode ``` -`make install` runs both. The `dev` dependency group in `pyproject.toml` holds maturin, pytest, pandas, ruff and -pyright, and `uv sync` installs it by default. +`make install` runs both. The `dev` dependency group in `pyproject.toml` holds maturin, pytest, pandas, ruff, +pyright and PyArrow's separate stubs, and `uv sync` installs it by default. - `--no-install-project` leaves the package to `maturin develop`. Without it, uv builds the package through maturin's build backend, in release mode. @@ -59,6 +59,26 @@ make format # ruff format make lint # ruff check, ruff format --check, pyright ``` +The public Python API is typed in `python/lance_csr/_internal/graph.pyi`, matching the compiled +`_internal.graph` module. `_internal/__init__.pyi` describes its parent module, and `py.typed` marks the +package as typed. Update the stubs, the hand-written API reference and the relevant guide whenever a +binding changes. Exceptions remain the built-in Python exceptions listed in the API reference. + +`make lint` checks `python/tests/typing/public_api.py` in strict mode: it exercises every class, checks +return types, and confirms invalid calls remain rejected. PyArrow's separate stubs are a development +dependency because Arrow tables and readers are part of this API. They are not a runtime dependency. + +After building the extension, check public members, enum variants, properties, constructors, argument +names and defaults against the stubs: + +```shell +uv run --no-sync pytest python/tests/test_typing.py -q +``` + +These tests fail when an exported class or public method has no stub. They also check that each top-level +export refers to the class declared by the extension. The wheel smoke check verifies that the stubs and +marker survive packaging; run it against the installed wheel as described below. + For Rust, from the repository root (CI runs the same checks): ```shell diff --git a/python/pyproject.toml b/python/pyproject.toml index bb2c8fa..10b9e3a 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -40,7 +40,14 @@ build-backend = "maturin" # Development tools, installed by `uv sync` (the `dev` group is uv's default). [dependency-groups] -dev = ["maturin[patchelf]>=1.9,<2", "pytest", "pandas", "ruff", "pyright"] +dev = [ + "maturin[patchelf]>=1.9,<2", + "pytest", + "pandas", + "ruff", + "pyright", + "pyarrow-stubs>=20.0.0.20260819", +] [tool.ruff] lint.select = ["F", "E", "W", "I", "G", "TCH", "PERF", "B019"] @@ -50,7 +57,9 @@ lint.select = ["F", "E", "W", "I", "G", "TCH", "PERF", "B019"] [tool.pyright] pythonVersion = "3.13" -include = ["python/lance_csr/__init__.py"] +include = ["python/lance_csr", "python/tests/typing"] +strict = ["python/tests/typing"] +reportUnnecessaryTypeIgnoreComment = "error" # Dependencies like pyarrow make this difficult to enforce strictly. reportMissingTypeStubs = "warning" reportImportCycles = "error" diff --git a/python/python/lance_csr/__init__.py b/python/python/lance_csr/__init__.py index 196aa72..8a29ec4 100644 --- a/python/python/lance_csr/__init__.py +++ b/python/python/lance_csr/__init__.py @@ -11,7 +11,7 @@ from __future__ import annotations try: - from . import _internal as _bindings # type: ignore[attr-defined] + from . import _internal as _bindings except ImportError as e: raise ImportError( "Unable to import the compiled lance-csr extension. " diff --git a/python/python/lance_csr/_internal/__init__.pyi b/python/python/lance_csr/_internal/__init__.pyi new file mode 100644 index 0000000..a2ac205 --- /dev/null +++ b/python/python/lance_csr/_internal/__init__.pyi @@ -0,0 +1,3 @@ +from . import graph as graph + +__version__: str diff --git a/python/python/lance_csr/_internal/graph.pyi b/python/python/lance_csr/_internal/graph.pyi new file mode 100644 index 0000000..1c6b13a --- /dev/null +++ b/python/python/lance_csr/_internal/graph.pyi @@ -0,0 +1,155 @@ +from collections.abc import Sequence +from datetime import date, datetime +from typing import ClassVar, Literal, Protocol, Self, TypedDict, TypeVar, final + +from lance import LanceDataset +from pyarrow import RecordBatch, RecordBatchReader, Table + +class _ArrowStream(Protocol): + def __arrow_c_stream__(self, requested_schema: object = None) -> object: ... + +_Dataset = Table | RecordBatch | RecordBatchReader | LanceDataset | _ArrowStream +_DatasetT = TypeVar("_DatasetT", bound=_Dataset) +_Parameter = int | float | str | bool | date | datetime | None +_ParameterT = TypeVar("_ParameterT", bound=_Parameter) +_Scalar = int | float | str | bool | None + +class _NodeIndexStats(TypedDict): + rows: int + version: int | None + build_ms: float + +class _CsrStats(TypedDict): + edges: int + delta_edges: int + tombstones: int + update: Literal["build", "rebuild", "delta", "reused"] + version: int | None + build_ms: float + +@final +class ExecutionStrategy: + DataFusion: ClassVar[ExecutionStrategy] + LanceNative: ClassVar[ExecutionStrategy] + def __int__(self) -> int: ... + +@final +class SqlDialect: + Default: ClassVar[SqlDialect] + Spark: ClassVar[SqlDialect] + PostgreSql: ClassVar[SqlDialect] + MySql: ClassVar[SqlDialect] + Sqlite: ClassVar[SqlDialect] + def __int__(self) -> int: ... + +@final +class DistanceMetric: + L2: ClassVar[DistanceMetric] + Cosine: ClassVar[DistanceMetric] + Dot: ClassVar[DistanceMetric] + def __int__(self) -> int: ... + +@final +class GraphConfig: + @staticmethod + def builder() -> GraphConfigBuilder: ... + def node_labels(self) -> list[str]: ... + def relationship_types(self) -> list[str]: ... + +@final +class GraphConfigBuilder: + def __init__(self) -> None: ... + def node_table(self, label: str, primary_key: str) -> Self: ... + def rel_table( + self, rel_type: str, from_table: str, to_table: str, src: str, dst: str + ) -> Self: ... + def with_node_label(self, label: str, id_field: str) -> Self: ... + def with_relationship( + self, + rel_type: str, + source_field: str, + target_field: str, + source_label: str | None = None, + target_label: str | None = None, + ) -> Self: ... + def build(self) -> GraphConfig: ... + +@final +class CypherQuery: + def __init__(self, query_text: str) -> None: ... + def with_config(self, config: GraphConfig) -> Self: ... + def with_parameter(self, key: str, value: _Parameter) -> Self: ... + def query_text(self) -> str: ... + def parameters(self) -> dict[str, _Scalar]: ... + def to_sql( + self, datasets: dict[str, _DatasetT], dialect: SqlDialect | None = None + ) -> str: ... + def execute( + self, + datasets: dict[str, _DatasetT], + strategy: ExecutionStrategy | None = None, + ) -> Table: ... + def execute_with_namespace( + self, namespace: DirNamespace, strategy: ExecutionStrategy | None = None + ) -> Table: ... + def explain(self, datasets: dict[str, _DatasetT]) -> str: ... + def variables(self) -> list[str]: ... + def node_labels(self) -> list[str]: ... + def relationship_types(self) -> list[str]: ... + def execute_with_vector_rerank( + self, datasets: dict[str, _DatasetT], vector_search: VectorSearch + ) -> Table: ... + +@final +class CypherEngine: + def __init__( + self, config: GraphConfig, datasets: dict[str, _DatasetT], use_csr: bool = False + ) -> None: ... + def execute( + self, query: str, params: dict[str, _ParameterT] | None = None + ) -> Table: ... + def execute_batches( + self, query: str, params: dict[str, _ParameterT] | None = None + ) -> RecordBatchReader: ... + def explain( + self, query: str, params: dict[str, _ParameterT] | None = None + ) -> str: ... + def refresh(self) -> None: ... + def node_index_stats(self) -> dict[str, _NodeIndexStats]: ... + def csr_stats(self) -> dict[str, _CsrStats]: ... + def dataset_versions(self) -> dict[str, int]: ... + def config(self) -> GraphConfig: ... + def execute_with_vector_rerank( + self, query: str, vector_search: VectorSearch + ) -> Table: ... + +@final +class SqlQuery: + def __init__(self, sql: str) -> None: ... + def sql(self) -> str: ... + def execute(self, datasets: dict[str, _DatasetT]) -> Table: ... + def explain(self, datasets: dict[str, _DatasetT]) -> str: ... + +@final +class SqlEngine: + def __init__(self, datasets: dict[str, _DatasetT]) -> None: ... + def execute(self, sql: str) -> Table: ... + +@final +class VectorSearch: + def __init__(self, column: str) -> None: ... + def query_vector(self, vector: Sequence[float]) -> Self: ... + def metric(self, metric: DistanceMetric) -> Self: ... + def top_k(self, k: int) -> Self: ... + def include_distance(self, include: bool) -> Self: ... + def distance_column_name(self, name: str) -> Self: ... + def use_lance_index(self, enabled: bool) -> Self: ... + def search( + self, table: Table | RecordBatch | RecordBatchReader | _ArrowStream + ) -> Table: ... + +@final +class DirNamespace: + def __init__(self, base_uri: str) -> None: ... + @property + def base_uri(self) -> str: ... diff --git a/python/python/lance_csr/py.typed b/python/python/lance_csr/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/python/python/tests/test_typing.py b/python/python/tests/test_typing.py new file mode 100644 index 0000000..4087b4d --- /dev/null +++ b/python/python/tests/test_typing.py @@ -0,0 +1,149 @@ +"""Keep shipped stubs aligned with the compiled public API.""" + +import ast +import inspect +from pathlib import Path + +import lance_csr +import pytest +from lance_csr import _internal + + +def stub_classes(): + """Read the installed stubs, so the same check also exercises a built wheel.""" + package = Path(lance_csr.__file__).parent + assert (package / "py.typed").is_file() + tree = ast.parse((package / "_internal" / "graph.pyi").read_text()) + return { + node.name: node + for node in tree.body + if isinstance(node, ast.ClassDef) and not node.name.startswith("_") + } + + +def public_members(node): + return { + item.name if isinstance(item, ast.FunctionDef) else item.target.id: item + for item in node.body + if (isinstance(item, ast.FunctionDef) and not item.name.startswith("_")) + or ( + isinstance(item, ast.AnnAssign) + and isinstance(item.target, ast.Name) + and not item.target.id.startswith("_") + ) + } + + +def stub_signature(node): + """Compare Python-callable names, kinds and defaults, ignoring annotations/self.""" + args = node.args + positional = [*args.posonlyargs, *args.args] + defaults = [inspect.Parameter.empty] * (len(positional) - len(args.defaults)) + defaults += [ast.literal_eval(value) for value in args.defaults] + parameters = [] + for index, (arg, default) in enumerate(zip(positional, defaults, strict=True)): + if arg.arg == "self": + continue + kind = ( + inspect.Parameter.POSITIONAL_ONLY + if index < len(args.posonlyargs) + else inspect.Parameter.POSITIONAL_OR_KEYWORD + ) + parameters.append(inspect.Parameter(arg.arg, kind, default=default)) + if args.vararg: + parameters.append( + inspect.Parameter(args.vararg.arg, inspect.Parameter.VAR_POSITIONAL) + ) + for arg, default in zip(args.kwonlyargs, args.kw_defaults, strict=True): + parameters.append( + inspect.Parameter( + arg.arg, + inspect.Parameter.KEYWORD_ONLY, + default=inspect.Parameter.empty + if default is None + else ast.literal_eval(default), + ) + ) + if args.kwarg: + parameters.append( + inspect.Parameter(args.kwarg.arg, inspect.Parameter.VAR_KEYWORD) + ) + return inspect.Signature(parameters) + + +def runtime_signature(value): + return inspect.Signature( + [ + parameter.replace(annotation=inspect.Parameter.empty) + for parameter in inspect.signature(value).parameters.values() + if parameter.name != "self" + ] + ) + + +def check_class(node, runtime): + members = public_members(node) + actual = {name for name in vars(runtime) if not name.startswith("_")} + assert actual == members.keys(), f"{node.name}: runtime/stub public members differ" + for name, member in members.items(): + value = getattr(runtime, name) + if isinstance(member, ast.AnnAssign): + assert isinstance(value, runtime), ( + f"{node.name}.{name}: expected enum member" + ) + elif any( + isinstance(d, ast.Name) and d.id == "property" + for d in member.decorator_list + ): + assert inspect.isgetsetdescriptor(value), ( + f"{node.name}.{name}: expected property" + ) + else: + assert runtime_signature(value) == stub_signature(member), ( + f"{node.name}.{name}" + ) + constructor = next( + ( + item + for item in node.body + if isinstance(item, ast.FunctionDef) and item.name == "__init__" + ), + None, + ) + if constructor is not None: + assert runtime_signature(runtime) == stub_signature(constructor), node.name + + +def test_public_stubs_match_compiled_extension(): + classes = stub_classes() + assert set(lance_csr.__all__) == classes.keys() + assert { + name for name in vars(_internal.graph) if not name.startswith("_") + } == classes.keys() + for name, node in classes.items(): + runtime = getattr(_internal.graph, name) + assert getattr(lance_csr, name) is runtime + check_class(node, runtime) + + +def test_stub_check_detects_missing_method(): + node = stub_classes()["CypherEngine"] + node.body = [ + item + for item in node.body + if not isinstance(item, ast.FunctionDef) or item.name != "refresh" + ] + with pytest.raises(AssertionError, match="public members differ"): + check_class(node, lance_csr.CypherEngine) + + +def test_stub_check_detects_changed_default(): + node = stub_classes()["CypherEngine"] + execute = next( + item + for item in node.body + if isinstance(item, ast.FunctionDef) and item.name == "execute" + ) + execute.args.defaults = [ast.Constant(False)] + with pytest.raises(AssertionError, match="CypherEngine.execute"): + check_class(node, lance_csr.CypherEngine) diff --git a/python/python/tests/typing/public_api.py b/python/python/tests/typing/public_api.py new file mode 100644 index 0000000..1edc328 --- /dev/null +++ b/python/python/tests/typing/public_api.py @@ -0,0 +1,95 @@ +"""Strict static consumer; checked by make lint, not executed by pytest.""" + +from datetime import date +from typing import assert_type + +from lance import LanceDataset +from lance_csr import ( + CypherEngine, + CypherQuery, + DirNamespace, + DistanceMetric, + ExecutionStrategy, + GraphConfig, + GraphConfigBuilder, + SqlDialect, + SqlEngine, + SqlQuery, + VectorSearch, + __version__, +) +from pyarrow import RecordBatch, RecordBatchReader, Table + + +def public_api( + table: Table, batch: RecordBatch, reader: RecordBatchReader, dataset: LanceDataset +) -> None: + assert_type(__version__, str) + builder = GraphConfig.builder().node_table("Person", "id") + assert_type(builder, GraphConfigBuilder) + builder = builder.rel_table("KNOWS", "Person", "Person", "src", "dst") + config = builder.build() + assert_type(config, GraphConfig) + assert_type(config.node_labels(), list[str]) + assert_type(config.relationship_types(), list[str]) + assert_type(GraphConfigBuilder().build(), GraphConfig) + assert_type(builder.with_node_label("City", "id"), GraphConfigBuilder) + assert_type(builder.with_relationship("R", "src", "dst"), GraphConfigBuilder) + + datasets = {"Person": table} + engine = CypherEngine(config, datasets, use_csr=True) + query = "MATCH (p:Person) WHERE p.id = $id RETURN p.id" + params = {"id": 1} + assert_type(engine.execute(query, params), Table) + assert_type(engine.execute_batches(query, params), RecordBatchReader) + assert_type(engine.explain(query, params), str) + assert_type(engine.refresh(), None) + assert_type(engine.config(), GraphConfig) + assert_type(engine.dataset_versions(), dict[str, int]) + assert_type(engine.node_index_stats()["person"]["rows"], int) + assert_type(engine.node_index_stats()["person"]["build_ms"], float) + assert_type(engine.csr_stats()["knows"]["edges"], int) + assert_type(engine.csr_stats()["knows"]["version"], int | None) + CypherEngine(config, {"Person": batch}) + CypherEngine(config, {"Person": reader}) + CypherEngine(config, {"Person": dataset}) + + one = CypherQuery(query).with_config(config).with_parameter("id", 1) + assert_type(one, CypherQuery) + one.with_parameter("day", date(2026, 1, 1)) + assert_type(one.query_text(), str) + assert_type(one.parameters(), dict[str, int | float | str | bool | None]) + assert_type(one.variables(), list[str]) + assert_type(one.node_labels(), list[str]) + assert_type(one.relationship_types(), list[str]) + assert_type(one.execute(datasets, ExecutionStrategy.DataFusion), Table) + assert_type(one.explain(datasets), str) + assert_type(one.to_sql(datasets, SqlDialect.PostgreSql), str) + namespace = DirNamespace("/tmp/graph") + assert_type(namespace.base_uri, str) + assert_type(one.execute_with_namespace(namespace), Table) + + vector = ( + VectorSearch("p.embedding") + .query_vector([1.0, 0.0]) + .metric(DistanceMetric.Cosine) + .top_k(5) + .include_distance(True) + .distance_column_name("distance") + .use_lance_index(False) + ) + assert_type(vector, VectorSearch) + assert_type(vector.search(table), Table) + assert_type(one.execute_with_vector_rerank(datasets, vector), Table) + assert_type(engine.execute_with_vector_rerank(query, vector), Table) + sql = SqlQuery("SELECT id FROM Person") + assert_type(sql.sql(), str) + assert_type(sql.execute(datasets), Table) + assert_type(sql.explain(datasets), str) + assert_type(SqlEngine(datasets).execute(sql.sql()), Table) + + # Rejected calls are checked too: removing a type restriction makes these unused. + engine.execute(query, {"id": object()}) # pyright: ignore[reportArgumentType] + vector.top_k("5") # pyright: ignore[reportArgumentType] + one.with_parameter("ids", [1, 2]) # pyright: ignore[reportArgumentType] + namespace.base_uri = "elsewhere" # pyright: ignore[reportAttributeAccessIssue] diff --git a/python/uv.lock b/python/uv.lock index 26fda10..5a277bf 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -51,6 +51,7 @@ dependencies = [ dev = [ { name = "maturin", extra = ["patchelf"] }, { name = "pandas" }, + { name = "pyarrow-stubs" }, { name = "pyright" }, { name = "pytest" }, { name = "ruff" }, @@ -66,6 +67,7 @@ requires-dist = [ dev = [ { name = "maturin", extras = ["patchelf"], specifier = ">=1.9,<2" }, { name = "pandas" }, + { name = "pyarrow-stubs", specifier = ">=20.0.0.20260819" }, { name = "pyright" }, { name = "pytest" }, { name = "ruff" }, @@ -352,6 +354,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/96/be/7b81a44d6a8e70581dcc1d6f01541f9000a973b1e5d75394aec91e7b179a/pyarrow-26.0.0-cp315-cp315t-win_amd64.whl", hash = "sha256:68cd662e9e2b00876a131950cf32336ace2d0865e1f9418763e3d3be8481dfa4", size = 29389953, upload-time = "2026-10-09T08:26:18.277Z" }, ] +[[package]] +name = "pyarrow-stubs" +version = "20.0.0.20260819" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyarrow" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/12/a7/8a2ca91ffe4c6576207f932de655d7d8a36485c520dccce70ce7d492b256/pyarrow_stubs-20.0.0.20260819.tar.gz", hash = "sha256:150710a72248bc834bf048d3092713f070904a4af76d40289c43afb3ee189823", size = 238222, upload-time = "2026-08-19T05:52:53.618Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/65/6c/eea1d03e475217aea95b1d52aee09c97575d05bbc592c39c085b71dab89f/pyarrow_stubs-20.0.0.20260819-py3-none-any.whl", hash = "sha256:297e60b6e5314739c082b4757d090d8be6047465510eb0684ca954ef7ea58be3", size = 235949, upload-time = "2026-08-19T05:52:54.711Z" }, +] + [[package]] name = "pydantic" version = "2.14.1" From 647018a82c1ad3cfdb0675f34b0d49f58ccd15e8 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:52:36 -0700 Subject: [PATCH 28/66] feat(graph): enforce cooperative query deadlines Add keyword-only timeout to the three Cypher execution entrypoints. Keep one deadline through streamed reads and release the owned stream while idle; preserve native Arrow reader errors and cooperative traversal cancellation. Issue #19. Submitted by David Hughes (@stugorf), with AI Collaboration. External benchmark runner adoption and performance measurements remain pending. --- crates/lance-graph-python/Cargo.toml | 2 +- crates/lance-graph-python/src/executor.rs | 71 ++++++- crates/lance-graph-python/src/graph.rs | 200 +++++++++++++++++++- crates/lance-graph/src/csr_expand.rs | 75 ++++++-- docs/api-reference.md | 17 +- docs/cypher.md | 16 ++ python/python/lance_csr/_internal/graph.pyi | 28 ++- python/python/tests/test_timeout.py | 82 ++++++++ python/python/tests/typing/public_api.py | 24 ++- 9 files changed, 479 insertions(+), 36 deletions(-) create mode 100644 python/python/tests/test_timeout.py diff --git a/crates/lance-graph-python/Cargo.toml b/crates/lance-graph-python/Cargo.toml index 3244cb3..f7a9862 100644 --- a/crates/lance-graph-python/Cargo.toml +++ b/crates/lance-graph-python/Cargo.toml @@ -21,4 +21,4 @@ lance-graph = { path = "../lance-graph" } serde = { version = "1", features = ["derive"] } serde_json = "1" pyo3 = { version = "0.28", features = ["extension-module", "abi3-py312", "py-clone"] } -tokio = { version = "1.37", features = ["rt-multi-thread", "macros"] } +tokio = { version = "1.37", features = ["rt-multi-thread", "macros", "sync", "time"] } diff --git a/crates/lance-graph-python/src/executor.rs b/crates/lance-graph-python/src/executor.rs index cf95aaf..6e619ad 100644 --- a/crates/lance-graph-python/src/executor.rs +++ b/crates/lance-graph-python/src/executor.rs @@ -13,7 +13,54 @@ // limitations under the License. use futures::Future; -use pyo3::{PyResult, Python}; +use pyo3::{ + exceptions::{PyTimeoutError, PyValueError}, + PyResult, Python, +}; +use std::time::Duration; + +pub(crate) const QUERY_TIMEOUT_MESSAGE: &str = "Query timeout exceeded"; + +/// One absolute deadline shared by planning, execution and all streamed reads. +#[derive(Clone, Copy)] +pub(crate) struct QueryDeadline(Option); + +impl QueryDeadline { + pub(crate) fn new(timeout: Option) -> PyResult { + let Some(seconds) = timeout else { + return Ok(Self(None)); + }; + let duration = Duration::try_from_secs_f64(seconds).map_err(|_| { + PyValueError::new_err("timeout must be a finite nonnegative number of seconds") + })?; + let instant = tokio::time::Instant::now() + .checked_add(duration) + .ok_or_else(|| PyValueError::new_err("timeout exceeds the supported deadline range"))?; + Ok(Self(Some(instant))) + } + + pub(crate) fn instant(self) -> Option { + self.0 + } + + pub(crate) fn check(self) -> PyResult<()> { + if self + .0 + .is_some_and(|deadline| tokio::time::Instant::now() >= deadline) + { + Err(PyTimeoutError::new_err(QUERY_TIMEOUT_MESSAGE)) + } else { + Ok(()) + } + } + + async fn expired(self) { + match self.0 { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => std::future::pending::<()>().await, + } + } +} pub const SIGNAL_CHECK_INTERVAL: std::time::Duration = std::time::Duration::from_millis(100); @@ -56,7 +103,21 @@ impl BackgroundExecutor { where F::Output: Send, { - let future = Self::result_or_interrupt(future); + self.block_on_until(py, QueryDeadline(None), future) + } + + /// Apply a query deadline without changing unrelated executor callers. + pub(crate) fn block_on_until( + &self, + py: Option>, + deadline: QueryDeadline, + future: F, + ) -> PyResult + where + F::Output: Send, + { + deadline.check()?; + let future = Self::result_or_interrupt(future, deadline); if let Some(py) = py { py.detach(move || self.runtime.block_on(future)) } else { @@ -65,7 +126,7 @@ impl BackgroundExecutor { } } - async fn result_or_interrupt(future: F) -> PyResult + async fn result_or_interrupt(future: F, deadline: QueryDeadline) -> PyResult where F: Future + Send, F::Output: Send, @@ -84,7 +145,9 @@ impl BackgroundExecutor { }; tokio::select! { - result = future => Ok(result), + biased; + _ = deadline.expired() => Err(PyTimeoutError::new_err(QUERY_TIMEOUT_MESSAGE)), + result = future => { deadline.check()?; Ok(result) }, err = interrupt_future => err, } } diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 4690d76..17dcd95 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -47,6 +47,7 @@ use pyo3::{ }; use serde_json::Value as JsonValue; +use crate::executor::{QueryDeadline, QUERY_TIMEOUT_MESSAGE}; use crate::namespace::PyDirNamespace; use crate::RT; @@ -662,13 +663,16 @@ impl CypherQuery { /// >>> # Explicit strategy /// >>> from lance.graph import ExecutionStrategy /// >>> result = query.execute(datasets, strategy=ExecutionStrategy.DataFusion) - #[pyo3(signature = (datasets, strategy=None))] + #[pyo3(signature = (datasets, strategy=None, *, timeout=None))] fn execute( &self, py: Python, datasets: &Bound<'_, PyDict>, strategy: Option, + timeout: Option, ) -> PyResult> { + let deadline = QueryDeadline::new(timeout)?; + deadline.check()?; // Convert datasets to Arrow batches while holding the GIL let arrow_datasets = python_datasets_to_batches(datasets)?; @@ -680,7 +684,11 @@ impl CypherQuery { // Use RT.block_on with Some(py) like the scanner to_pyarrow method let result_batch = RT - .block_on(Some(py), inner_query.execute(arrow_datasets, rust_strategy))? + .block_on_until( + Some(py), + deadline, + inner_query.execute(arrow_datasets, rust_strategy), + )? .map_err(graph_error_to_pyerr)?; record_batch_to_python_table(py, &result_batch) @@ -1946,6 +1954,120 @@ impl RecordBatchReader for BlockingBatchReader { } } +/// A timed reader owns its stream independently of calls into Arrow's C interface. +struct TimedBatchState { + stream: Option, + expired: bool, +} + +struct DeadlineBatchReader { + schema: arrow_schema::SchemaRef, + state: Arc>, + deadline: tokio::time::Instant, + watchdog: tokio::task::JoinHandle<()>, + finished: bool, +} + +fn stream_timeout_error() -> arrow_schema::ArrowError { + arrow_schema::ArrowError::ExternalError(Box::new(std::io::Error::new( + std::io::ErrorKind::TimedOut, + QUERY_TIMEOUT_MESSAGE, + ))) +} + +impl Iterator for DeadlineBatchReader { + type Item = Result; + + fn next(&mut self) -> Option { + if self.finished { + return None; + } + let result = RT.runtime.block_on(async { + let mut state = self.state.lock().await; + if state.expired || tokio::time::Instant::now() >= self.deadline { + state.stream.take(); + state.expired = true; + return Some(Err(stream_timeout_error())); + } + let stream = state.stream.as_mut()?; + let result = tokio::select! { + biased; + _ = tokio::time::sleep_until(self.deadline) => None, + result = stream.next() => Some(result), + }; + match result { + None => { + state.stream.take(); + state.expired = true; + Some(Err(stream_timeout_error())) + } + Some(result) => { + if tokio::time::Instant::now() >= self.deadline { + state.stream.take(); + state.expired = true; + return Some(Err(stream_timeout_error())); + } + if result.is_none() { + state.stream.take(); + } + result.map(|r| { + r.map_err(|e| arrow_schema::ArrowError::ExternalError(Box::new(e))) + }) + } + } + }); + if result.as_ref().is_none_or(|r| r.is_err()) { + self.finished = true; + self.watchdog.abort(); + } + result + } +} + +impl RecordBatchReader for DeadlineBatchReader { + fn schema(&self) -> arrow_schema::SchemaRef { + self.schema.clone() + } +} + +impl Drop for DeadlineBatchReader { + fn drop(&mut self) { + self.watchdog.abort(); + if let Ok(mut state) = self.state.try_lock() { + state.stream.take(); + } + } +} + +fn reader_with_deadline( + stream: datafusion::physical_plan::SendableRecordBatchStream, + deadline: QueryDeadline, +) -> Box { + let schema = stream.schema(); + let Some(deadline) = deadline.instant() else { + return Box::new(BlockingBatchReader { schema, stream }); + }; + let state = Arc::new(tokio::sync::Mutex::new(TimedBatchState { + stream: Some(stream), + expired: false, + })); + let watched = state.clone(); + let watchdog = RT.runtime.spawn(async move { + tokio::time::sleep_until(deadline).await; + let mut state = watched.lock().await; + if state.stream.take().is_some() { + state.expired = true; + } + }); + Box::new(DeadlineBatchReader { + schema, + state, + deadline, + watchdog, + finished: false, + }) +} + /// Cypher query engine with cached catalog for efficient multi-query execution /// /// The engine registers each dataset once and reuses the catalog and DataFusion @@ -2126,6 +2248,8 @@ impl CypherEngine { /// Execute a Cypher query and return the full result as a ``pyarrow.Table``. /// /// The whole result is held in memory; use ``execute_batches`` for large outputs. + /// ``timeout`` is an optional nonnegative number of seconds for the whole query; + /// expiry raises ``TimeoutError``. Cancellation is cooperative. /// /// Parameters /// ---------- @@ -2142,18 +2266,22 @@ impl CypherEngine { /// If the query is invalid /// RuntimeError /// If query execution fails - #[pyo3(signature = (query, params = None))] + #[pyo3(signature = (query, params = None, *, timeout=None))] fn execute( &self, py: Python, query: &str, params: Option<&Bound<'_, PyDict>>, + timeout: Option, ) -> PyResult> { + let deadline = QueryDeadline::new(timeout)?; + deadline.check()?; let state = self.snapshot(); let cypher_query = self.parse(py, query, params, &state, true)?; let result_batch = RT - .block_on( + .block_on_until( Some(py), + deadline, cypher_query .execute_with_catalog_and_context(state.catalog.clone(), state.context.clone()), )? @@ -2168,30 +2296,34 @@ 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. - #[pyo3(signature = (query, params = None))] + /// ``timeout`` keeps one deadline including pauses. Initial expiry raises + /// ``TimeoutError``; expired reads raise ArrowInvalid: Query timeout exceeded. + #[pyo3(signature = (query, params = None, *, timeout=None))] fn execute_batches( &self, py: Python, query: &str, params: Option<&Bound<'_, PyDict>>, + timeout: Option, ) -> PyResult> { use arrow::pyarrow::IntoPyArrow; + let deadline = QueryDeadline::new(timeout)?; + deadline.check()?; let state = self.snapshot(); let cypher_query = self.parse(py, query, params, &state, true)?; let stream = RT - .block_on( + .block_on_until( Some(py), + deadline, cypher_query.execute_stream_with_catalog_and_context( state.catalog.clone(), state.context.clone(), ), )? .map_err(graph_error_to_pyerr)?; - let reader: Box = Box::new(BlockingBatchReader { - schema: stream.schema(), - stream, - }); + deadline.check()?; + let reader = reader_with_deadline(stream, deadline); Ok(reader.into_pyarrow(py)?.unbind()) } @@ -2742,6 +2874,54 @@ mod tests { } } + fn retained_stream() -> ( + datafusion::physical_plan::SendableRecordBatchStream, + std::sync::Weak, + ) { + let values = Arc::new(arrow_array::Int64Array::from(vec![1_i64])); + let retained = Arc::downgrade(&values); + let schema = Arc::new(Schema::new(vec![arrow_schema::Field::new( + "id", + arrow_schema::DataType::Int64, + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap(); + let stream = datafusion::physical_plan::stream::RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(vec![Ok(batch)]), + ); + (Box::pin(stream), retained) + } + + #[test] + fn timed_reader_releases_unconsumed_stream_at_deadline() { + let (stream, retained) = retained_stream(); + let mut reader = reader_with_deadline(stream, QueryDeadline::new(Some(0.02)).unwrap()); + let stop = std::time::Instant::now() + std::time::Duration::from_secs(1); + while retained.upgrade().is_some() && std::time::Instant::now() < stop { + std::thread::sleep(std::time::Duration::from_millis(2)); + } + assert!( + retained.upgrade().is_none(), + "idle deadline must drop the owned stream" + ); + assert!(reader + .next() + .unwrap() + .unwrap_err() + .to_string() + .contains(QUERY_TIMEOUT_MESSAGE)); + assert!(reader.next().is_none()); + } + + #[test] + fn closing_timed_reader_releases_stream_before_deadline() { + let (stream, retained) = retained_stream(); + let reader = reader_with_deadline(stream, QueryDeadline::new(Some(10.0)).unwrap()); + drop(reader); + assert!(retained.upgrade().is_none()); + } + #[test] fn test_split_vector_column_with_alias() { let (alias, column) = split_vector_column("d.embedding"); diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 66b009f..afe1987 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -1234,19 +1234,27 @@ impl BatchExpansion { } /// Collect only the next chunk; target properties are fetched after this returns. - fn next_chunk( + async fn next_chunk( &mut self, adjacency: &Neighbors, target: Option<&TargetLookup>, batch_size: usize, ) -> ExpandedRows { let mut chunk = ExpandedRows::default(); + let mut work = 0_usize; while self.row < self.keys.len() { let i = self.row; if self.keys.is_valid(i) && self.bound.as_ref().is_none_or(|b| b.is_valid(i)) { let key = self.keys.value(i); let required = self.bound.as_ref().map(|b| b.value(i)); - for neighbor in adjacency.neighbors_from(key, &mut self.position) { + while let Some(candidate) = adjacency.next_candidate(key, &mut self.position) { + work += 1; + if work.is_multiple_of(1024) { + cooperative_yield().await; + } + let Some(neighbor) = candidate else { + continue; + }; let neighbor = neighbor as i64; if required.is_some_and(|r| r != neighbor) { continue; @@ -1580,7 +1588,11 @@ impl ExecutionPlan for CsrExpandExec { .next_path_chunk(walk, target.as_deref(), batch_size) .await } - None => expansion.next_chunk(&adjacency, target.as_deref(), batch_size), + None => { + expansion + .next_chunk(&adjacency, target.as_deref(), batch_size) + .await + } }; if chunk.rows.is_empty() { pending = None; @@ -2036,21 +2048,58 @@ mod tests { ])); let adj = Neighbors::directed(Arc::new(adj)); let mut expansion = BatchExpansion::try_new(input.clone(), 0, None).unwrap(); - let chunk = expansion.next_chunk(&adj, None, 8192); + let chunk = expansion.next_chunk(&adj, None, 8192).await; let out = materialize(&input, chunk, None, None, &out_schema) .await .unwrap(); assert_eq!(out.num_rows(), 2); let dst = out.column(2).as_primitive::(); assert_eq!(dst.values().to_vec(), vec![2, 3]); - assert!(expansion.next_chunk(&adj, None, 8192).rows.is_empty()); + assert!(expansion.next_chunk(&adj, None, 8192).await.rows.is_empty()); // Output is split at `batch_size` rows. let mut expansion = BatchExpansion::try_new(input, 0, None).unwrap(); - assert_eq!(expansion.next_chunk(&adj, None, 1).neighbor_values, vec![2]); + assert_eq!( + expansion.next_chunk(&adj, None, 1).await.neighbor_values, + vec![2] + ); assert_eq!(expansion.row, 0); assert_eq!(expansion.position.forward, 1); - assert_eq!(expansion.next_chunk(&adj, None, 1).neighbor_values, vec![3]); - assert!(expansion.next_chunk(&adj, None, 1).rows.is_empty()); + assert_eq!( + expansion.next_chunk(&adj, None, 1).await.neighbor_values, + vec![3] + ); + assert!(expansion.next_chunk(&adj, None, 1).await.rows.is_empty()); + } + + #[tokio::test] + async fn expansion_yields_before_scanning_all_rejected_neighbors() { + let adj = Neighbors::directed(Arc::new( + Adjacency::build( + edge_stream(vec![Some(1); 100_000], (2..100_002).map(Some).collect()), + "src", + "dst", + ) + .await + .unwrap(), + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("bound", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1])), + Arc::new(Int64Array::from(vec![-1])), + ], + ) + .unwrap(); + let mut expansion = BatchExpansion::try_new(batch, 0, Some(1)).unwrap(); + let mut next = Box::pin(expansion.next_chunk(&adj, None, 1)); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert!(expansion.position.forward <= 1024); + assert!(expansion.position.forward > 0); } #[tokio::test] @@ -2069,7 +2118,7 @@ mod tests { RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(vec![0])) as ArrayRef)]) .unwrap(); let mut expansion = BatchExpansion::try_new(input, 0, None).unwrap(); - let chunk = expansion.next_chunk(&adj, None, 32); + let chunk = expansion.next_chunk(&adj, None, 32).await; assert_eq!( chunk.neighbor_values, (2..=64).step_by(2).collect::>() @@ -2078,7 +2127,10 @@ mod tests { // The rest of the hub, including its tombstones, has not been visited. assert_eq!(expansion.row, 0); assert_eq!(expansion.position.forward, 64); - assert_eq!(expansion.next_chunk(&adj, None, 32).neighbor_values[0], 66); + assert_eq!( + expansion.next_chunk(&adj, None, 32).await.neighbor_values[0], + 66 + ); assert_eq!(neighbors(&base, 0).len(), n); } @@ -2113,13 +2165,14 @@ mod tests { ); let mut expansion = BatchExpansion::try_new(input, 0, Some(1)).unwrap(); for _ in 0..2 { - let chunk = expansion.next_chunk(&adjacency, Some(&target), 1); + let chunk = expansion.next_chunk(&adjacency, Some(&target), 1).await; assert_eq!(chunk.rows, vec![0]); assert_eq!(chunk.neighbor_values, vec![2]); assert_eq!(chunk.target_positions, vec![0]); } assert!(expansion .next_chunk(&adjacency, Some(&target), 1) + .await .rows .is_empty()); } diff --git a/docs/api-reference.md b/docs/api-reference.md index 0162102..c964eab 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -68,6 +68,17 @@ Keys match labels case-insensitively. --- +`CypherEngine.execute`, `CypherEngine.execute_batches` and `CypherQuery.execute` accept keyword-only +`timeout` in seconds (`None` means no deadline). Zero expires immediately; negative, nonfinite or +unrepresentable durations raise `ValueError`. One deadline includes planning, every batch and pauses +between reads. Before returning a reader, expiry raises `TimeoutError`. Reading an expired native Arrow +reader raises `pyarrow.ArrowInvalid` containing `Query timeout exceeded`. Earlier batches remain valid. +Expiry or reader close releases the stream even while the caller is idle; the engine remains usable. + +Cancellation is cooperative. Traversal loops yield at bounded work intervals, but synchronous parsing, +Python data providers and foreign blocking calls cannot be forcibly preempted. Deadlines are checked +between planning stages and before accepting completed results; this is not a hard real-time limit. + ## `GraphConfig` An immutable mapping from tables to a property graph. Create one with `GraphConfig.builder()`. @@ -132,14 +143,14 @@ Raises `ValueError` if `datasets` is empty, a configured key or endpoint column was opened through a namespace client and can't be reopened by URI. Raises `RuntimeError` if a dataset can't be opened. -`#!python execute(query: str, params: dict | None = None) -> pyarrow.Table` +`#!python execute(query: str, params: dict | None = None, *, timeout: float | None = None) -> pyarrow.Table` : Runs `query` and returns the whole result. `params` binds `$name` parameters: `int`, `float`, `str`, `bool`, `None`, lists, string-keyed nested maps, `datetime.datetime` and `datetime.date`. Integers must fit signed 64-bit values; nesting beyond 128 levels, cycles and unsupported values raise `ValueError`. Direct counts from `size()` of an unmapped list comprehension have a non-nullable Arrow field; see [Cypher functions](cypher.md#functions). -`#!python execute_batches(query: str, params: dict | None = None) -> pyarrow.RecordBatchReader` +`#!python execute_batches(query: str, params: dict | None = None, *, timeout: float | None = None) -> pyarrow.RecordBatchReader` : Streams the result. Planning errors raise immediately; execution errors raise while reading. The reader keeps its snapshot alive across `refresh()`. @@ -190,7 +201,7 @@ A parsed Cypher query. Raises `ValueError` if `query_text` doesn't parse. Builde `#!python with_parameter(key: str, value: object) -> CypherQuery` : Binds `$key` to `value`. -`#!python execute(datasets: dict, strategy: ExecutionStrategy | None = None) -> pyarrow.Table` +`#!python execute(datasets: dict, strategy: ExecutionStrategy | None = None, *, timeout: float | None = None) -> pyarrow.Table` : Runs the query. Datasets are read into memory for this call. `#!python execute_with_namespace(namespace: DirNamespace, strategy: ExecutionStrategy | None = None) -> pyarrow.Table` diff --git a/docs/cypher.md b/docs/cypher.md index e363fca..480797b 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -124,6 +124,22 @@ print(rows) # 4 - The reader keeps its snapshot alive, so a concurrent [`refresh()`](lance.md#refresh) doesn't affect it. - Dropping the reader cancels the remaining work. +### Query deadlines + +Pass `timeout=5.0` to `execute`, `execute_batches` or `CypherQuery.execute` to allow five seconds from the +call's start. Streams keep that same deadline while you pause between reads. The default is `None`. + +```python +try: + engine.execute("RETURN 1 AS n", timeout=0) +except TimeoutError: + print("Query timeout exceeded") # Query timeout exceeded +``` + +Initial expiry raises `TimeoutError`. Reading an expired native Arrow reader raises `pyarrow.ArrowInvalid` +with `Query timeout exceeded`; partial output already read remains valid. Close readers when you stop +consuming them. See the [API reference](api-reference.md) for validation and cooperative cancellation limits. + ### `explain(query, params=None)` Returns the plan as text: the graph logical plan, DataFusion's logical plan and the physical plan. diff --git a/python/python/lance_csr/_internal/graph.pyi b/python/python/lance_csr/_internal/graph.pyi index 1c6b13a..e6544f7 100644 --- a/python/python/lance_csr/_internal/graph.pyi +++ b/python/python/lance_csr/_internal/graph.pyi @@ -10,9 +10,19 @@ class _ArrowStream(Protocol): _Dataset = Table | RecordBatch | RecordBatchReader | LanceDataset | _ArrowStream _DatasetT = TypeVar("_DatasetT", bound=_Dataset) -_Parameter = int | float | str | bool | date | datetime | None +_Parameter = ( + int + | float + | str + | bool + | date + | datetime + | None + | list["_Parameter"] + | dict[str, "_Parameter"] +) _ParameterT = TypeVar("_ParameterT", bound=_Parameter) -_Scalar = int | float | str | bool | None +_Scalar = int | float | str | bool | None | list["_Scalar"] | dict[str, "_Scalar"] class _NodeIndexStats(TypedDict): rows: int @@ -88,6 +98,8 @@ class CypherQuery: self, datasets: dict[str, _DatasetT], strategy: ExecutionStrategy | None = None, + *, + timeout: float | None = None, ) -> Table: ... def execute_with_namespace( self, namespace: DirNamespace, strategy: ExecutionStrategy | None = None @@ -106,10 +118,18 @@ class CypherEngine: self, config: GraphConfig, datasets: dict[str, _DatasetT], use_csr: bool = False ) -> None: ... def execute( - self, query: str, params: dict[str, _ParameterT] | None = None + self, + query: str, + params: dict[str, _ParameterT] | None = None, + *, + timeout: float | None = None, ) -> Table: ... def execute_batches( - self, query: str, params: dict[str, _ParameterT] | None = None + self, + query: str, + params: dict[str, _ParameterT] | None = None, + *, + timeout: float | None = None, ) -> RecordBatchReader: ... def explain( self, query: str, params: dict[str, _ParameterT] | None = None diff --git a/python/python/tests/test_timeout.py b/python/python/tests/test_timeout.py new file mode 100644 index 0000000..87db8f9 --- /dev/null +++ b/python/python/tests/test_timeout.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""A single deadline covers execution and idle time in native Arrow readers.""" + +import time + +import lance +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, CypherQuery, GraphConfig + +CONFIG = GraphConfig.builder().node_table("Person", "id").build() +TABLES = {"Person": pa.table({"id": list(range(20_000))})} +QUERY = "MATCH (p:Person) RETURN p.id AS id" + + +@pytest.fixture +def engine(): + return CypherEngine(CONFIG, TABLES) + + +@pytest.mark.parametrize("timeout", [-1, float("nan"), float("inf"), 1e300]) +@pytest.mark.parametrize("method", ["execute", "execute_batches", "query"]) +def test_invalid_timeout(engine, timeout, method): + with pytest.raises(ValueError, match="timeout"): + if method == "query": + CypherQuery(QUERY).with_config(CONFIG).execute(TABLES, timeout=timeout) + else: + getattr(engine, method)(QUERY, timeout=timeout) + + +@pytest.mark.parametrize("method", ["execute", "execute_batches", "query"]) +def test_zero_timeout_then_engine_reuse(engine, method): + with pytest.raises(TimeoutError, match="Query timeout exceeded"): + if method == "query": + CypherQuery(QUERY).with_config(CONFIG).execute(TABLES, timeout=0) + else: + getattr(engine, method)(QUERY, timeout=0) + assert engine.execute("RETURN 1 AS n").to_pylist() == [{"n": 1}] + + +def test_native_reader_deadline_expires_while_idle(engine): + reader = engine.execute_batches(QUERY, timeout=1.0) + assert isinstance(reader, pa.RecordBatchReader) + time.sleep(1.05) + with pytest.raises(pa.ArrowInvalid, match="Query timeout exceeded"): + reader.read_next_batch() + reader.close() + assert engine.execute("RETURN 1 AS n", timeout=10).to_pylist() == [{"n": 1}] + + +def test_stream_keeps_one_deadline_between_reads(tmp_path): + dataset = lance.write_dataset( + TABLES["Person"], tmp_path / "people.lance", max_rows_per_file=4000 + ) + engine = CypherEngine(CONFIG, {"Person": dataset}) + reader = engine.execute_batches(QUERY, timeout=1.0) + first = reader.read_next_batch() + assert 0 < first.num_rows < 20_000 + time.sleep(1.05) + with pytest.raises(pa.ArrowInvalid, match="Query timeout exceeded"): + reader.read_all() + + +def test_no_deadline_and_normal_execution_errors(engine): + assert engine.execute(QUERY, timeout=None).num_rows == 20_000 + assert engine.execute_batches(QUERY, timeout=10).read_all().num_rows == 20_000 + with pytest.raises(ValueError): + engine.execute("RETURN unknown", timeout=10) + + +def test_expensive_query_stops_and_engine_stays_usable(engine): + started = time.monotonic() + with pytest.raises(TimeoutError, match="Query timeout exceeded"): + engine.execute( + "MATCH (a:Person), (b:Person), (c:Person) " + "WHERE a.id + b.id + c.id > 30000 RETURN count(*) AS n", + timeout=0.02, + ) + assert time.monotonic() - started < 2.0 + assert engine.execute("RETURN 2 AS n").to_pylist() == [{"n": 2}] diff --git a/python/python/tests/typing/public_api.py b/python/python/tests/typing/public_api.py index 1edc328..c4e3704 100644 --- a/python/python/tests/typing/public_api.py +++ b/python/python/tests/typing/public_api.py @@ -1,7 +1,7 @@ """Strict static consumer; checked by make lint, not executed by pytest.""" from datetime import date -from typing import assert_type +from typing import TYPE_CHECKING, TypeAlias, assert_type from lance import LanceDataset from lance_csr import ( @@ -20,6 +20,19 @@ ) from pyarrow import RecordBatch, RecordBatchReader, Table +if TYPE_CHECKING: + from collections.abc import Mapping + +ParameterResult: TypeAlias = ( + int + | float + | str + | bool + | None + | list["ParameterResult"] + | dict[str, "ParameterResult"] +) + def public_api( table: Table, batch: RecordBatch, reader: RecordBatchReader, dataset: LanceDataset @@ -56,9 +69,14 @@ def public_api( one = CypherQuery(query).with_config(config).with_parameter("id", 1) assert_type(one, CypherQuery) + assert_type(one.with_parameter("ids", [1, 2]), CypherQuery) + assert_type(engine.execute(query, timeout=1.0), Table) + assert_type(engine.execute_batches(query, timeout=1.0), RecordBatchReader) + assert_type(one.execute(datasets, timeout=1.0), Table) one.with_parameter("day", date(2026, 1, 1)) assert_type(one.query_text(), str) - assert_type(one.parameters(), dict[str, int | float | str | bool | None]) + parameters: Mapping[str, ParameterResult] = one.parameters() + assert_type(len(parameters), int) assert_type(one.variables(), list[str]) assert_type(one.node_labels(), list[str]) assert_type(one.relationship_types(), list[str]) @@ -91,5 +109,5 @@ def public_api( # Rejected calls are checked too: removing a type restriction makes these unused. engine.execute(query, {"id": object()}) # pyright: ignore[reportArgumentType] vector.top_k("5") # pyright: ignore[reportArgumentType] - one.with_parameter("ids", [1, 2]) # pyright: ignore[reportArgumentType] + one.with_parameter("ids", {1, 2}) # pyright: ignore[reportArgumentType] namespace.base_uri = "elsewhere" # pyright: ignore[reportAttributeAccessIssue] From b05ed30432f9c27b0f154f2b7ff70406cfa46f47 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:28:10 -0700 Subject: [PATCH 29/66] fix(query): retain scope and snapshot boundaries across features Keep one cooperative scheduling helper, bind and restore lambda-local types, and bypass cached templates for extension-owned path predicates and query-local inventories. Add CALL/UNION/path/fold, refresh isolation, parameter rebinding and deadline regressions. Combined runtime validation follows the inventory partition fix. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/src/apply.rs | 2 + crates/lance-graph/src/csr_expand.rs | 14 --- crates/lance-graph/src/query.rs | 4 +- crates/lance-graph/src/semantic.rs | 23 +++- crates/lance-graph/src/semantic/types.rs | 11 +- docs/api-reference.md | 2 +- docs/cypher.md | 2 + python/python/tests/test_call_subqueries.py | 117 ++++++++++++++++++++ 8 files changed, 152 insertions(+), 23 deletions(-) diff --git a/crates/lance-graph/src/apply.rs b/crates/lance-graph/src/apply.rs index d2cf84f..50dfe32 100644 --- a/crates/lance-graph/src/apply.rs +++ b/crates/lance-graph/src/apply.rs @@ -468,6 +468,8 @@ mod tests { "WITH 1 AS x WITH x AS y CALL { WITH y RETURN y AS n } RETURN n", "RETURN 1 AS n UNION CALL { RETURN 2 AS n } RETURN n", "RETURN 1 AS n UNION RETURN 2 AS n LIMIT $take", + "MATCH path = (p:Person)-[:KNOWS]->(f:Person) RETURN path", + "MATCH path = shortestPath((p:Person)-[:KNOWS*..3]->(f:Person)) RETURN relationships(path)", ] { let mut query = CypherQuery::new(text).unwrap(); assert!(!query.prepare_cached_plan(catalog.clone()).unwrap()); diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index afe1987..0e8a6b3 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -1447,20 +1447,6 @@ fn retain_by_mask(values: &mut Vec, keep: &[bool]) { values.retain(|_| *keep.next().expect("one mask entry per value")); } -/// Cooperate with the stream's executor without requiring a particular async runtime. -pub(crate) async fn cooperative_yield() { - let mut yielded = false; - std::future::poll_fn(move |cx| { - if std::mem::replace(&mut yielded, true) { - std::task::Poll::Ready(()) - } else { - cx.waker().wake_by_ref(); - std::task::Poll::Pending - } - }) - .await -} - impl DisplayAs for CsrExpandExec { fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { write!( diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 082adb6..0b33a65 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -235,7 +235,9 @@ impl CypherQuery { &mut self, catalog: Arc, ) -> Result { - let mut shape_sensitive = false; + // Rich path properties live in an extension-owned spec, and its inventory is + // query-local. Keep ordinary substitution and a fresh inventory for each query. + let mut shape_sensitive = crate::path_value::needs_values(&self.ast); self.ast.for_each_query(&mut |query| { shape_sensitive |= query .limit diff --git a/crates/lance-graph/src/semantic.rs b/crates/lance-graph/src/semantic.rs index ea7a9eb..9a8533c 100644 --- a/crates/lance-graph/src/semantic.rs +++ b/crates/lance-graph/src/semantic.rs @@ -917,7 +917,14 @@ impl SemanticAnalyzer { location: snafu::Location::new(file!(), line!(), column!()), }); } - self.analyze_lambda(&[accumulator, variable], expression, false)?; + self.analyze_lambda( + &[ + (accumulator, self.expression_type(initial)), + (variable, self.expression_type(list).element()), + ], + expression, + false, + )?; } ValueExpression::ListPredicate { variable, @@ -926,7 +933,11 @@ impl SemanticAnalyzer { .. } => { self.analyze_value_expression(list)?; - self.analyze_lambda(&[variable], predicate, true)?; + self.analyze_lambda( + &[(variable, self.expression_type(list).element())], + predicate, + true, + )?; } ValueExpression::ListComprehension { variable, @@ -961,14 +972,17 @@ impl SemanticAnalyzer { fn analyze_lambda( &mut self, - names: &[&String], + names: &[(&String, ExpressionType)], expression: &ValueExpression, allow_patterns: bool, ) -> Result<()> { let before = self.variables.clone(); - for name in names { + let before_types = self.expression_types.clone(); + for (name, value_type) in names { self.variables.remove(&name.to_lowercase()); self.register_projection_alias(name); + self.expression_types + .insert(name.to_lowercase(), value_type.clone()); } let mut unsupported = !allow_patterns && crate::ast::needs_relational_evaluation(expression); @@ -984,6 +998,7 @@ impl SemanticAnalyzer { self.analyze_value_expression(expression) }; self.variables = before; + self.expression_types = before_types; result } diff --git a/crates/lance-graph/src/semantic/types.rs b/crates/lance-graph/src/semantic/types.rs index 7331815..fad1621 100644 --- a/crates/lance-graph/src/semantic/types.rs +++ b/crates/lance-graph/src/semantic/types.rs @@ -90,7 +90,9 @@ impl SemanticAnalyzer { pub(super) fn expression_type(&self, expression: &ValueExpression) -> ExpressionType { use ExpressionType::{List, Map, Scalar, Unknown}; match expression { - ValueExpression::Variable(name) => self.variable_expression_type(name), + ValueExpression::Variable(name) | ValueExpression::LocalVariable(name) => { + self.variable_expression_type(name) + } ValueExpression::Property(property) | ValueExpression::Literal(PropertyValue::Property(property)) => self .property_expression_type( @@ -165,7 +167,9 @@ impl SemanticAnalyzer { .chain(otherwise.iter().map(|value| self.expression_type(value))) .collect(), ), - ValueExpression::Boolean(_) => Scalar(DataType::Boolean), + ValueExpression::Boolean(_) | ValueExpression::ListPredicate { .. } => { + Scalar(DataType::Boolean) + } ValueExpression::Map(fields) => Map(fields .iter() .map(|(name, value)| (name.clone(), self.expression_type(value))) @@ -178,7 +182,8 @@ impl SemanticAnalyzer { Scalar(DataType::Float64) } // Their local scopes have ended; do not infer from similarly named outer variables. - ValueExpression::ListComprehension { .. } + ValueExpression::Reduce { .. } + | ValueExpression::ListComprehension { .. } | ValueExpression::PatternComprehension { .. } | ValueExpression::Parameter(_) => Unknown, } diff --git a/docs/api-reference.md b/docs/api-reference.md index c964eab..43ccf0f 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -129,7 +129,7 @@ CypherEngine(config: GraphConfig, datasets: dict, use_csr: bool = False) Registers `datasets` once and reuses the catalog, DataFusion session and (with `use_csr=True`) CSR index across queries. Safe to share between threads. `execute()` and `execute_batches()` reuse up to 128 unoptimized query templates per snapshot, keyed by exact text and scalar parameter names/types. Values are -bound before optimization on every call. Containers, `CALL` subqueries, parameterized `SKIP`/`LIMIT` and unsupported template +bound before optimization on every call. Containers, rich path values, `CALL` subqueries, parameterized `SKIP`/`LIMIT` and unsupported template shapes use ordinary planning. `explain()` always plans directly; successful `refresh()` clears the cache. See [Repeated queries](lance.md#repeated-queries) for the scope and limits. diff --git a/docs/cypher.md b/docs/cypher.md index 480797b..886d0b0 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -357,6 +357,8 @@ property types fail explicitly. Expanding a graph pattern from a mixed-label val A path value needs a query-local inventory of the participating node and relationship tables, including properties. It therefore materializes those tables, even with `use_csr=True`. Traversal output is emitted in batches and named ordinary paths keep a cursor per hop, but inventory memory grows with the graph. +Rich path queries bypass the query-template cache: each execution binds its current parameters and +creates an inventory from its own snapshot. A reader opened before `refresh()` retains the old snapshot. Queries asking only for shortest-path length or node keys, including existing homogeneous node-property comprehensions, retain the usual CSR execution route. diff --git a/python/python/tests/test_call_subqueries.py b/python/python/tests/test_call_subqueries.py index 8da6d32..b01679a 100644 --- a/python/python/tests/test_call_subqueries.py +++ b/python/python/tests/test_call_subqueries.py @@ -324,3 +324,120 @@ def test_call_union_coerces_numeric_and_null_output_schema(graph): reader = engine.execute_batches(query, {"value": value}) assert reader.schema.field("n").type == expected_type assert reader.read_all().column(0).to_pylist() == expected + + +def test_call_imported_shortest_path_keeps_relationship_records(graph): + engine, _, _ = graph + query = ( + "MATCH path = allShortestPaths((a:Person {id: 10})-[:KNOWS*..3]->" + "(b:Person {id: 40})) CALL { WITH path " + "RETURN [edge IN relationships(path) | endNode(edge).id] AS targets, " + "reduce(total = 0, edge IN relationships(path) | total + endNode(edge).id) AS total } " + "RETURN targets, total ORDER BY targets" + ) + assert engine.execute(query).to_pylist() == [ + {"targets": [20, 40], "total": 60}, + {"targets": [20, 40], "total": 60}, + ] + + +def test_call_union_shares_path_occurrence_identity(graph): + engine, _, _ = graph + query = ( + "CALL { MATCH path = (a:Person {id: 10})-[:KNOWS]->(b:Person) " + "RETURN path UNION MATCH path = (a:Person {id: 10})-[:KNOWS]->(b:Person) " + "RETURN path } WITH collect(path) AS paths UNWIND paths AS path " + "RETURN length(path) AS hops, [n IN nodes(path) | n.name] AS names, " + "[r IN relationships(path) | endNode(r).id] AS targets ORDER BY targets" + ) + assert engine.execute(query).to_pylist() == [ + {"hops": 1, "names": ["A", "B"], "targets": [20]}, + {"hops": 1, "names": ["A", "B"], "targets": [20]}, + {"hops": 1, "names": ["A", "C"], "targets": [30]}, + ] + + +def test_call_keeps_path_alias_across_nested_union(graph): + engine, _, _ = graph + query = ( + "MATCH path = (a:Person {id: 20})-[:KNOWS]->(b:Person) " + "CALL { WITH path CALL { WITH path RETURN path AS value " + "UNION ALL WITH path RETURN path AS value LIMIT 0 } RETURN value AS result } " + "RETURN length(result) AS hops, [n IN nodes(result) | labels(n)] AS labels" + ) + assert engine.execute(query).to_pylist() == [ + {"hops": 1, "labels": [["Person"], ["Person"]]} + ] + + +def test_call_empty_children_respect_query_deadline(graph): + engine, _, _ = graph + with pytest.raises(TimeoutError): + engine.execute( + "UNWIND range(1, 1000000) AS x CALL { WITH x UNWIND [] AS n RETURN n } RETURN n", + timeout=0.02, + ) + assert engine.execute("RETURN 1 AS ok").to_pylist() == [{"ok": 1}] + + +def test_rich_path_parameters_bind_freshly(graph): + engine, _, _ = graph + query = ( + "MATCH path = (a:Person {id: $start})-[:KNOWS]->(b:Person) " + "RETURN nodes(path) AS ids ORDER BY ids" + ) + for start, expected in [ + (10, [[10, 20], [10, 20], [10, 30]]), + (20, [[20, 40]]), + (10, [[10, 20], [10, 20], [10, 30]]), + ]: + assert engine.execute(query, {"start": start}).column(0).to_pylist() == expected + + +def test_rich_path_reader_retains_snapshot_after_refresh(graph, tmp_path): + import lance + + original, config, data = graph + paths = {name: str(tmp_path / f"{name}.lance") for name in data} + datasets = { + name: lance.write_dataset(table, paths[name]) for name, table in data.items() + } + engine = CypherEngine(config, datasets, use_csr=bool(original.csr_stats())) + query = ( + "MATCH path = (a:Person {id: $start})-[:KNOWS]->(b:Person) " + "RETURN nodes(path) AS ids ORDER BY ids" + ) + old = engine.execute_batches(query, {"start": 20}) + lance.write_dataset( + pa.table({"src": [20], "dst": [30]}), paths["KNOWS"], mode="append" + ) + engine.refresh() + assert engine.execute(query, {"start": 20}).column(0).to_pylist() == [ + [20, 30], + [20, 40], + ] + assert old.read_all().column(0).to_pylist() == [[20, 40]] + assert engine.execute(query, {"start": 10}).column(0).to_pylist() == [ + [10, 20], + [10, 20], + [10, 30], + ] + + +def test_lambda_types_shadow_and_restore_outer_bindings(graph): + engine, _, _ = graph + for initial, expected in [(0, 3), (0.5, 3.5), (None, None), (0, 3)]: + rows = engine.execute( + "WITH 'outer' AS total RETURN total AS outer, " + "reduce(total = $initial, item IN [1, 2] | total + item) AS value", + {"initial": initial}, + ).to_pylist() + assert rows == [{"outer": "outer", "value": expected}] + assert engine.execute( + "WITH 'outer' AS item RETURN all(item IN [1, 2] WHERE item > 0) AS valid, " + "item AS outer" + ).to_pylist() == [{"valid": True, "outer": "outer"}] + with pytest.raises(ValueError, match="Type mismatch"): + engine.execute( + "RETURN reduce(total = 0, item IN ['text'] | total + item) AS value" + ) From 97ce7d66ca079aa16285ed894541cf6a22512fe7 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:27:15 -0700 Subject: [PATCH 30/66] fix(cypher): read every rich path source partition Preserve inventory source coalescing and refresh output partitioning after optimizer rewrites. Cover Arrow and Lance scans beyond the first source partition and the unmodified IC14 query over real relationship records. --- crates/lance-graph/src/path_match.rs | 38 +++++++++++- docs/cypher.md | 4 +- python/python/tests/queries/NOTICE.txt | 13 +++++ python/python/tests/queries/README.md | 11 ++++ python/python/tests/queries/ic14.cypher | 26 +++++++++ python/python/tests/test_path_values.py | 78 +++++++++++++++++++++++++ 6 files changed, 165 insertions(+), 5 deletions(-) create mode 100644 python/python/tests/queries/NOTICE.txt create mode 100644 python/python/tests/queries/README.md create mode 100644 python/python/tests/queries/ic14.cypher diff --git a/crates/lance-graph/src/path_match.rs b/crates/lance-graph/src/path_match.rs index d63e606..7d99585 100644 --- a/crates/lance-graph/src/path_match.rs +++ b/crates/lance-graph/src/path_match.rs @@ -63,7 +63,7 @@ pub(crate) struct PathSpec { pub shortest: Option, } -#[derive(Debug, Clone)] +#[derive(Clone)] pub(crate) struct PathMatchNode { inputs: Vec, spec: Arc, @@ -78,6 +78,16 @@ enum Output { Relationship { step: usize, property: usize }, } +// DataFusion formats extension nodes during planning; never print cached graph batches. +impl fmt::Debug for PathMatchNode { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("PathMatchNode") + .field("variable", &self.spec.variable) + .field("sources", &(self.inputs.len() - 1)) + .finish() + } +} + impl PathMatchNode { pub fn try_new(inputs: Vec, spec: PathSpec) -> Result { let node_fields = property_fields(spec.nodes.iter().map(|s| s.schema.clone()))?; @@ -207,7 +217,6 @@ struct Ready { shortest: Option<(Neighbors, Neighbors)>, } -#[derive(Debug)] struct PathMatchExec { inputs: Vec>, spec: Arc, @@ -276,6 +285,15 @@ pub(crate) fn plan_path_match( })) } +impl fmt::Debug for PathMatchExec { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("PathMatchExec") + .field("variable", &self.spec.variable) + .field("sources", &(self.inputs.len() - 1)) + .finish() + } +} + impl DisplayAs for PathMatchExec { fn fmt_as(&self, _: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "PathMatchExec: {} edge inventory", self.spec.variable) @@ -291,16 +309,30 @@ impl ExecutionPlan for PathMatchExec { fn children(&self) -> Vec<&Arc> { self.inputs.iter().collect() } + // Inventory builds consume one stream per source; keep the coalescers through optimization. + fn required_input_distribution(&self) -> Vec { + use datafusion::physical_plan::Distribution; + std::iter::once(Distribution::UnspecifiedDistribution) + .chain((1..self.inputs.len()).map(|_| Distribution::SinglePartition)) + .collect() + } + fn with_new_children( self: Arc, inputs: Vec>, ) -> Result> { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(self.schema.clone()), + inputs[0].output_partitioning().clone(), + EmissionType::Incremental, + Boundedness::Bounded, + )); Ok(Arc::new(Self { inputs, spec: self.spec.clone(), outputs: self.outputs.clone(), schema: self.schema.clone(), - properties: self.properties.clone(), + properties, predicates: self.predicates.clone(), ready: self.ready.clone(), })) diff --git a/docs/cypher.md b/docs/cypher.md index 886d0b0..35a99f6 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -197,8 +197,8 @@ implemented; `ExecutionStrategy.LanceNative` is reserved and raises an error. The Graph Data Council (GDC), formerly known as the Linked Data Benchmark Council (LDBC), maintains the [GDC SNB Interactive](https://ldbcouncil.org/benchmarks/snb/) complex queries. lance-csr implements a large -read-only subset of openCypher. The tables below specify its supported surface; relationship-path -operations needed by the unchanged IC14 query are not yet supported. +read-only subset of openCypher, including named path values and the relationship-path operations +used by IC14. The tables below specify its supported surface. ### Clauses diff --git a/python/python/tests/queries/NOTICE.txt b/python/python/tests/queries/NOTICE.txt new file mode 100644 index 0000000..69d3211 --- /dev/null +++ b/python/python/tests/queries/NOTICE.txt @@ -0,0 +1,13 @@ + Copyright [2020-]2025 Linked Data Benchmark Council + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. \ No newline at end of file diff --git a/python/python/tests/queries/README.md b/python/python/tests/queries/README.md new file mode 100644 index 0000000..280f6db --- /dev/null +++ b/python/python/tests/queries/README.md @@ -0,0 +1,11 @@ +# Query fixtures + +`ic14.cypher` is the unchanged query string from `run_query14` in +[graph-benchmark-ldbc-complex](https://github.com/hyperdimensionalcomputingai/graph-benchmark-ldbc-complex/blob/c185123edd58b37b417bc4deb2fec643fdbd5e61/neo4j/query.py). +The fixture covers collected path values, relationship endpoints, correlated pattern +comprehensions and ordered folds using a small generated graph, including parallel edges. + +The source adapts the Apache-2.0 LDBC SNB Interactive reference implementation. +The upstream copyright and attribution notice is retained in `NOTICE.txt`; the +Apache-2.0 license is at the repository root. No benchmark dataset or expected +SF1 answers are distributed with this fixture. diff --git a/python/python/tests/queries/ic14.cypher b/python/python/tests/queries/ic14.cypher new file mode 100644 index 0000000..684741c --- /dev/null +++ b/python/python/tests/queries/ic14.cypher @@ -0,0 +1,26 @@ + + MATCH path = allShortestPaths((person1:Person {ID: $person1Id})-[:knows*0..]-(person2:Person {ID: $person2Id})) + WITH collect(path) AS paths + UNWIND paths AS path + WITH path, relationships(path) AS rels_in_path + WITH + [n IN nodes(path) | n.ID] AS personIdsInPath, + [r IN rels_in_path | + reduce(w = 0.0, v IN [ + (a:Person)<-[:commentHasCreator]-(:Comment)-[:replyOfPost]->(:Post)-[:postHasCreator]->(b:Person) + WHERE (a.ID = startNode(r).ID AND b.ID = endNode(r).ID) OR (a.ID = endNode(r).ID AND b.ID = startNode(r).ID) + | 1.0] | w + v) + ] AS weight1, + [r IN rels_in_path | + reduce(w = 0.0, v IN [ + (a:Person)<-[:commentHasCreator]-(:Comment)-[:replyOfComment]->(:Comment)-[:commentHasCreator]->(b:Person) + WHERE (a.ID = startNode(r).ID AND b.ID = endNode(r).ID) OR (a.ID = endNode(r).ID AND b.ID = startNode(r).ID) + | 0.5] | w + v) + ] AS weight2 + WITH + personIdsInPath, + reduce(w = 0.0, v IN weight1 | w + v) AS w1, + reduce(w = 0.0, v IN weight2 | w + v) AS w2 + RETURN personIdsInPath, (w1 + w2) AS pathWeight + ORDER BY pathWeight DESC + \ No newline at end of file diff --git a/python/python/tests/test_path_values.py b/python/python/tests/test_path_values.py index fb90c00..f5f6c63 100644 --- a/python/python/tests/test_path_values.py +++ b/python/python/tests/test_path_values.py @@ -348,3 +348,81 @@ def test_multiple_path_operators_share_occurrence_identity(engine): identity = edge["__lance_graph_identity"] assert ids.setdefault(weight, identity) == identity assert ids[5] != ids[7] + + +@pytest.mark.parametrize("use_csr", [False, True]) +def test_unmodified_ic14_with_real_relationship_values(use_csr): + from pathlib import Path + + # Exact query text from the benchmark's Neo4j IC14, including collect(path), + # relationship endpoints, nested correlated pattern comprehensions, and reduce. + query = (Path(__file__).parent / "queries" / "ic14.cypher").read_text() + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .node_table("Comment", "id") + .node_table("Post", "id") + .rel_table("knows", "Person", "Person", "src", "dst") + .rel_table("commentHasCreator", "Comment", "Person", "src", "dst") + .rel_table("postHasCreator", "Post", "Person", "src", "dst") + .rel_table("replyOfPost", "Comment", "Post", "src", "dst") + .rel_table("replyOfComment", "Comment", "Comment", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table({"id": [1, 2, 3, 4]}), + "Comment": pa.table({"id": [10, 11, 12, 13]}), + "Post": pa.table({"id": [20, 21]}), + "knows": pa.table({"src": [1, 1, 2, 1, 4], "dst": [2, 2, 3, 4, 3]}), + "commentHasCreator": pa.table( + {"src": [10, 11, 12, 13], "dst": [1, 2, 2, 4]} + ), + "postHasCreator": pa.table({"src": [20, 21], "dst": [2, 3]}), + "replyOfPost": pa.table({"src": [10, 11, 13], "dst": [20, 21, 21]}), + "replyOfComment": pa.table({"src": [12], "dst": [10]}), + }, + use_csr=use_csr, + ) + rows = engine.execute(query, {"person1Id": 1, "person2Id": 3}).to_pylist() + assert [(r["personidsinpath"], r["pathweight"]) for r in rows] == [ + ([1, 2, 3], 2.5), + ([1, 2, 3], 2.5), + ([1, 4, 3], 1.0), + ] + + +@pytest.mark.parametrize("storage", ["arrow", "lance"]) +@pytest.mark.parametrize("shortest", [False, True]) +def test_rich_inventory_reads_every_partition(tmp_path, storage, shortest): + import lance + + n = 20_000 + people = pa.table({"id": pa.array(range(n + 1), pa.int64())}) + edges = pa.table( + { + "src": pa.array(range(n), pa.int64()), + "dst": pa.array(range(1, n + 1), pa.int64()), + } + ) + datasets = {"Person": people, "R": edges} + if storage == "lance": + datasets = { + name: lance.write_dataset(table, tmp_path / f"{name}.lance") + for name, table in datasets.items() + } + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + ) + engine = CypherEngine(config, datasets, use_csr=True) + pattern = f"(a:Person {{id: {n - 10}}})-[:R*1..5]->(b:Person {{id: {n - 5}}})" + if shortest: + pattern = f"allShortestPaths({pattern})" + rows = engine.execute(f"MATCH p = {pattern} RETURN p").to_pylist() + assert len(rows) == 1 + assert rows[0]["p"]["nodes"] == list(range(n - 10, n - 4)) + assert len(rows[0]["p"]["relationships"]) == 5 From 33da59d3cf695598075ceb0e7e19f72bf03b3c7a Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:33:56 -0700 Subject: [PATCH 31/66] fix(query): yield while CALL consumes empty batches Bound every Apply loop iteration so empty parent or child batches cannot defer cancellation indefinitely. Exercise the yield with a real MemTable and retain path alias imports across CALL. --- crates/lance-graph/src/apply.rs | 54 +++++++++++++++++---- python/python/tests/test_call_subqueries.py | 20 +++++++- 2 files changed, 63 insertions(+), 11 deletions(-) diff --git a/crates/lance-graph/src/apply.rs b/crates/lance-graph/src/apply.rs index 50dfe32..f454e5d 100644 --- a/crates/lance-graph/src/apply.rs +++ b/crates/lance-graph/src/apply.rs @@ -185,7 +185,7 @@ impl ExecutionPlan for ApplyExec { input, batch: None, next_row: 0, - rows_since_yield: 0, + iterations_since_yield: 0, current: vec![], child: None, inner: self.inner.clone(), @@ -196,6 +196,12 @@ impl ExecutionPlan for ApplyExec { }; let stream = futures::stream::try_unfold(state, |mut state| async move { loop { + // Bound all ready work, including empty outer and child batches. + if state.iterations_since_yield == 64 { + crate::csr_expand::cooperative_yield().await; + state.iterations_since_yield = 0; + } + state.iterations_since_yield += 1; if let Some(child) = &mut state.child { if let Some(batch) = child.next().await { let batch = batch?; @@ -231,13 +237,6 @@ impl ExecutionPlan for ApplyExec { continue; } } - // Empty children can complete synchronously for many outer rows. - // Bound that work so cancellation and other tasks can be polled. - if state.rows_since_yield == 64 { - crate::csr_expand::cooperative_yield().await; - state.rows_since_yield = 0; - } - state.rows_since_yield += 1; let batch = state.batch.as_ref().expect("a nonempty outer batch"); state.current = batch .columns() @@ -266,7 +265,7 @@ struct ApplyStreamState { input: SendableRecordBatchStream, batch: Option, next_row: usize, - rows_since_yield: usize, + iterations_since_yield: usize, current: Vec, child: Option, inner: LogicalPlan, @@ -347,6 +346,43 @@ mod tests { .to_vec() } + #[tokio::test] + async fn call_empty_memory_batches_yield_before_reaching_the_result() { + use datafusion::datasource::{MemTable, TableProvider}; + use datafusion::execution::context::SessionContext; + use datafusion::logical_expr::{lit, LogicalPlanBuilder}; + use futures::TryStreamExt; + use std::task::{Context, Poll}; + + let batch = RecordBatch::try_from_iter(vec![( + "x", + Arc::new(Int64Array::from(vec![7])) as arrow_array::ArrayRef, + )]) + .unwrap(); + let mut batches = vec![RecordBatch::new_empty(batch.schema()); 10_000]; + batches.push(batch.clone()); + let table = MemTable::try_new(batch.schema(), vec![batches]).unwrap(); + let state = SessionContext::new().state(); + let input = table.scan(&state, None, &[], None).await.unwrap(); + let inner = LogicalPlanBuilder::empty(true) + .project(vec![lit(1_i64).alias("n")]) + .unwrap() + .build() + .unwrap(); + let schema = Arc::new(arrow_schema::Schema::new(vec![ + batch.schema().field(0).clone(), + inner.schema().as_arrow().field(0).clone(), + ])); + let plan = ApplyExec::new(input, inner, vec![], Arc::new(state.clone()), schema); + let mut stream = plan.execute(0, state.task_ctx()).unwrap(); + let mut cx = Context::from_waker(futures::task::noop_waker_ref()); + assert!(matches!(stream.poll_next_unpin(&mut cx), Poll::Pending)); + let batches: Vec<_> = stream.try_collect().await.unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!(ints(&batches[0], "x"), [7]); + assert_eq!(ints(&batches[0], "n"), [1]); + } + #[tokio::test] async fn call_per_row_top_k_and_entity_export() { let result = run("MATCH (p:Person) CALL { WITH p MATCH (p)-[:KNOWS]->(f:Person) RETURN f AS friend ORDER BY friend.age DESC LIMIT 1 } RETURN p.id AS id, friend.id AS friend ORDER BY id").await.unwrap(); diff --git a/python/python/tests/test_call_subqueries.py b/python/python/tests/test_call_subqueries.py index b01679a..b2d08f8 100644 --- a/python/python/tests/test_call_subqueries.py +++ b/python/python/tests/test_call_subqueries.py @@ -332,7 +332,8 @@ def test_call_imported_shortest_path_keeps_relationship_records(graph): "MATCH path = allShortestPaths((a:Person {id: 10})-[:KNOWS*..3]->" "(b:Person {id: 40})) CALL { WITH path " "RETURN [edge IN relationships(path) | endNode(edge).id] AS targets, " - "reduce(total = 0, edge IN relationships(path) | total + endNode(edge).id) AS total } " + "reduce(total = 0, edge IN relationships(path) | " + "total + endNode(edge).id) AS total } " "RETURN targets, total ORDER BY targets" ) assert engine.execute(query).to_pylist() == [ @@ -374,7 +375,8 @@ def test_call_empty_children_respect_query_deadline(graph): engine, _, _ = graph with pytest.raises(TimeoutError): engine.execute( - "UNWIND range(1, 1000000) AS x CALL { WITH x UNWIND [] AS n RETURN n } RETURN n", + "UNWIND range(1, 1000000) AS x " + "CALL { WITH x UNWIND [] AS n RETURN n } RETURN n", timeout=0.02, ) assert engine.execute("RETURN 1 AS ok").to_pylist() == [{"ok": 1}] @@ -441,3 +443,17 @@ def test_lambda_types_shadow_and_restore_outer_bindings(graph): engine.execute( "RETURN reduce(total = 0, item IN ['text'] | total + item) AS value" ) + + +def test_shortest_path_alias_import_demands_relationship_values(graph): + engine, _, _ = graph + query = ( + "MATCH path = allShortestPaths((a:Person {id: 10})-[:KNOWS*..3]->" + "(b:Person {id: 40})) WITH path AS carried " + "CALL { WITH carried RETURN relationships(carried) AS edges } " + "RETURN [edge IN edges | endNode(edge).id] AS targets ORDER BY targets" + ) + assert engine.execute(query).to_pylist() == [ + {"targets": [20, 40]}, + {"targets": [20, 40]}, + ] From 65ae2a7cd2b4da25483409ffe8d51c2e771fba1a Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:32:45 -0700 Subject: [PATCH 32/66] perf(cypher): keep compact paths through WITH aliases Materialize path records only when their consumers need values or relationships. Track shortest-path aliases before checking demand and retain node labels for homogeneous property comprehensions. --- crates/lance-graph/src/logical_plan.rs | 3 ++ crates/lance-graph/src/path_value.rs | 34 +++++++++++++++++++-- docs/cypher.md | 4 +-- python/python/tests/test_path_values.py | 40 +++++++++++++++++++++++++ 4 files changed, 76 insertions(+), 5 deletions(-) diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index d77ed26..c031014 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -1937,6 +1937,9 @@ impl<'a> LogicalPlanner<'a> { if let Some(label) = self.variables.get(var).cloned() { self.variables.insert(alias.clone(), label); } + if let Some(label) = self.path_labels.get(var).cloned() { + self.path_labels.insert(alias.clone(), label); + } Some(alias.clone()) } (Some(alias), _) => Some(alias.clone()), diff --git a/crates/lance-graph/src/path_value.rs b/crates/lance-graph/src/path_value.rs index c4d5d14..a41ca21 100644 --- a/crates/lance-graph/src/path_value.rs +++ b/crates/lance-graph/src/path_value.rs @@ -123,8 +123,8 @@ pub(crate) fn paths_array( ))) } -/// A direct length/nodes query can keep its existing CSR plan. Values carried through -/// containers and node-property comprehensions need the actual path records. +/// Length/nodes queries, including WITH pass-through aliases, keep their CSR plan. +/// Returned path values, relationship access and containers need actual path records. pub(crate) fn needs_values(query: &crate::ast::CypherQuery) -> bool { use crate::ast::{GraphPattern, ReadingClause, ValueExpression as V}; let clauses: Vec<_> = query @@ -153,6 +153,20 @@ pub(crate) fn needs_values(query: &crate::ast::CypherQuery) -> bool { } } } + // A bare WITH only renames/carries a path; inspect its eventual consumers instead. + for with in query + .with_clause + .iter() + .chain(query.further_with_stages.iter().map(|s| &s.with_clause)) + { + for item in &with.items { + if let (V::Variable(name), Some(alias)) = (&item.expression, &item.alias) { + if names.contains(&name.to_lowercase()) { + names.insert(alias.to_lowercase()); + } + } + } + } let mut needed = query .union_branches .iter() @@ -216,7 +230,10 @@ pub(crate) fn needs_values(query: &crate::ast::CypherQuery) -> bool { .chain(query.further_with_stages.iter().map(|s| &s.with_clause)) { for item in &with.items { - inspect(&item.expression); + if !matches!(&item.expression, V::Variable(name) if names.contains(&name.to_lowercase())) + { + inspect(&item.expression); + } } if let Some(order) = &with.order_by { for item in &order.items { @@ -261,6 +278,17 @@ mod tests { for (suffix, rich) in [ ("RETURN length(p), nodes(p)", false), ("RETURN [n IN nodes(p) | n.id]", false), + ("WITH p RETURN nodes(p)", false), + ( + "WITH p AS q WITH q AS r RETURN [n IN nodes(r) | n.id]", + false, + ), + ("WITH p AS q RETURN q", true), + ("WITH p AS q RETURN relationships(q)", true), + ( + "WITH p AS q WITH collect(q) AS ps UNWIND ps AS r RETURN nodes(r)", + true, + ), ("RETURN p", true), ("RETURN relationships(p)", true), ("RETURN length(head(collect(p)))", true), diff --git a/docs/cypher.md b/docs/cypher.md index 35a99f6..80f8fe1 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -359,8 +359,8 @@ properties. It therefore materializes those tables, even with `use_csr=True`. Tr in batches and named ordinary paths keep a cursor per hop, but inventory memory grows with the graph. Rich path queries bypass the query-template cache: each execution binds its current parameters and creates an inventory from its own snapshot. A reader opened before `refresh()` retains the old snapshot. -Queries asking only for shortest-path length or node keys, including existing homogeneous node-property -comprehensions, retain the usual CSR execution route. +Queries asking only for shortest-path length or node keys, including homogeneous node-property +comprehensions and paths carried through `WITH` aliases, retain the usual CSR execution route. ### Expressions and predicates diff --git a/python/python/tests/test_path_values.py b/python/python/tests/test_path_values.py index f5f6c63..82623d7 100644 --- a/python/python/tests/test_path_values.py +++ b/python/python/tests/test_path_values.py @@ -426,3 +426,43 @@ def test_rich_inventory_reads_every_partition(tmp_path, storage, shortest): assert len(rows) == 1 assert rows[0]["p"]["nodes"] == list(range(n - 10, n - 4)) assert len(rows[0]["p"]["relationships"]) == 5 + + +@pytest.mark.parametrize( + ("suffix", "rich"), + [ + ("WITH p RETURN nodes(p) AS ids", False), + ("WITH p AS q WITH q AS r RETURN nodes(r) AS ids", False), + ("WITH p AS q RETURN [n IN nodes(q) | n.id] AS ids", False), + ( + "WITH p AS q WITH collect(q) AS ps UNWIND ps AS r RETURN nodes(r) AS ids", + True, + ), + ], +) +def test_path_pass_through_only_materializes_consumed_values(engine, suffix, rich): + query = ( + "MATCH p = allShortestPaths((a:Person {id: 1})-[:KNOWS*]->(b:Person {id: 3})) " + + suffix + ) + assert engine.execute(query).to_pylist() == [{"ids": [1, 2, 3]}] * 2 + assert ("PathMatchExec" in engine.explain(query)) is rich + + +@pytest.mark.parametrize( + "suffix", ["RETURN relationships(q) AS edges", "RETURN q AS path"] +) +def test_path_alias_retains_actual_relationship_values(engine, suffix): + query = ( + "MATCH p = allShortestPaths((a:Person {id: 1})-[:KNOWS*]->(b:Person {id: 3})) " + "WITH p AS q " + suffix + ) + rows = engine.execute(query).to_pylist() + edges = [row.get("edges", row.get("path", {}).get("relationships")) for row in rows] + assert sorted( + [r["__lance_graph_properties"]["weight"] for r in path] for path in edges + ) == [ + [5, 11], + [7, 11], + ] + assert "PathMatchExec" in engine.explain(query) From 6f1787fffe0a4d6a99d654d1d40aa3b20de66413 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 10:22:52 -0700 Subject: [PATCH 33/66] docs(python): show standalone bridge tests on macOS Resolve libpython from the active environment and execute the exact test artifact returned by Cargo. Preserve extension packaging. Submitted by David Hughes (@stugorf), with AI Collaboration. --- python/DEVELOPMENT.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/python/DEVELOPMENT.md b/python/DEVELOPMENT.md index 87d7401..7f57095 100644 --- a/python/DEVELOPMENT.md +++ b/python/DEVELOPMENT.md @@ -52,6 +52,23 @@ The Rust tests run from the repository root: cargo test --workspace ``` +The PyO3 crate is excluded from that workspace. Its private Rust tests cover engine snapshot and cache +lifecycles. On macOS, `extension-module` leaves Python symbols for the interpreter to provide, so a standalone +test executable needs an explicit link to the active Python library. From `python/`, after dependency setup: + +```shell +lance_python_library=$(uv run --no-sync python -c 'import pathlib, sysconfig; print(pathlib.Path(sysconfig.get_config_var("LIBDIR")) / sysconfig.get_config_var("LDLIBRARY"))') +lance_test_artifacts=$(mktemp) +cargo rustc --manifest-path ../crates/lance-graph-python/Cargo.toml --tests --message-format=json -- -C "link-arg=$lance_python_library" > "$lance_test_artifacts" +lance_test_binary=$(uv run --no-sync python -c 'import json, sys; artifacts = [json.loads(line) for line in open(sys.argv[1])]; print(next(a["executable"] for a in artifacts if a.get("reason") == "compiler-artifact" and a.get("profile", {}).get("test") and a.get("executable")))' "$lance_test_artifacts") +"$lance_test_binary" +rm "$lance_test_artifacts" +``` + +This uses Cargo's emitted executable path rather than selecting an older binary from `target/`. It does not +change extension packaging or replace the Python tests. Use a separate `CARGO_TARGET_DIR` for each worktree +when validating multiple revisions, so each test binary is built from that worktree's sources. + ## Formatting and linting ```shell From 32edb0c113ec1bd58011279ca1c95163f3700283 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:35:41 -0700 Subject: [PATCH 34/66] test(cypher): cover path aliases through CALL and UNION --- crates/lance-graph/src/path_value.rs | 8 ++++++++ python/python/tests/test_call_subqueries.py | 16 ++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/crates/lance-graph/src/path_value.rs b/crates/lance-graph/src/path_value.rs index a41ca21..11e276e 100644 --- a/crates/lance-graph/src/path_value.rs +++ b/crates/lance-graph/src/path_value.rs @@ -285,6 +285,14 @@ mod tests { ), ("WITH p AS q RETURN q", true), ("WITH p AS q RETURN relationships(q)", true), + ( + "WITH p AS q CALL { WITH q RETURN relationships(q) AS edges } RETURN edges", + true, + ), + ( + "WITH p AS q CALL { WITH q RETURN q AS result UNION WITH q RETURN q AS result } RETURN nodes(result)", + true, + ), ( "WITH p AS q WITH collect(q) AS ps UNWIND ps AS r RETURN nodes(r)", true, diff --git a/python/python/tests/test_call_subqueries.py b/python/python/tests/test_call_subqueries.py index b2d08f8..87668e5 100644 --- a/python/python/tests/test_call_subqueries.py +++ b/python/python/tests/test_call_subqueries.py @@ -457,3 +457,19 @@ def test_shortest_path_alias_import_demands_relationship_values(graph): {"targets": [20, 40]}, {"targets": [20, 40]}, ] + + +def test_shortest_path_alias_survives_call_union(graph): + engine, _, _ = graph + query = ( + "MATCH path = allShortestPaths((a:Person {id: 10})-[:KNOWS*..3]->" + "(b:Person {id: 40})) WITH path AS carried " + "CALL { WITH carried RETURN carried AS result " + "UNION WITH carried RETURN carried AS result } " + "RETURN [edge IN relationships(result) | endNode(edge).id] AS targets " + "ORDER BY targets" + ) + assert engine.execute(query).to_pylist() == [ + {"targets": [20, 40]}, + {"targets": [20, 40]}, + ] From 68fdcb40289d29cfa6bee27bdfa20aad7df943e5 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:35:24 -0700 Subject: [PATCH 35/66] fix(cypher): reject ambiguous rich path endpoint IDs --- crates/lance-graph/src/path_inventory.rs | 39 ++++++++++++++++++++++-- docs/cypher.md | 4 ++- python/python/tests/test_path_values.py | 26 ++++++++++++++++ 3 files changed, 66 insertions(+), 3 deletions(-) diff --git a/crates/lance-graph/src/path_inventory.rs b/crates/lance-graph/src/path_inventory.rs index 1741c22..713a1db 100644 --- a/crates/lance-graph/src/path_inventory.rs +++ b/crates/lance-graph/src/path_inventory.rs @@ -164,8 +164,12 @@ impl Inventory { "path inventory exceeds u32 row references".into(), ) })?; - if keys.is_valid(row) { - by_key.insert((label, keys.value(row)), dense); + if keys.is_valid(row) && by_key.insert((label, keys.value(row)), dense).is_some() { + return exec_err!( + "path values require unique node IDs within label '{}'; duplicate ID {}", + source.label, + keys.value(row) + ); } node_labels.push(label); labels.push(source.label.as_str()); @@ -280,3 +284,34 @@ impl Inventory { }) } } + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::prelude::SessionContext; + + #[tokio::test] + async fn duplicate_node_keys_do_not_choose_an_arbitrary_endpoint() { + let batch = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from(vec![1, 1])) as ArrayRef, + )]) + .unwrap(); + let source = NodeSource { + label: "Person".into(), + id: "id".into(), + schema: batch.schema(), + }; + let ctx = SessionContext::new(); + let plan = ctx + .read_batch(batch) + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let error = Inventory::build(&[source], &[], &[plan], ctx.task_ctx(), &Default::default()) + .await + .unwrap_err(); + assert!(error.to_string().contains("duplicate ID 1")); + } +} diff --git a/docs/cypher.md b/docs/cypher.md index 80f8fe1..cb4fdba 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -352,7 +352,9 @@ same name. `keys(entity)` lists only non-null user properties; `keys(map)` inclu Prefer the Cypher property and endpoint functions when writing queries. Relationship identities identify actual occurrences within this query's inventory, not persistent IDs that can be compared across queries or refreshes. Rich paths support the same signed integer and `uint8`/`uint16`/`uint32` keys as the CSR. Mixed labels may share compatible properties; conflicting Arrow -property types fail explicitly. Expanding a graph pattern from a mixed-label value remains unsupported. +property types fail explicitly. Non-null node IDs must be unique within each participating label so +relationship endpoints identify one node; duplicate IDs raise an error. Expanding a graph pattern from +a mixed-label value remains unsupported. A path value needs a query-local inventory of the participating node and relationship tables, including properties. It therefore materializes those tables, even with `use_csr=True`. Traversal output is emitted diff --git a/python/python/tests/test_path_values.py b/python/python/tests/test_path_values.py index 82623d7..e19cc62 100644 --- a/python/python/tests/test_path_values.py +++ b/python/python/tests/test_path_values.py @@ -466,3 +466,29 @@ def test_path_alias_retains_actual_relationship_values(engine, suffix): [7, 11], ] assert "PathMatchExec" in engine.explain(query) + + +@pytest.mark.parametrize("storage", ["arrow", "lance"]) +def test_rich_path_rejects_ambiguous_node_ids(tmp_path, storage): + import lance + + tables = { + "Person": pa.table({"id": [1, 1, 2], "name": ["one", "duplicate", "two"]}), + "R": pa.table({"src": [1], "dst": [2]}), + } + if storage == "lance": + tables = { + name: lance.write_dataset(table, tmp_path / f"{name}.lance") + for name, table in tables.items() + } + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .build() + ) + engine = CypherEngine(config, tables, use_csr=True) + with pytest.raises( + Exception, match="unique node IDs within label 'Person'; duplicate ID 1" + ): + engine.execute("MATCH p = (a:Person)-[:R]->(b:Person) RETURN p") From f4f5bfa04f49ce962a95d9e9d5e2900273f9ae71 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:36:54 -0700 Subject: [PATCH 36/66] fix(python): release timed streams after execution errors Drop the owned DataFusion stream when its first error finishes the Arrow reader, preserving the original error. A real CAST failure and Weak reference to a later Arrow batch fail before the fix and pass after it. All 15 native binding tests pass; the Python reader-error regression and complete native rebuild are queued. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph-python/src/graph.rs | 45 +++++++++++++++++++++++++- python/python/tests/test_timeout.py | 10 ++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 17dcd95..6f7323f 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -2007,7 +2007,7 @@ impl Iterator for DeadlineBatchReader { state.expired = true; return Some(Err(stream_timeout_error())); } - if result.is_none() { + if result.as_ref().is_none_or(|result| result.is_err()) { state.stream.take(); } result.map(|r| { @@ -2922,6 +2922,49 @@ mod tests { assert!(retained.upgrade().is_none()); } + #[test] + fn timed_reader_releases_stream_after_an_execution_error() { + let (stream, retained) = RT.runtime.block_on(async { + let values = Arc::new(arrow_array::StringArray::from(vec!["2"])); + let retained = Arc::downgrade(&values); + let schema = Arc::new(Schema::new(vec![arrow_schema::Field::new( + "value", + arrow_schema::DataType::Utf8, + false, + )])); + let bad = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(arrow_array::StringArray::from(vec!["invalid"]))], + ) + .unwrap(); + let remaining = RecordBatch::try_new(schema.clone(), vec![values]).unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "input", + Arc::new(MemTable::try_new(schema, vec![vec![bad, remaining]]).unwrap()), + ) + .unwrap(); + let stream = ctx + .sql("SELECT CAST(value AS BIGINT) AS value FROM input") + .await + .unwrap() + .execute_stream() + .await + .unwrap(); + (stream, retained) + }); + let mut reader = reader_with_deadline(stream, QueryDeadline::new(Some(10.0)).unwrap()); + assert!(retained.upgrade().is_some()); + let error = reader.next().unwrap().unwrap_err().to_string(); + assert!(error.contains("invalid"), "{error}"); + assert!(!error.contains(QUERY_TIMEOUT_MESSAGE), "{error}"); + assert!( + retained.upgrade().is_none(), + "error must release the owned stream" + ); + assert!(reader.next().is_none()); + } + #[test] fn test_split_vector_column_with_alias() { let (alias, column) = split_vector_column("d.embedding"); diff --git a/python/python/tests/test_timeout.py b/python/python/tests/test_timeout.py index 87db8f9..9717815 100644 --- a/python/python/tests/test_timeout.py +++ b/python/python/tests/test_timeout.py @@ -70,6 +70,16 @@ def test_no_deadline_and_normal_execution_errors(engine): engine.execute("RETURN unknown", timeout=10) +def test_execution_error_in_a_timed_stream_is_preserved(engine): + reader = engine.execute_batches("MATCH (p:Person) RETURN 1 / p.id AS n", timeout=10) + with pytest.raises(pa.ArrowInvalid) as raised: + reader.read_next_batch() + assert "Query timeout exceeded" not in str(raised.value) + with pytest.raises(StopIteration): + reader.read_next_batch() + assert engine.execute("RETURN 3 AS n").to_pylist() == [{"n": 3}] + + def test_expensive_query_stops_and_engine_stays_usable(engine): started = time.monotonic() with pytest.raises(TimeoutError, match="Query timeout exceeded"): From 5f4d13901a6464b6f8863209e5a40240dcc90058 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:44:58 -0700 Subject: [PATCH 37/66] test(query): isolate CALL cooperative scheduling bound Use 65 real empty MemTable batches with Tokio inherited budgeting disabled for the first poll. The regression fails against the old per-row checkpoint and passes with the loop-wide checkpoint. --- crates/lance-graph/src/apply.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/crates/lance-graph/src/apply.rs b/crates/lance-graph/src/apply.rs index f454e5d..719ccf9 100644 --- a/crates/lance-graph/src/apply.rs +++ b/crates/lance-graph/src/apply.rs @@ -347,11 +347,11 @@ mod tests { } #[tokio::test] - async fn call_empty_memory_batches_yield_before_reaching_the_result() { + async fn call_empty_memory_batches_yield_before_end_of_stream() { use datafusion::datasource::{MemTable, TableProvider}; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::{lit, LogicalPlanBuilder}; - use futures::TryStreamExt; + use futures::{Future, TryStreamExt}; use std::task::{Context, Poll}; let batch = RecordBatch::try_from_iter(vec![( @@ -359,8 +359,7 @@ mod tests { Arc::new(Int64Array::from(vec![7])) as arrow_array::ArrayRef, )]) .unwrap(); - let mut batches = vec![RecordBatch::new_empty(batch.schema()); 10_000]; - batches.push(batch.clone()); + let batches = vec![RecordBatch::new_empty(batch.schema()); 65]; let table = MemTable::try_new(batch.schema(), vec![batches]).unwrap(); let state = SessionContext::new().state(); let input = table.scan(&state, None, &[], None).await.unwrap(); @@ -376,11 +375,12 @@ mod tests { let plan = ApplyExec::new(input, inner, vec![], Arc::new(state.clone()), schema); let mut stream = plan.execute(0, state.task_ctx()).unwrap(); let mut cx = Context::from_waker(futures::task::noop_waker_ref()); - assert!(matches!(stream.poll_next_unpin(&mut cx), Poll::Pending)); + // Isolate Apply's own bound from Tokio's task-wide scheduling budget. + let mut next = Box::pin(tokio::task::unconstrained(stream.next())); + assert!(matches!(next.as_mut().poll(&mut cx), Poll::Pending)); + drop(next); let batches: Vec<_> = stream.try_collect().await.unwrap(); - assert_eq!(batches.len(), 1); - assert_eq!(ints(&batches[0], "x"), [7]); - assert_eq!(ints(&batches[0], "n"), [1]); + assert!(batches.is_empty()); } #[tokio::test] From c3c799c873e30c761f3ca5aa96b4f008c665a499 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:42:15 -0700 Subject: [PATCH 38/66] refactor(cypher): reuse strict DataFusion math kernels Apply the numeric submission split to the integrated feature stack: route numeric functions through the strict math wrapper, reuse abs/floor/ceil/sqrt kernels, and remove the duplicate custom implementation. Keep Cypher rounding ties and explicitly reject out-of-range unsigned conversion. The identical wrapper passed 343 Rust library tests, clippy, 96 focused Python tests, and 548 total Python tests on the separate math branch. This integration follow-up has formatting and whitespace checks; integrated runtime validation remains pending. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/Cargo.toml | 1 + .../src/datafusion_planner/expression.rs | 12 +- .../src/datafusion_planner/math_functions.rs | 164 ++++++++++++++++++ .../lance-graph/src/datafusion_planner/mod.rs | 1 + .../datafusion_planner/scalar_functions.rs | 60 +------ docs/cypher.md | 2 +- python/python/tests/test_math_functions.py | 132 ++++++++++++++ 7 files changed, 308 insertions(+), 64 deletions(-) create mode 100644 crates/lance-graph/src/datafusion_planner/math_functions.rs create mode 100644 python/python/tests/test_math_functions.py diff --git a/crates/lance-graph/Cargo.toml b/crates/lance-graph/Cargo.toml index b054f60..aa40f55 100644 --- a/crates/lance-graph/Cargo.toml +++ b/crates/lance-graph/Cargo.toml @@ -18,6 +18,7 @@ arrow-schema = "58" async-trait = "0.1" datafusion = { version = "54", default-features = false, features = [ "nested_expressions", + "math_expressions", "regex_expressions", "unicode_expressions", "crypto_expressions", diff --git a/crates/lance-graph/src/datafusion_planner/expression.rs b/crates/lance-graph/src/datafusion_planner/expression.rs index baddb70..7eba94a 100644 --- a/crates/lance-graph/src/datafusion_planner/expression.rs +++ b/crates/lance-graph/src/datafusion_planner/expression.rs @@ -275,11 +275,13 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { VE::Variable(path) => col(qualify_column(path, &name.to_lowercase())), _ => Expr::Literal(datafusion::scalar::ScalarValue::Null, None), }, - "abs" | "round" | "floor" | "ceil" | "sqrt" | "left" | "right" | "substring" - | "replace" | "split" | "trim" | "ltrim" | "rtrim" | "size" | "reverse" - | "last" | "tail" | "keys" | "__entity_keys" | "labels" | "startnode" - | "endnode" | "__predicate_all" | "__predicate_any" | "__predicate_none" - | "__predicate_single" => { + "abs" | "round" | "floor" | "ceil" | "sqrt" => { + super::math_functions::call(name, args.iter().map(to_df_value_expr).collect()) + } + "left" | "right" | "substring" | "replace" | "split" | "trim" | "ltrim" + | "rtrim" | "size" | "reverse" | "last" | "tail" | "keys" | "__entity_keys" + | "labels" | "startnode" | "endnode" | "__predicate_all" | "__predicate_any" + | "__predicate_none" | "__predicate_single" => { super::scalar_functions::call(name, args.iter().map(to_df_value_expr).collect()) } "range" => { diff --git a/crates/lance-graph/src/datafusion_planner/math_functions.rs b/crates/lance-graph/src/datafusion_planner/math_functions.rs new file mode 100644 index 0000000..e82f427 --- /dev/null +++ b/crates/lance-graph/src/datafusion_planner/math_functions.rs @@ -0,0 +1,164 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Strict Cypher numeric arguments and rounding over DataFusion's math kernels. + +use std::sync::Arc; + +use arrow::array::{AsArray, Float64Array}; +use arrow::compute::{cast_with_options, CastOptions}; +use arrow_schema::{DataType, Field}; +use datafusion::common::plan_err; +use datafusion::error::Result; +use datafusion::functions::math; +use datafusion::logical_expr::{ + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; + +#[derive(Debug, PartialEq, Eq, Hash)] +struct CypherMath { + name: String, + signature: Signature, +} + +impl ScalarUDFImpl for CypherMath { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, types: &[DataType]) -> Result { + let input = &types[0]; + if !input.is_integer() && !input.is_floating() && !input.is_null() { + return plan_err!("{} does not accept {input} for argument 1", self.name); + } + Ok(if self.name == "abs" && input.is_integer() { + DataType::Int64 + } else { + DataType::Float64 + }) + } + + fn invoke_with_args(&self, mut args: ScalarFunctionArgs) -> Result { + let output_type = args.return_type().clone(); + let input = cast_with_options( + &args.args[0].to_array(args.number_rows)?, + &output_type, + &CastOptions { + safe: false, + ..Default::default() + }, + )?; + if self.name == "round" { + let rounded = input + .as_primitive::() + .iter() + .map(|n| { + n.map(|n| { + // Cypher rounds ties toward positive infinity and returns positive zero. + let rounded = if n.fract() == -0.5 { + n.ceil() + } else { + n.round() + }; + if rounded == 0.0 { + 0.0 + } else { + rounded + } + }) + }) + .collect::(); + return Ok(ColumnarValue::Array(Arc::new(rounded))); + } + let udf = match self.name.as_str() { + "abs" => math::abs(), + "floor" => math::floor(), + "ceil" => math::ceil(), + "sqrt" => math::sqrt(), + _ => unreachable!("validated math function"), + }; + args.args = vec![ColumnarValue::Array(input)]; + args.arg_fields = vec![Arc::new(Field::new("arg", output_type, true))]; + udf.invoke_with_args(args) + } +} + +pub(super) fn call(name: &str, args: Vec) -> Expr { + ScalarUDF::new_from_impl(CypherMath { + name: name.to_lowercase(), + signature: Signature::any(1, Volatility::Immutable), + }) + .call(args) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{config::GraphConfig, query::CypherQuery}; + use arrow::array::{Array, ArrayRef, Int64Array, RecordBatch}; + use std::collections::HashMap; + + async fn run(expression: &str) -> crate::error::Result { + let batch = + RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(vec![1])) as ArrayRef)]) + .unwrap(); + CypherQuery::new(&format!("RETURN {expression} AS value"))? + .with_config( + GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(), + ) + .execute(HashMap::from([("Person".to_string(), batch)]), None) + .await + } + + #[tokio::test] + async fn cypher_math_rounding_and_overflow() { + for (expression, expected) in [ + ("abs(-1.5)", 1.5), + ("round(-1.5)", -1.0), + ("round(1.5)", 2.0), + ("round(4503599627370495.5)", 4503599627370496.0), + ("round(4503599627370497.0)", 4503599627370497.0), + ("floor(-1.5)", -2.0), + ("ceil(-1.5)", -1.0), + ("sqrt(9)", 3.0), + ] { + let out = run(expression).await.unwrap(); + assert_eq!( + out.column(0) + .as_primitive::() + .value(0), + expected, + "{expression}" + ); + } + let out = run("abs(-5)").await.unwrap(); + assert_eq!( + out.column(0) + .as_primitive::() + .value(0), + 5 + ); + assert!(run("abs(-9223372036854775807 - 1)").await.is_err()); + } + + #[tokio::test] + async fn cypher_math_validates_types_arity_and_nulls() { + for name in ["abs", "round", "floor", "ceil", "sqrt"] { + let out = run(&format!("{name}(null)")).await.unwrap(); + assert_eq!(out.column(0).null_count(), 1, "{name}"); + for arg in ["'1'", "true", "[1]", "1, 2"] { + assert!( + run(&format!("{name}({arg})")).await.is_err(), + "{name}({arg})" + ); + } + } + } +} diff --git a/crates/lance-graph/src/datafusion_planner/mod.rs b/crates/lance-graph/src/datafusion_planner/mod.rs index c67009d..ee85dc4 100644 --- a/crates/lance-graph/src/datafusion_planner/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/mod.rs @@ -21,6 +21,7 @@ mod cypher_udf; pub(crate) mod expression; mod join_ops; mod list_folds; +mod math_functions; mod scalar_functions; mod scan_ops; mod udf; diff --git a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs index 5ff9818..33bb32e 100644 --- a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs +++ b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs @@ -7,8 +7,7 @@ use std::sync::Arc; use arrow::array::{ - new_null_array, Array, ArrayRef, AsArray, Float64Array, Int64Array, ListBuilder, StringArray, - StringBuilder, + new_null_array, Array, AsArray, Int64Array, ListBuilder, StringArray, StringBuilder, }; use arrow::compute::cast; use arrow_schema::{DataType, Field}; @@ -48,7 +47,6 @@ impl ScalarUDFImpl for CypherScalar { fn return_type(&self, types: &[DataType]) -> Result { for (i, t) in types.iter().enumerate() { let valid = match self.name.as_str() { - "abs" | "round" | "floor" | "ceil" | "sqrt" => t.is_integer() || t.is_floating(), "range" => t.is_integer(), "__predicate_all" | "__predicate_any" | "__predicate_none" | "__predicate_single" => list_field(t) @@ -76,8 +74,6 @@ impl ScalarUDFImpl for CypherScalar { DataType::Boolean } "size" => DataType::Int64, - "abs" if types[0].is_integer() => DataType::Int64, - "abs" | "round" | "floor" | "ceil" | "sqrt" => DataType::Float64, "range" => DataType::List(Arc::new(Field::new("item", DataType::Int64, true))), "split" | "keys" | "__entity_keys" | "labels" => { DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))) @@ -247,58 +243,6 @@ impl ScalarUDFImpl for CypherScalar { Ok(result) }; } - if matches!( - self.name.as_str(), - "abs" | "round" | "floor" | "ceil" | "sqrt" - ) { - let input = cast(&args.args[0].to_array(args.number_rows)?, &output_type)?; - let result: ArrayRef = if output_type == DataType::Int64 { - Arc::new( - input - .as_primitive::() - .iter() - .map(|n| { - n.map(|n| { - n.checked_abs().ok_or_else(|| { - datafusion::common::exec_datafusion_err!("abs overflow") - }) - }) - .transpose() - }) - .collect::>()?, - ) - } else { - Arc::new( - input - .as_primitive::() - .iter() - .map(|n| { - n.map(|n| match self.name.as_str() { - "abs" => n.abs(), - // Cypher rounds ties toward positive infinity and returns positive zero. - "round" => { - let rounded = if n.fract() == -0.5 { - n.ceil() - } else { - n.round() - }; - if rounded == 0.0 { - 0.0 - } else { - rounded - } - } - "floor" => n.floor(), - "ceil" => n.ceil(), - "sqrt" => n.sqrt(), - _ => unreachable!(), - }) - }) - .collect::(), - ) - }; - return Ok(ColumnarValue::Array(result)); - } let arrays = args .args .iter() @@ -421,7 +365,7 @@ pub(super) fn call(name: &str, args: Vec) -> Expr { mod tests { use super::*; use crate::{config::GraphConfig, query::CypherQuery}; - use arrow::array::RecordBatch; + use arrow::array::{ArrayRef, RecordBatch}; use std::collections::HashMap; async fn run(expression: &str) -> crate::error::Result { diff --git a/docs/cypher.md b/docs/cypher.md index cb4fdba..24683bf 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -414,7 +414,7 @@ steps. A step pointing away from the end gives an empty list; zero steps and non raise an error. Null arguments return null. `abs` preserves integer results; the other math functions return floating point values. Integer -absolute-value overflow raises an error. `round` accepts one argument, with ties toward positive +absolute-value overflow and unsigned inputs outside Int64 range raise an error. `round` accepts one argument, with ties toward positive infinity (`round(-1.5)` is `-1.0`), as specified in the [Neo4j 4.4 numeric reference](https://neo4j.com/docs/cypher-manual/4.4/functions/mathematical-numeric/). `sqrt` of a negative number returns NaN. The math functions return null for null input and reject diff --git a/python/python/tests/test_math_functions.py b/python/python/tests/test_math_functions.py new file mode 100644 index 0000000..4cb066c --- /dev/null +++ b/python/python/tests/test_math_functions.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Cypher math types, nulls, and rounding through both execution backends.""" + +import math + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture(params=[False, True], ids=["join", "csr"]) +def engine(request): + config = GraphConfig.builder().node_table("Person", "id").build() + people = pa.table( + { + "id": [1, 2, 3], + "integer": pa.array([-128, None, 2], type=pa.int8()), + "number": pa.array([-1.5, None, 2.0], type=pa.float32()), + } + ) + return CypherEngine(config, {"Person": people}, use_csr=request.param) + + +@pytest.mark.parametrize( + ("expression", "expected"), + [ + ("abs(-5)", 5), + ("abs(-1.5)", 1.5), + ("round(-1.5)", -1.0), + ("round(1.5)", 2.0), + ("round(4503599627370495.5)", 4503599627370496.0), + ("round(4503599627370497.0)", 4503599627370497.0), + ("floor(-1.5)", -2.0), + ("ceil(-1.5)", -1.0), + ("sqrt(9)", 3.0), + ] + + [(f"{name}(null)", None) for name in ("abs", "round", "floor", "ceil", "sqrt")], +) +def test_math_value(engine, expression, expected): + result = engine.execute(f"RETURN {expression} AS value") + assert result.to_pylist() == [{"value": expected}] + assert result.schema.field("value").type == ( + pa.int64() if isinstance(expected, int) else pa.float64() + ) + + +@pytest.mark.parametrize("name", ["abs", "round", "floor", "ceil", "sqrt"]) +@pytest.mark.parametrize("argument", ["'1'", "true", "[1]", "{x: 1}", "", "1, 2"]) +def test_math_invalid_arguments(engine, name, argument): + with pytest.raises((ValueError, RuntimeError)): + engine.execute(f"RETURN {name}({argument}) AS value") + + +def test_math_integer_overflow(engine): + with pytest.raises((ValueError, RuntimeError), match="overflow"): + engine.execute("RETURN abs(-9223372036854775807 - 1) AS value") + + +def test_math_arrow_columns_and_nulls(engine): + result = engine.execute( + "MATCH (p:Person) RETURN p.id AS id, abs(p.integer) AS magnitude, " + "abs(p.number) AS floating, round(p.number) AS rounded, " + "floor(p.number) AS floored, ceil(p.number) AS ceiling, " + "sqrt(abs(p.integer)) AS root ORDER BY id" + ) + assert result.to_pylist() == [ + { + "id": 1, + "magnitude": 128, + "floating": 1.5, + "rounded": -1.0, + "floored": -2.0, + "ceiling": -1.0, + "root": math.sqrt(128), + }, + { + "id": 2, + "magnitude": None, + "floating": None, + "rounded": None, + "floored": None, + "ceiling": None, + "root": None, + }, + { + "id": 3, + "magnitude": 2, + "floating": 2.0, + "rounded": 2.0, + "floored": 2.0, + "ceiling": 2.0, + "root": math.sqrt(2), + }, + ] + assert result.schema.field("magnitude").type == pa.int64() + for column in ("floating", "rounded", "floored", "ceiling", "root"): + assert result.schema.field(column).type == pa.float64() + + +def test_math_nan_and_signed_zero(engine): + result = engine.execute("RETURN sqrt(-1) AS root, round(-0.5) AS zero").to_pylist()[ + 0 + ] + assert math.isnan(result["root"]) + assert math.copysign(1, result["zero"]) == 1 + + +def test_math_out_of_range_unsigned_integer(): + config = GraphConfig.builder().node_table("Person", "id").build() + engine = CypherEngine( + config, + {"Person": pa.table({"id": [1], "large": pa.array([2**63], type=pa.uint64())})}, + ) + with pytest.raises((ValueError, RuntimeError), match="cast|Cast"): + engine.execute("MATCH (p:Person) RETURN abs(p.large) AS value") + + +def test_math_non_finite_arrow_values(): + config = GraphConfig.builder().node_table("Person", "id").build() + engine = CypherEngine( + config, + {"Person": pa.table({"id": [1, 2], "number": [math.inf, math.nan]})}, + ) + rows = engine.execute( + "MATCH (p:Person) RETURN p.id AS id, round(p.number) AS rounded, " + "abs(p.number) AS magnitude, sqrt(p.number) AS root ORDER BY id" + ).to_pylist() + for name in ("rounded", "magnitude", "root"): + assert rows[0][name] == math.inf + assert math.isnan(rows[1][name]) From 99655beaf7e6db9b99147324370122127b034a12 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 11:55:47 -0700 Subject: [PATCH 39/66] fix(cypher): check scalar result types across query scopes Infer known math, string and list return types in the existing semantic match so WITH and UNWIND retain numeric/string mismatch checks. Preserve numeric promotion, nulls, Unicode and list elements. Validation: 638 Rust tests passed (10 ignored), 1146 Python tests passed, strict docs, snippets, Ruff, Pyright and typos passed. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph/src/semantic/types.rs | 39 +++++++++++++++++++-- docs/cypher.md | 2 +- python/python/tests/test_engine_types.py | 44 ++++++++++++++++++++++++ 3 files changed, 81 insertions(+), 4 deletions(-) diff --git a/crates/lance-graph/src/semantic/types.rs b/crates/lance-graph/src/semantic/types.rs index fad1621..9dcdfc2 100644 --- a/crates/lance-graph/src/semantic/types.rs +++ b/crates/lance-graph/src/semantic/types.rs @@ -132,8 +132,15 @@ impl SemanticAnalyzer { "count" | "size" | "length" | "tointeger" | "to_epoch_ms" => { Scalar(DataType::Int64) } - "tofloat" => Scalar(DataType::Float64), - "tolower" | "lower" | "toupper" | "upper" | "tostring" => { + "abs" => match first() { + Scalar(value) if value.is_integer() => Scalar(DataType::Int64), + Scalar(value) if value.is_floating() => Scalar(DataType::Float64), + Scalar(DataType::Null) => Scalar(DataType::Null), + _ => Unknown, + }, + "tofloat" | "round" | "floor" | "ceil" | "sqrt" => Scalar(DataType::Float64), + "tolower" | "lower" | "toupper" | "upper" | "tostring" | "left" | "right" + | "substring" | "trim" | "ltrim" | "rtrim" | "replace" => { Scalar(DataType::Utf8) } "date" => Scalar(DataType::Date32), @@ -142,7 +149,16 @@ impl SemanticAnalyzer { Some("UTC".into()), )), "collect" => List(Box::new(first())), - "head" | "__index" => first().element(), + "head" | "last" | "__index" => first().element(), + "split" | "keys" | "labels" => List(Box::new(Scalar(DataType::Utf8))), + "tail" | "reverse" => match first() { + Scalar(DataType::FixedSizeList(field, _)) + if name.eq_ignore_ascii_case("tail") => + { + Scalar(DataType::List(field)) + } + value => value, + }, "range" => List(Box::new(Scalar(DataType::Int64))), "nodes" => List(Box::new(Unknown)), "__list" => List(Box::new(common_type( @@ -363,6 +379,23 @@ mod tests { } } + #[test] + fn scalar_return_types_follow_with_and_unwind_scopes() { + for query in [ + "RETURN left('12', 1) + 1", + "RETURN abs(1) = '1'", + "WITH sqrt(4) AS value RETURN value = '2'", + "WITH split('1,2', ',') AS values UNWIND values AS value RETURN value + 1", + "RETURN last(reverse(['1'])) + 1", + "WITH tail(['1','2']) AS values UNWIND values AS value RETURN value + 1", + "UNWIND keys({one: 1}) AS key RETURN key = 1", + "MATCH (p:Person) UNWIND labels(p) AS label RETURN label = 1", + ] { + let errors = analyze(query).errors.join("\n"); + assert!(errors.contains("Type mismatch"), "{query}: {errors}"); + } + } + #[test] fn preserves_coercions_nulls_and_unknown_types() { for query in [ diff --git a/docs/cypher.md b/docs/cypher.md index 24683bf..609ccdc 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -525,7 +525,7 @@ aggregation and `OPTIONAL MATCH`. | Exception | Raised when | | --- | --- | -| `ValueError` | The query doesn't parse, refers to an unknown label, a dataset is missing a configured column, or a comparison/arithmetic expression has incompatible known types. Type errors name the expression and its operand types. | +| `ValueError` | The query doesn't parse, refers to an unknown label, a dataset is missing a configured column, or a comparison/arithmetic expression has incompatible known types. Type errors name the expression and its operand types. Known scalar-function return types are checked through `WITH` and `UNWIND` scopes. | | `NotImplementedError` | The query parses but uses an unsupported feature (for example a function that isn't implemented, or more than 20 hops). | | `RuntimeError` | Execution fails, for example when a Lance dataset can't be read. | diff --git a/python/python/tests/test_engine_types.py b/python/python/tests/test_engine_types.py index a629b00..bcd1e22 100644 --- a/python/python/tests/test_engine_types.py +++ b/python/python/tests/test_engine_types.py @@ -95,3 +95,47 @@ def test_comprehension_shadow_does_not_change_outer_type(engine): {"ages": [31], "name": "Ada"}, {"ages": [41], "name": "Bob"}, ] + + +@pytest.mark.parametrize( + "query", + [ + "RETURN left('12', 1) + 1", + "RETURN abs(1) = '1'", + "WITH sqrt(4) AS value RETURN value = '2'", + "WITH split('1,2', ',') AS values UNWIND values AS value RETURN value + 1", + "RETURN last(reverse(['1'])) + 1", + "WITH tail(['1', '2']) AS values UNWIND values AS value RETURN value + 1", + "UNWIND keys({one: 1}) AS key RETURN key = 1", + "MATCH (p:Person) UNWIND labels(p) AS label RETURN label = 1", + ], +) +def test_scalar_return_types_reject_string_numeric_coercion(engine, query): + with pytest.raises(ValueError, match="Type mismatch"): + engine.execute(query) + + +@pytest.mark.parametrize( + ("query", "expected"), + [ + ("RETURN abs(-1) + sqrt(4) AS value", [3.0]), + ("RETURN round(null) + 1 AS value", [None]), + ("RETURN abs(null) = '1' AS value", [None]), + ("RETURN left('猫😀', 1) = '猫' AS value", [True]), + ("RETURN reverse('猫😀') = '😀猫' AS value", [True]), + ( + "WITH split('猫,😀', ',') AS values UNWIND values AS value " + "RETURN value AS value ORDER BY value", + ["猫", "😀"], + ), + ("RETURN last(reverse([1, 2])) + 1 AS value", [2]), + ( + "WITH tail([1, 2, 3]) AS values UNWIND values AS value " + "RETURN value + 0.5 AS value ORDER BY value", + [2.5, 3.5], + ), + ("RETURN left(null, 1) AS value", [None]), + ], +) +def test_scalar_return_types_preserve_values_and_nulls(engine, query, expected): + assert engine.execute(query).column("value").to_pylist() == expected From 3465806c9e070f089b872e7f8f19f1e82703b93c Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 12:44:28 -0700 Subject: [PATCH 40/66] fix(cypher): preserve cross-label shortest path identity --- crates/lance-graph/src/logical_plan.rs | 2 +- crates/lance-graph/src/path_value.rs | 78 ++++++++++++++-- crates/lance-graph/src/query.rs | 31 ++++++- docs/csr.md | 4 +- docs/cypher.md | 5 +- python/python/tests/test_path_values.py | 113 ++++++++++++++++++++++++ 6 files changed, 221 insertions(+), 12 deletions(-) diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index c031014..75e1097 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -424,7 +424,7 @@ impl<'a> LogicalPlanner<'a> { /// Convert a Cypher AST to a logical plan pub fn plan(&mut self, query: &CypherQuery) -> Result { - self.rich_paths = crate::path_value::needs_values(query); + self.rich_paths = crate::path_value::needs_values(query, self.config); self.plan_query(query, None, false) } diff --git a/crates/lance-graph/src/path_value.rs b/crates/lance-graph/src/path_value.rs index 11e276e..a1b35be 100644 --- a/crates/lance-graph/src/path_value.rs +++ b/crates/lance-graph/src/path_value.rs @@ -123,9 +123,12 @@ pub(crate) fn paths_array( ))) } -/// Length/nodes queries, including WITH pass-through aliases, keep their CSR plan. +/// Homogeneous length/nodes queries, including WITH aliases, keep their CSR plan. /// Returned path values, relationship access and containers need actual path records. -pub(crate) fn needs_values(query: &crate::ast::CypherQuery) -> bool { +pub(crate) fn needs_values( + query: &crate::ast::CypherQuery, + config: &crate::config::GraphConfig, +) -> bool { use crate::ast::{GraphPattern, ReadingClause, ValueExpression as V}; let clauses: Vec<_> = query .reading_clauses @@ -146,8 +149,37 @@ pub(crate) fn needs_values(query: &crate::ast::CypherQuery) -> bool { if path.variable.is_some() { return true; } - if let Some(name) = path.shortest.as_ref().and_then(|s| s.variable.as_ref()) { - names.insert(name.to_lowercase()); + if path.shortest.is_some() { + // Numeric keys alone cannot distinguish nodes in different labels, + // even when both endpoints have the same label on an undirected path. + let crosses_labels = path.segments.iter().any(|segment| { + segment.relationship.types.iter().any(|kind| { + config + .get_relationship_mapping(kind) + .is_some_and(|mapping| { + mapping + .source_label + .as_deref() + .zip(mapping.target_label.as_deref()) + .is_some_and(|(source, target)| { + !source.eq_ignore_ascii_case(target) + }) + }) + }) + }); + let different_ends = path + .start_node + .labels + .first() + .zip(path.segments.last().and_then(|s| s.end_node.labels.first())) + .is_some_and(|(source, target)| !source.eq_ignore_ascii_case(target)); + if crosses_labels || different_ends { + return true; + } + if let Some(name) = path.shortest.as_ref().and_then(|s| s.variable.as_ref()) + { + names.insert(name.to_lowercase()); + } } } } @@ -170,10 +202,10 @@ pub(crate) fn needs_values(query: &crate::ast::CypherQuery) -> bool { let mut needed = query .union_branches .iter() - .any(|branch| needs_values(&branch.query)) + .any(|branch| needs_values(&branch.query, config)) || clauses.iter().any(|clause| match clause { ReadingClause::Call(inner) => { - needs_values(inner) + needs_values(inner, config) || inner .call_union_imports() .unwrap_or_default() @@ -274,6 +306,13 @@ mod tests { #[test] fn direct_shortest_queries_keep_the_csr_route() { + let config = crate::config::GraphConfig::builder() + .node_table("Person", "id") + .node_table("Place", "id") + .rel_table("R", "Person", "Person", "src", "dst") + .rel_table("LIVES", "Person", "Place", "src", "dst") + .build() + .unwrap(); let prefix = "MATCH p = shortestPath((a:Person)-[:R*]-(b:Person)) "; for (suffix, rich) in [ ("RETURN length(p), nodes(p)", false), @@ -303,7 +342,32 @@ mod tests { ("WITH collect(p) AS ps UNWIND ps AS p RETURN nodes(p)", true), ] { let query = crate::parser::parse_cypher_query(&format!("{prefix}{suffix}")).unwrap(); - assert_eq!(needs_values(&query), rich, "{suffix}"); + assert_eq!(needs_values(&query, &config), rich, "{suffix}"); + } + } + + #[test] + fn cross_label_shortest_paths_need_typed_identity() { + let config = crate::config::GraphConfig::builder() + .node_table("Person", "id") + .node_table("Place", "id") + .rel_table("LIVES", "Person", "Place", "src", "dst") + .build() + .unwrap(); + for ends in ["(b:Place)", "(b:Person)", "(b)"] { + let query = crate::parser::parse_cypher_query(&format!( + "MATCH p = shortestPath((a:Person)-[:LIVES*0..]-{ends}) RETURN nodes(p)" + )) + .unwrap(); + assert!(needs_values(&query, &config)); + } + for text in [ + "CALL { MATCH p = shortestPath((a:Person)-[:LIVES*0..]->(b:Place)) RETURN length(p) AS n } RETURN n", + "RETURN 0 AS n UNION MATCH p = shortestPath((a:Person)-[:LIVES*0..]->(b:Place)) RETURN length(p) AS n", + "CALL { RETURN 0 AS n UNION MATCH p = shortestPath((a:Person)-[:LIVES*0..]-(b:Person)) RETURN length(p) AS n } RETURN n", + ] { + let query = crate::parser::parse_cypher_query(text).unwrap(); + assert!(needs_values(&query, &config), "{text}"); } } } diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 0b33a65..f7dded4 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -237,7 +237,7 @@ impl CypherQuery { ) -> Result { // Rich path properties live in an extension-owned spec, and its inventory is // query-local. Keep ordinary substitution and a fresh inventory for each query. - let mut shape_sensitive = crate::path_value::needs_values(&self.ast); + let mut shape_sensitive = false; self.ast.for_each_query(&mut |query| { shape_sensitive |= query .limit @@ -260,7 +260,10 @@ impl CypherQuery { .all_reading_clauses() .any(|clause| matches!(clause, ReadingClause::Call(_))); }); - if shape_sensitive || self.scalar_parameters().is_none() { + if shape_sensitive + || self.scalar_parameters().is_none() + || crate::path_value::needs_values(&self.ast, self.require_config()?) + { // CALL binds its own child templates; row counts alter plan shape. return Ok(false); } @@ -1732,6 +1735,30 @@ mod tests { } } + #[test] + fn cross_label_shortest_queries_bypass_cached_inventories() { + let config = GraphConfig::builder() + .node_table("Person", "id") + .node_table("Place", "id") + .rel_table("LIVES", "Person", "Place", "src", "dst") + .build() + .unwrap(); + for text in [ + "MATCH p = shortestPath((a:Person {id: $id})-[:LIVES*0..]->(b:Place)) RETURN length(p)", + "MATCH p = shortestPath((a:Person {id: $id})-[:LIVES*0..]-(b:Person)) WITH p AS q RETURN nodes(q)", + "RETURN 0 AS n UNION MATCH p = shortestPath((a:Person {id: $id})-[:LIVES*0..]->(b:Place)) RETURN length(p) AS n", + ] { + let mut query = CypherQuery::new(text) + .unwrap() + .with_config(config.clone()) + .with_parameter("id", 1); + assert!(!query + .prepare_cached_plan(Arc::new(crate::InMemoryCatalog::new())) + .unwrap()); + assert!(query.cached_plan.is_none(), "{text}"); + } + } + // Opt-in diagnostic: reports stage means for a fixed local scan, with no timing // assertions. Physical planning calls the query planner directly, avoiding the // extra optimize() performed by SessionState::create_physical_plan(). diff --git a/docs/csr.md b/docs/csr.md index 9db5951..97e2244 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -136,7 +136,9 @@ Named path values and `relationships(path)` use a query-local inventory of actua CSR neighbors do not carry edge identity or properties. This also preserves parallel edges when a path passes through `collect()` and `UNWIND`. The inventory materializes the participating node and relationship tables; bounded traversal batches do not bound that inventory. See [path values](cypher.md#path-values). -Shortest-path queries that only read length or node keys keep the CSR route below. +Shortest-path queries within one node label that only read length or node keys keep the CSR route below. +Relationships crossing labels use typed identities, even for these queries, to distinguish equal keys +from different labels. For `shortestPath` and `allShortestPaths`, a source with one target in an input batch searches from both ends using the outgoing and incoming CSR. It expands the smaller frontier and finishes the whole layer diff --git a/docs/cypher.md b/docs/cypher.md index 609ccdc..2b080fd 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -362,7 +362,10 @@ in batches and named ordinary paths keep a cursor per hop, but inventory memory Rich path queries bypass the query-template cache: each execution binds its current parameters and creates an inventory from its own snapshot. A reader opened before `refresh()` retains the old snapshot. Queries asking only for shortest-path length or node keys, including homogeneous node-property -comprehensions and paths carried through `WITH` aliases, retain the usual CSR execution route. +comprehensions and paths carried through `WITH` aliases, retain the usual CSR execution route when +the relationship joins the same node label. Cross-label shortest paths use typed identities even for +length or node keys, so equal numeric keys from different labels cannot become a false zero-hop path. +This also applies inside `CALL` and `UNION`, and when equal-label endpoints pass through another label. ### Expressions and predicates diff --git a/python/python/tests/test_path_values.py b/python/python/tests/test_path_values.py index e19cc62..512e46f 100644 --- a/python/python/tests/test_path_values.py +++ b/python/python/tests/test_path_values.py @@ -492,3 +492,116 @@ def test_rich_path_rejects_ambiguous_node_ids(tmp_path, storage): Exception, match="unique node IDs within label 'Person'; duplicate ID 1" ): engine.execute("MATCH p = (a:Person)-[:R]->(b:Person) RETURN p") + + +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize("search", ["shortestPath", "allShortestPaths"]) +@pytest.mark.parametrize("minimum", [0, 1]) +def test_shortest_scalar_consumers_preserve_cross_label_identity( + use_csr, search, minimum +): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .node_table("Place", "id") + .rel_table("LIVES", "Person", "Place", "src", "dst") + .build() + ) + engine = CypherEngine( + config, + { + "Person": pa.table({"id": [1, 2], "name": ["one", "two"]}), + "Place": pa.table({"id": [1], "name": ["place"]}), + "LIVES": pa.table({"src": [1, 2], "dst": [1, 1]}), + }, + use_csr=use_csr, + ) + query = ( + f"MATCH p = {search}((a:Person {{id: 1}})-[:LIVES*{minimum}..]->(b:Place)) " + "WITH p AS q RETURN length(q) AS hops, nodes(q) AS ids, " + "[n IN nodes(q) | n.name] AS names, [n IN nodes(q) | labels(n)] AS labels" + ) + assert engine.execute(query).to_pylist() == [ + { + "hops": 1, + "ids": [1, 1], + "names": ["one", "place"], + "labels": [["Person"], ["Place"]], + } + ] + assert "PathMatchExec" in engine.explain(query) + # Equal endpoint labels do not make Place 1 interchangeable with Person 1. + rows = engine.execute( + f"MATCH p = {search}((a:Person {{id: 1}})" + f"-[:LIVES*{minimum}..]-(b:Person {{id: 2}})) " + "RETURN length(p) AS hops, nodes(p) AS ids, [n IN nodes(p) | n.name] AS names" + ).to_pylist() + assert rows == [{"hops": 2, "ids": [1, 1, 2], "names": ["one", "place", "two"]}] + + +@pytest.fixture(params=[False, True]) +def cross_label_graph(request): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .node_table("Place", "id") + .rel_table("LIVES", "Person", "Place", "src", "dst") + .build() + ) + tables = { + "Person": pa.table({"id": [1, 2]}), + "Place": pa.table({"id": [1, 2]}), + "LIVES": pa.table({"src": [1, 2], "dst": [1, 1]}), + } + return config, tables, request.param + + +@pytest.mark.parametrize("scope", ["direct", "alias", "call", "union", "call_union"]) +def test_cross_label_shortest_scopes_rebind_each_parameter(cross_label_graph, scope): + config, tables, use_csr = cross_label_graph + engine = CypherEngine(config, tables, use_csr=use_csr) + prefix = "MATCH p = shortestPath((a:Person {id: $start})-[:LIVES*0..]->(b:Place)) " + body = prefix + "RETURN length(p) AS hops, nodes(p) AS ids" + queries = { + "direct": body, + "alias": prefix + "WITH p AS q RETURN length(q) AS hops, nodes(q) AS ids", + "call": "CALL { " + body + " } RETURN hops, ids", + "union": body + " UNION ALL " + body, + "call_union": "CALL { " + body + " UNION ALL " + body + " } RETURN hops, ids", + } + query = queries[scope] + copies = 2 if "union" in scope else 1 + for start in [1, 2, None, 1]: + expected = [] if start is None else [{"hops": 1, "ids": [start, 1]}] * copies + params = {"start": start} + assert engine.execute(query, params).to_pylist() == expected + assert engine.execute_batches(query, params).read_all().to_pylist() == expected + + +def test_cross_label_shortest_inventory_isolates_refresh(cross_label_graph, tmp_path): + import lance + + config, tables, use_csr = cross_label_graph + paths = {name: str(tmp_path / f"{name}.lance") for name in tables} + datasets = { + name: lance.write_dataset(table, paths[name]) for name, table in tables.items() + } + engine = CypherEngine(config, datasets, use_csr=use_csr) + query = ( + "MATCH p = shortestPath((a:Person {id: $start})-[:LIVES*0..]->(b:Place)) " + "WITH p AS q RETURN nodes(q) AS ids" + ) + for start in [1, 2, 1]: + assert engine.execute(query, {"start": start}).to_pylist() == [ + {"ids": [start, 1]} + ] + reader = engine.execute_batches(query, {"start": 1}) + lance.write_dataset( + pa.table({"src": [1], "dst": [2]}), paths["LIVES"], mode="overwrite" + ) + engine.refresh() + assert reader.read_all().to_pylist() == [{"ids": [1, 1]}] + for start in [1, 2, 1]: + assert engine.execute(query, {"start": start}).to_pylist() == ( + [{"ids": [1, 2]}] if start == 1 else [] + ) From ff9108f8da498a8bc83781b05f32f5f19fde1dd6 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 12:50:38 -0700 Subject: [PATCH 41/66] fix(cypher): reject unsigned range argument overflow Return an error when an Arrow UInt64 range endpoint or step exceeds Int64 instead of silently converting it to null. Preserve actual null arguments. Add Rust and CSR/join Python regressions and document the boundary. --- .../datafusion_planner/scalar_functions.rs | 41 +++++++++++++++++-- docs/cypher.md | 3 +- python/python/tests/test_scalar_functions.py | 20 +++++++++ 3 files changed, 60 insertions(+), 4 deletions(-) diff --git a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs index 33bb32e..465729d 100644 --- a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs +++ b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs @@ -9,7 +9,7 @@ use std::sync::Arc; use arrow::array::{ new_null_array, Array, AsArray, Int64Array, ListBuilder, StringArray, StringBuilder, }; -use arrow::compute::cast; +use arrow::compute::{cast, cast_with_options, CastOptions}; use arrow_schema::{DataType, Field}; use datafusion::common::{exec_err, plan_err, ScalarValue}; use datafusion::error::Result; @@ -184,8 +184,14 @@ impl ScalarUDFImpl for CypherScalar { } if self.name == "range" { for arg in &mut args.args { - *arg = - ColumnarValue::Array(cast(&arg.to_array(args.number_rows)?, &DataType::Int64)?); + *arg = ColumnarValue::Array(cast_with_options( + &arg.to_array(args.number_rows)?, + &DataType::Int64, + &CastOptions { + safe: false, + ..Default::default() + }, + )?); } return range::gen_series_udf().invoke_with_args(args); } @@ -368,6 +374,35 @@ mod tests { use arrow::array::{ArrayRef, RecordBatch}; use std::collections::HashMap; + #[tokio::test] + async fn range_rejects_unsigned_overflow() { + let batch = RecordBatch::try_from_iter(vec![ + ("id", Arc::new(Int64Array::from(vec![1])) as ArrayRef), + ( + "large", + Arc::new(arrow::array::UInt64Array::from(vec![1_u64 << 63])) as ArrayRef, + ), + ]) + .unwrap(); + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + for expression in [ + "range(p.large, 1)", + "range(1, p.large)", + "range(1, 2, p.large)", + ] { + let error = CypherQuery::new(&format!("MATCH (p:Person) RETURN {expression} AS value")) + .unwrap() + .with_config(config.clone()) + .execute(HashMap::from([("Person".to_string(), batch.clone())]), None) + .await + .unwrap_err(); + assert!(error.to_string().contains("cast"), "{expression}: {error}"); + } + } + async fn run(expression: &str) -> crate::error::Result { let batch = RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(vec![1])) as ArrayRef)]) diff --git a/docs/cypher.md b/docs/cypher.md index 2b080fd..df005b4 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -414,7 +414,8 @@ Splitting an empty string returns `['']`. `last([])` returns null; `tail([])` and `reverse([])` return empty lists. `reverse` also accepts strings. `range(start, end [, step])` includes the end when reached, defaults to step 1, and accepts negative steps. A step pointing away from the end gives an empty list; zero steps and non-integer arguments -raise an error. Null arguments return null. +raise an error. Unsigned Arrow arguments outside Int64 range also raise an error; they are not +converted to null. Null arguments return null. `abs` preserves integer results; the other math functions return floating point values. Integer absolute-value overflow and unsigned inputs outside Int64 range raise an error. `round` accepts one argument, with ties toward positive diff --git a/python/python/tests/test_scalar_functions.py b/python/python/tests/test_scalar_functions.py index 196b76d..b52b308 100644 --- a/python/python/tests/test_scalar_functions.py +++ b/python/python/tests/test_scalar_functions.py @@ -10,6 +10,26 @@ from lance_csr import CypherEngine, GraphConfig +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize( + "expression", ["range(p.value, 1)", "range(1, p.value)", "range(1, 2, p.value)"] +) +def test_range_unsigned_overflow_and_null(use_csr, expression): + config = GraphConfig.builder().node_table("Person", "id").build() + query = f"MATCH (p:Person) RETURN {expression} AS value" + for value in (None, 2**63): + engine = CypherEngine( + config, + {"Person": pa.table({"id": [1], "value": pa.array([value], pa.uint64())})}, + use_csr=use_csr, + ) + if value is None: + assert engine.execute(query).to_pylist() == [{"value": None}] + else: + with pytest.raises((ValueError, RuntimeError), match="cast|Cast"): + engine.execute(query) + + @pytest.fixture(params=[False, True], ids=["join", "csr"]) def engine(request): config = GraphConfig.builder().node_table("Person", "id").build() From c9b3c97ad8206a0b5df54dab2cfbe4ec97c09e2b Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 12:51:32 -0700 Subject: [PATCH 42/66] fix(cypher): guard SQL export of runtime-only functions Reject custom Cypher scalar implementations that differ from registered SQL functions after constant folding. Preserve list size and inclusive range exports through native expressions. Compare real Cypher and SqlEngine results for rounding, integer width, Unicode, lists and nulls; document the existing NotImplementedError boundary. --- crates/lance-graph-python/src/graph.rs | 4 +- .../lance-graph/src/datafusion_planner/mod.rs | 1 + .../datafusion_planner/scalar_functions.rs | 28 ++++++ crates/lance-graph/src/query.rs | 74 +++++++++++++++ docs/api-reference.md | 3 +- docs/sql.md | 14 ++- python/python/tests/test_to_sql.py | 94 ++++++++++++++++++- 7 files changed, 214 insertions(+), 4 deletions(-) diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 6f7323f..5d6ee3d 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -595,7 +595,7 @@ impl CypherQuery { Ok(dict.unbind()) } - /// Convert query to SQL + /// Export supported Cypher expressions as SQL /// /// Parameters /// ---------- @@ -611,6 +611,8 @@ impl CypherQuery { /// /// Raises /// ------ + /// NotImplementedError + /// If a runtime-only Cypher function has no faithful SQL translation /// RuntimeError /// If SQL generation fails #[pyo3(signature = (datasets, dialect=None))] diff --git a/crates/lance-graph/src/datafusion_planner/mod.rs b/crates/lance-graph/src/datafusion_planner/mod.rs index ee85dc4..0698ea7 100644 --- a/crates/lance-graph/src/datafusion_planner/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/mod.rs @@ -32,6 +32,7 @@ mod test_fixtures; // Re-export public types pub use analysis::{PlanningContext, QueryAnalysis, RelationshipInstance}; +pub(crate) use scalar_functions::sql_expr; use crate::config::GraphConfig; use crate::error::Result; diff --git a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs index 465729d..d877529 100644 --- a/crates/lance-graph/src/datafusion_planner/scalar_functions.rs +++ b/crates/lance-graph/src/datafusion_planner/scalar_functions.rs @@ -367,6 +367,34 @@ pub(super) fn call(name: &str, args: Vec) -> Expr { .call(args) } +/// SQL equivalents for the list operations supported before Cypher-specific wrappers. +pub(crate) fn sql_expr( + function: &datafusion::logical_expr::expr::ScalarFunction, + schema: &datafusion::common::DFSchema, +) -> Result> { + use datafusion::logical_expr::{expr_fn::cast, ExprSchemable}; + + if !function.func.inner().is::() { + return Ok(None); + } + let args = &function.args; + Ok(match function.name() { + "size" if list_field(&args[0].get_type(schema)?).is_some() => Some(cast( + length::array_length_udf().call(args.clone()), + DataType::Int64, + )), + "range" => Some( + range::gen_series_udf().call( + args.iter() + .cloned() + .map(|arg| cast(arg, DataType::Int64)) + .collect(), + ), + ), + _ => None, + }) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index f7dded4..17dbd91 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -474,11 +474,20 @@ impl CypherQuery { /// /// # Returns /// A SQL string representing the query in the specified dialect + /// + /// # Errors + /// Returns `GraphError::UnsupportedFeature` for runtime-only Cypher functions + /// without a faithful SQL translation. Constant expressions may be folded first. pub async fn to_sql( &self, datasets: HashMap, dialect: Option, ) -> Result { + use datafusion::common::tree_node::{Transformed, TreeNode}; + use datafusion::common::{not_impl_err, DFSchema}; + use datafusion::logical_expr::expr_rewriter::NamePreserver; + use datafusion::logical_expr::registry::FunctionRegistry; + use datafusion::logical_expr::{utils::merge_schema, Expr, LogicalPlan}; use std::sync::Arc; let dialect = dialect.unwrap_or_default(); @@ -502,6 +511,40 @@ impl CypherQuery { location: snafu::Location::new(file!(), line!(), column!()), })?; + // An SQL function with the same name need not implement Cypher semantics. + let optimized_plan = optimized_plan + .transform_up_with_subqueries(|plan| { + let schema = if let LogicalPlan::TableScan(scan) = &plan { + DFSchema::try_from_qualified_schema(scan.table_name.clone(), &scan.source.schema())? + } else { + merge_schema(&plan.inputs()) + }; + let names = NamePreserver::new(&plan); + plan.map_expressions(|expr| { + let name = names.save(&expr); + let rewritten = expr.transform_up(|expr| { + if let Expr::ScalarFunction(function) = &expr { + if let Some(native) = crate::datafusion_planner::sql_expr(function, &schema)? { + return Ok(Transformed::yes(native)); + } + if !ctx.udf(function.name()).is_ok_and(|registered| registered == function.func) { + return not_impl_err!( + "SQL export of Cypher function '{}' is unsupported; execute the Cypher query directly", + function.name() + ); + } + } + Ok(Transformed::no(expr)) + })?; + Ok(rewritten.update_data(|expr| name.restore(expr))) + }) + }) + .map_err(|error| GraphError::UnsupportedFeature { + feature: error.to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + })? + .data; + // Unparse to SQL using the specified dialect let dialect_unparser = dialect.unparser(); let unparser = dialect_unparser.as_unparser(); @@ -2510,6 +2553,37 @@ mod tests { assert!(sql.contains("p.name")); } + #[tokio::test] + async fn test_to_sql_rejects_runtime_only_cypher_function() { + let batch = RecordBatch::try_from_iter(vec![ + ( + "id", + Arc::new(arrow_array::Int64Array::from(vec![1])) as arrow_array::ArrayRef, + ), + ( + "number", + Arc::new(arrow_array::Float64Array::from(vec![-1.5])) as arrow_array::ArrayRef, + ), + ]) + .unwrap(); + let query = CypherQuery::new("MATCH (p:Person) RETURN round(p.number) AS value") + .unwrap() + .with_config( + GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(), + ); + let error = query + .to_sql(HashMap::from([("Person".to_string(), batch)]), None) + .await + .unwrap_err(); + assert!(matches!(error, GraphError::UnsupportedFeature { .. })); + assert!(error + .to_string() + .contains("SQL export of Cypher function 'round'")); + } + async fn write_lance_dataset(path: &std::path::Path, batch: arrow_array::RecordBatch) { use arrow_array::{RecordBatch, RecordBatchIterator}; use lance::dataset::{Dataset, WriteParams}; diff --git a/docs/api-reference.md b/docs/api-reference.md index 43ccf0f..dd2e790 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -211,7 +211,8 @@ A parsed Cypher query. Raises `ValueError` if `query_text` doesn't parse. Builde : The plans, as for `CypherEngine.explain`. `#!python to_sql(datasets: dict, dialect: SqlDialect | None = None) -> str` -: The equivalent SQL, as joins. See [From Cypher to SQL](sql.md#from-cypher-to-sql). +: Exports supported expressions as SQL joins. Raises `NotImplementedError` for runtime-only Cypher + functions without a faithful SQL translation. See [From Cypher to SQL](sql.md#from-cypher-to-sql). `#!python execute_with_vector_rerank(datasets: dict, vector_search: VectorSearch) -> pyarrow.Table` : Runs the query, then reranks by vector similarity. diff --git a/docs/sql.md b/docs/sql.md index 57f60c3..e2cd70e 100644 --- a/docs/sql.md +++ b/docs/sql.md @@ -59,7 +59,7 @@ The SQL dialect is DataFusion's; see the ## From Cypher to SQL -`CypherQuery.to_sql(datasets, dialect=None)` shows the SQL that a Cypher query is equivalent to. It's handy +`CypherQuery.to_sql(datasets, dialect=None)` exports supported Cypher expressions as SQL. It's handy for understanding a query, or for running a graph pattern on another SQL engine: ```python @@ -82,6 +82,18 @@ print(query.to_sql(graph, SqlDialect.PostgreSql)) `SqlDialect` has `Default` (DataFusion's), `Spark`, `PostgreSql`, `MySql` and `Sqlite`. +Export is experimental and supports a subset of Cypher execution. Runtime-only Cypher functions +raise `NotImplementedError` when no faithful SQL expression is available. This includes math wrappers such +as `round` and `abs`, and string functions such as `left` and string `size`, when they depend on +table values. Cypher `round(-1.5)` produces `-1.0`; DataFusion SQL `round(-1.5)` produces `-2.0`. +Integer `abs` also normalizes its result to Int64 in Cypher, while SQL preserves the input width. +Use `CypherQuery.execute` or `CypherEngine.execute` for these expressions. + +List `size` and inclusive `range` retain SQL export support through native expressions, including +null arguments, integer casts and the optional range step. Constant expressions that the optimizer +evaluates completely can also export their resulting literals. These checks establish DataFusion +semantics; other SQL dialects still require validation on the destination engine. + !!! info "This is the join path" The generated SQL expresses each hop as a join. It describes what the query computes, not how diff --git a/python/python/tests/test_to_sql.py b/python/python/tests/test_to_sql.py index 8806d8f..db86e81 100644 --- a/python/python/tests/test_to_sql.py +++ b/python/python/tests/test_to_sql.py @@ -5,7 +5,99 @@ import pyarrow as pa import pytest -from lance_csr import CypherQuery, GraphConfig +from lance_csr import CypherQuery, GraphConfig, SqlEngine + + +@pytest.fixture +def scalar_export_env(): + config = GraphConfig.builder().node_table("Person", "id").build() + datasets = { + "Person": pa.table( + { + "id": [1, 2, 3], + "number": pa.array([-1.5, None, 2.0], pa.float32()), + "integer": pa.array([-100, None, 2], pa.int8()), + "text": pa.array(["a猫😀z", None, ""], pa.large_string()), + "items": pa.array([[1, None], None, []], pa.list_(pa.int64())), + "nested": pa.array( + [[[1, 2], [3]], None, []], pa.list_(pa.list_(pa.int64())) + ), + "end": pa.array([3, None, -1], pa.int8()), + } + ) + } + return config, datasets + + +@pytest.mark.parametrize( + "expression", + [ + "size(p.items)", + "size(p.nested)", + "range(1, p.end)", + "range(p.end, 1, -1)", + "size(range(1, p.end))", + "toUpper(p.text)", + "round(-1.5)", + "left('a猫😀z', 3)", + ], +) +def test_sql_export_preserves_scalar_values_and_types(scalar_export_env, expression): + config, datasets = scalar_export_env + query = CypherQuery( + f"MATCH (p:Person) RETURN p.id AS id, {expression} AS value ORDER BY id" + ).with_config(config) + cypher = query.execute(datasets) + sql = SqlEngine(datasets).execute(query.to_sql(datasets)) + assert sql.to_pylist() == cypher.to_pylist() + assert sql.schema.types == cypher.schema.types + + +@pytest.mark.parametrize( + ("expression", "function", "expected"), + [ + ("round(p.number)", "round", [-1.0, None, 2.0]), + ("abs(p.integer)", "abs", [100, None, 2]), + ("left(p.text, 3)", "left", ["a猫😀", None, ""]), + ("size(p.text)", "size", [4, None, 0]), + ("reverse(p.text)", "reverse", ["z😀猫a", None, ""]), + ], +) +def test_sql_export_rejects_runtime_only_cypher_functions( + scalar_export_env, expression, function, expected +): + config, datasets = scalar_export_env + query = CypherQuery( + f"MATCH (p:Person) RETURN p.id AS id, {expression} AS value ORDER BY id" + ).with_config(config) + result = query.execute(datasets) + assert result.column("value").to_pylist() == expected + with pytest.raises( + NotImplementedError, match=f"SQL export of Cypher function '{function}'" + ): + query.to_sql(datasets) + + # A same-named SQL function can differ in values or Arrow result width. + engine = SqlEngine(datasets) + if function == "round": + assert engine.execute( + "SELECT round(number) AS value FROM person ORDER BY id" + ).column("value").to_pylist() == [-2.0, None, 2.0] + if function == "abs": + sql = engine.execute("SELECT abs(integer) AS value FROM person ORDER BY id") + assert sql.schema.field("value").type == pa.int8() + assert result.schema.field("value").type == pa.int64() + + +def test_sql_export_checks_filter_functions(scalar_export_env): + config, datasets = scalar_export_env + query = CypherQuery( + "MATCH (p:Person) WHERE round(p.number) = -1 RETURN p.id" + ).with_config(config) + with pytest.raises( + NotImplementedError, match="SQL export of Cypher function 'round'" + ): + query.to_sql(datasets) @pytest.fixture From 91d6b482a625b671c1e2eef923a4105331dc7ef4 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 12:54:02 -0700 Subject: [PATCH 43/66] test(query): configure graph for path cache eligibility --- crates/lance-graph/src/apply.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/lance-graph/src/apply.rs b/crates/lance-graph/src/apply.rs index 719ccf9..66e9fd1 100644 --- a/crates/lance-graph/src/apply.rs +++ b/crates/lance-graph/src/apply.rs @@ -497,6 +497,7 @@ mod tests { #[test] fn call_bypasses_snapshot_templates_in_every_reading_stage() { + let (config, _) = graph(); let catalog = Arc::new(lance_graph_catalog::InMemoryCatalog::new()); for text in [ "CALL { RETURN 1 AS n } RETURN n", @@ -507,7 +508,7 @@ mod tests { "MATCH path = (p:Person)-[:KNOWS]->(f:Person) RETURN path", "MATCH path = shortestPath((p:Person)-[:KNOWS*..3]->(f:Person)) RETURN relationships(path)", ] { - let mut query = CypherQuery::new(text).unwrap(); + let mut query = CypherQuery::new(text).unwrap().with_config(config.clone()); assert!(!query.prepare_cached_plan(catalog.clone()).unwrap()); } } From 2dac56b110eda9db471c0dd43e3e39ea64ada2aa Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 14:24:43 -0700 Subject: [PATCH 44/66] feat(python): allow query partition target overrides --- crates/lance-graph-python/src/graph.rs | 100 ++++++++++++++++-- docs/api-reference.md | 9 +- docs/cypher.md | 5 +- docs/lance.md | 22 ++++ examples/lance_snapshot_refresh.py | 2 + python/python/lance_csr/_internal/graph.pyi | 7 +- python/python/tests/test_target_partitions.py | 92 ++++++++++++++++ python/python/tests/typing/public_api.py | 4 + 8 files changed, 227 insertions(+), 14 deletions(-) create mode 100644 python/python/tests/test_target_partitions.py diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 5d6ee3d..f53246e 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -22,7 +22,7 @@ use arrow::ffi_stream::ArrowArrayStreamReader; use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::Schema; use datafusion::datasource::{DefaultTableSource, MemTable, TableProvider}; -use datafusion::execution::context::SessionContext; +use datafusion::execution::context::{SessionConfig, SessionContext}; use datafusion::execution::session_state::SessionStateBuilder; use futures::StreamExt; use lance::datafusion::LanceTableProvider; @@ -40,9 +40,9 @@ use lance_graph::{ SqlDialect as RustSqlDialect, SqlQuery as RustSqlQuery, VectorSearch as RustVectorSearch, }; use pyo3::{ - exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, + exceptions::{PyNotImplementedError, PyRuntimeError, PyTypeError, PyValueError}, prelude::*, - types::{PyDict, PyInt, PyList}, + types::{PyBool, PyDict, PyInt, PyList}, IntoPyObject, }; use serde_json::Value as JsonValue; @@ -1537,12 +1537,18 @@ fn build_engine_state( config: &RustGraphConfig, sources: HashMap, use_csr: bool, + target_partitions: Option, previous: Option<&EngineState>, ) -> PyResult { + let mut session_config = SessionConfig::new(); + if let Some(partitions) = target_partitions { + session_config = session_config.with_target_partitions(partitions); + } // The CSR planner only changes plans that contain CSR nodes, so install it always. let ctx = SessionContext::new_with_state( SessionStateBuilder::new() .with_default_features() + .with_config(session_config) .with_query_planner(Arc::new(CsrQueryPlanner)) // Move target-only filters into CSR nodes first, so pruning can then drop // the filter columns from the take. @@ -2109,6 +2115,7 @@ pub struct CypherEngine { config: RustGraphConfig, state: RwLock>, use_csr: bool, + target_partitions: Option, } impl CypherEngine { @@ -2213,23 +2220,48 @@ impl CypherEngine { /// datasets : dict /// Dictionary mapping table names to Lance datasets or PyArrow tables. /// Lance datasets are not read here; PyArrow inputs are held in memory. + /// target_partitions : int or None, optional + /// Positive target for DataFusion query partitioning, preserved by refresh(). + /// None keeps DataFusion's default. This is not a thread or memory limit. /// /// Raises /// ------ /// ValueError /// If no datasets are given, a configured id/endpoint column is missing, or a /// Lance dataset cannot be reopened by URI (for example, namespace-vended - /// credentials) + /// credentials), or target_partitions is zero, negative or out of range + /// TypeError + /// If target_partitions is a boolean or a non-integer value other than None /// RuntimeError /// If a dataset cannot be opened or registered #[new] - #[pyo3(signature = (config, datasets, use_csr = false))] + #[pyo3(signature = (config, datasets, use_csr = false, *, target_partitions = None))] fn new( py: Python, config: &GraphConfig, datasets: &Bound<'_, PyDict>, use_csr: bool, + target_partitions: Option<&Bound<'_, PyAny>>, ) -> PyResult { + let target_partitions = target_partitions + .map(|value| { + if value.is_instance_of::() || !value.is_instance_of::() { + return Err(PyTypeError::new_err( + "target_partitions must be an integer or None, not a boolean", + )); + } + value + .extract::() + .ok() + .filter(|value| *value > 0) + .ok_or_else(|| { + PyValueError::new_err(format!( + "target_partitions must be between 1 and {}", + usize::MAX + )) + }) + }) + .transpose()?; if datasets.is_empty() { return Err(PyValueError::new_err("No input datasets provided")); } @@ -2239,11 +2271,13 @@ impl CypherEngine { let source = python_value_to_source(py, &name, &value)?; sources.insert(name, source); } - let state = build_engine_state(py, &config.inner, sources, use_csr, None)?; + let state = + build_engine_state(py, &config.inner, sources, use_csr, target_partitions, None)?; Ok(Self { config: config.inner.clone(), state: RwLock::new(Arc::new(state)), use_csr, + target_partitions, }) } @@ -2372,8 +2406,14 @@ impl CypherEngine { } } } - let state = - build_engine_state(py, &self.config, sources, self.use_csr, Some(¤t))?; + let state = build_engine_state( + py, + &self.config, + sources, + self.use_csr, + self.target_partitions, + Some(¤t), + )?; if self.try_publish_refresh(¤t, state) { return Ok(()); } @@ -2751,6 +2791,48 @@ mod tests { } } + #[test] + fn target_partitions_survive_refresh_without_changing_other_defaults() { + Python::initialize(); + Python::attach(|py| { + let defaults = SessionConfig::new(); + let config = RustGraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + for target_partitions in [None, Some(1), Some(2)] { + let state = build_engine_state( + py, + &config, + test_state(1).sources, + false, + target_partitions, + None, + ) + .unwrap(); + let engine = CypherEngine { + config: config.clone(), + state: RwLock::new(Arc::new(state)), + use_csr: false, + target_partitions, + }; + let before = engine.snapshot(); + engine.refresh(py).unwrap(); + for snapshot in [before, engine.snapshot()] { + let state = snapshot.context.state(); + assert_eq!( + state.config().target_partitions(), + target_partitions.unwrap_or(defaults.target_partitions()) + ); + assert_eq!( + state.config().round_robin_repartition(), + defaults.round_robin_repartition() + ); + } + } + }); + } + #[test] fn test_plan_cache_reuses_bounds_and_refreshes() { let engine = CypherEngine { @@ -2760,6 +2842,7 @@ mod tests { .unwrap(), state: RwLock::new(Arc::new(test_state(1))), use_csr: false, + target_partitions: None, }; let snapshot = engine.snapshot(); let text = "MATCH (p:Person) WHERE p.id = $value RETURN p.id"; @@ -2835,6 +2918,7 @@ mod tests { .unwrap(), state: RwLock::new(Arc::new(test_state(1))), use_csr: false, + target_partitions: None, }; let reader_snapshot = engine.snapshot(); let (started_tx, started_rx) = mpsc::channel(); diff --git a/docs/api-reference.md b/docs/api-reference.md index dd2e790..968ffb1 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -123,7 +123,7 @@ See [Graph schema](graph-schema.md). ```python -CypherEngine(config: GraphConfig, datasets: dict, use_csr: bool = False) +CypherEngine(config: GraphConfig, datasets: dict, use_csr: bool = False, *, target_partitions: int | None = None) ``` Registers `datasets` once and reuses the catalog, DataFusion session and (with `use_csr=True`) CSR index @@ -138,10 +138,12 @@ See [Repeated queries](lance.md#repeated-queries) for the scope and limits. | `config` | The graph mapping. | | `datasets` | See [Datasets](#datasets). Lance datasets are pinned to the version passed in and not read here. | | `use_csr` | Build the experimental [CSR index](csr.md) for relationship types with integer endpoints, and node indexes for labels with unique integer keys. Default `False`. | +| `target_partitions` | Optional positive integer target for DataFusion query partitioning. `None` keeps DataFusion's default. Keyword-only; preserved by `refresh()`. It is not a thread or memory limit. See [Query parallelism](lance.md#query-parallelism). | Raises `ValueError` if `datasets` is empty, a configured key or endpoint column is missing, or a Lance dataset was opened through a namespace client and can't be reopened by URI. Raises `RuntimeError` if a dataset can't -be opened. +be opened. For `target_partitions`, booleans and non-integers other than `None` raise `TypeError`; +zero, negative values and integers outside the platform's unsigned integer range raise `ValueError`. `#!python execute(query: str, params: dict | None = None, *, timeout: float | None = None) -> pyarrow.Table` : Runs `query` and returns the whole result. `params` binds `$name` parameters: `int`, `float`, `str`, @@ -160,7 +162,8 @@ be opened. `#!python refresh() -> None` : Moves every Lance dataset to its latest version and updates the CSR, node indexes and statistics. All-or-nothing; running queries keep their snapshot. If another refresh publishes first, retries from - that snapshot so an older candidate cannot replace it. See [Lance datasets](lance.md#refresh). + that snapshot so an older candidate cannot replace it. Preserves `target_partitions`. + See [Lance datasets](lance.md#refresh). `#!python dataset_versions() -> dict[str, int]` : The pinned version of each Lance dataset. PyArrow inputs are not listed. diff --git a/docs/cypher.md b/docs/cypher.md index df005b4..210f248 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -40,12 +40,13 @@ engine = CypherEngine(config, datasets, use_csr=True) ## `CypherEngine` ```python -CypherEngine(config, datasets, use_csr=False) +CypherEngine(config, datasets, use_csr=False, target_partitions=None) ``` The constructor validates the datasets against the config, registers them with a DataFusion session, and, with `use_csr=True`, builds the [CSR index](csr.md). After that, every query reuses the same catalog and -session. +session. The optional keyword-only `target_partitions` sets a positive target for query partitioning; +`None` keeps the current default. See [Query parallelism](lance.md#query-parallelism) before tuning it. !!! tip "Share one engine" diff --git a/docs/lance.md b/docs/lance.md index 4e09671..5a6a83b 100644 --- a/docs/lance.md +++ b/docs/lance.md @@ -95,6 +95,28 @@ change the plan. Queries whose operators cannot produce a reusable template fall new snapshot, including a refresh whose datasets did not change; active readers retain their existing plans. These rules apply with and without CSR. They do not promise a speedup for a particular workload. +## Query parallelism + +`CypherEngine(..., target_partitions=None)` keeps DataFusion's default query partitioning. If you have +measured your workload, you can choose a positive target explicitly: + +```python +serial_engine = CypherEngine( + config, + {"Person": lance.dataset("person.lance"), "KNOWS": lance.dataset("knows.lance")}, + use_csr=True, + target_partitions=1, +) +print(serial_engine.execute("MATCH (p:Person) RETURN count(*) AS n").to_pylist()) # [{'n': 4}] +serial_engine.refresh() # retains target_partitions=1 +``` + +The target is a planning preference, not a thread or memory limit; operators may use a different number +of partitions. A lower target can help short queries while making larger joins and scans slower. Compare +your full workload in the storage mode you use before choosing a value. The default remains unchanged, +and `refresh()` preserves your choice. Booleans and non-integers other than `None` raise `TypeError`; +zero, negative and out-of-range integers raise `ValueError`. + ## Streaming large results `execute()` returns one table, held in memory. For big outputs, stream: diff --git a/examples/lance_snapshot_refresh.py b/examples/lance_snapshot_refresh.py index 1588a0c..efae887 100644 --- a/examples/lance_snapshot_refresh.py +++ b/examples/lance_snapshot_refresh.py @@ -46,6 +46,8 @@ def main() -> None: "Person": lance.dataset(person_uri), "KNOWS": lance.dataset(root / "knows.lance"), }, + # Optional: choose this target after measuring your own queries; refresh keeps it. + target_partitions=2, ) count = "MATCH (p:Person) RETURN count(*) AS n" diff --git a/python/python/lance_csr/_internal/graph.pyi b/python/python/lance_csr/_internal/graph.pyi index e6544f7..0c921a6 100644 --- a/python/python/lance_csr/_internal/graph.pyi +++ b/python/python/lance_csr/_internal/graph.pyi @@ -115,7 +115,12 @@ class CypherQuery: @final class CypherEngine: def __init__( - self, config: GraphConfig, datasets: dict[str, _DatasetT], use_csr: bool = False + self, + config: GraphConfig, + datasets: dict[str, _DatasetT], + use_csr: bool = False, + *, + target_partitions: int | None = None, ) -> None: ... def execute( self, diff --git a/python/python/tests/test_target_partitions.py b/python/python/tests/test_target_partitions.py new file mode 100644 index 0000000..bd608af --- /dev/null +++ b/python/python/tests/test_target_partitions.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""The optional partition target preserves query results and snapshot behavior.""" + +import inspect + +import lance +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture +def graph(): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + return config, { + "Person": pa.table({"id": [1, 2, 3, 4]}), + "KNOWS": pa.table({"src": [1, 1, 2], "dst": [2, 3, 4]}), + } + + +@pytest.mark.parametrize("partitions", [None, 1, 2]) +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize("storage", ["arrow", "lance"]) +def test_partition_targets_preserve_queries_and_refresh( + graph, partitions, use_csr, storage, tmp_path +): + config, data = graph + if storage == "lance": + data = { + name: lance.write_dataset(table, tmp_path / f"{name}.lance") + for name, table in data.items() + } + engine = CypherEngine(config, data, use_csr, target_partitions=partitions) + query = ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE a.id = $start " + "RETURN b.id AS id ORDER BY id" + ) + for start, ids in [(1, [2, 3]), (2, [4]), (1, [2, 3])]: + expected = [{"id": value} for value in ids] + assert engine.execute(query, {"start": start}).to_pylist() == expected + assert ( + engine.execute_batches(query, {"start": start}).read_all().to_pylist() + == expected + ) + assert engine.execute( + "MATCH (p:Person) RETURN count(*) AS n, sum(p.id) AS total" + ).to_pylist() == [{"n": 4, "total": 10}] + scan = "MATCH (p:Person) RETURN p.id AS id ORDER BY id" + reader = engine.execute_batches(scan) + if storage == "lance": + lance.write_dataset( + pa.table({"id": [5]}), tmp_path / "Person.lance", mode="append" + ) + engine.refresh() + assert reader.read_all().to_pylist() == [{"id": value} for value in range(1, 5)] + count = 5 if storage == "lance" else 4 + assert engine.execute(scan).to_pylist() == [ + {"id": value} for value in range(1, count + 1) + ] + + +@pytest.mark.parametrize("value", [True, False, 1.0, "1", [], {}]) +def test_partition_target_rejects_non_integers(graph, value): + config, data = graph + with pytest.raises(TypeError, match="target_partitions"): + CypherEngine(config, data, target_partitions=value) + + +@pytest.mark.parametrize("value", [0, -1, -(1 << 128), 1 << 128]) +def test_partition_target_rejects_nonpositive_and_overflow(graph, value): + config, data = graph + with pytest.raises(ValueError, match="target_partitions"): + CypherEngine(config, data, target_partitions=value) + + +def test_partition_target_is_keyword_only_and_default_is_none(graph): + config, data = graph + parameter = inspect.signature(CypherEngine).parameters["target_partitions"] + assert parameter.kind is inspect.Parameter.KEYWORD_ONLY + assert parameter.default is None + with pytest.raises(TypeError): + CypherEngine(config, data, False, 1) + assert CypherEngine(config, data, False).execute("RETURN 1 AS n").to_pylist() == [ + {"n": 1} + ] diff --git a/python/python/tests/typing/public_api.py b/python/python/tests/typing/public_api.py index c4e3704..27f2965 100644 --- a/python/python/tests/typing/public_api.py +++ b/python/python/tests/typing/public_api.py @@ -66,6 +66,8 @@ def public_api( CypherEngine(config, {"Person": batch}) CypherEngine(config, {"Person": reader}) CypherEngine(config, {"Person": dataset}) + CypherEngine(config, datasets, target_partitions=1) + CypherEngine(config, datasets, target_partitions=None) one = CypherQuery(query).with_config(config).with_parameter("id", 1) assert_type(one, CypherQuery) @@ -107,6 +109,8 @@ def public_api( assert_type(SqlEngine(datasets).execute(sql.sql()), Table) # Rejected calls are checked too: removing a type restriction makes these unused. + CypherEngine(config, datasets, target_partitions="1") # pyright: ignore[reportArgumentType] + CypherEngine(config, datasets, False, 1) # pyright: ignore[reportCallIssue] engine.execute(query, {"id": object()}) # pyright: ignore[reportArgumentType] vector.top_k("5") # pyright: ignore[reportArgumentType] one.with_parameter("ids", {1, 2}) # pyright: ignore[reportArgumentType] From 7461a7774d6cdf74bbd6961efa7ec86b229fa945 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 14:25:46 -0700 Subject: [PATCH 45/66] docs(python): locate stdlib for embedded bridge tests --- python/DEVELOPMENT.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/DEVELOPMENT.md b/python/DEVELOPMENT.md index 7f57095..d5e73a7 100644 --- a/python/DEVELOPMENT.md +++ b/python/DEVELOPMENT.md @@ -58,14 +58,16 @@ test executable needs an explicit link to the active Python library. From `pytho ```shell lance_python_library=$(uv run --no-sync python -c 'import pathlib, sysconfig; print(pathlib.Path(sysconfig.get_config_var("LIBDIR")) / sysconfig.get_config_var("LDLIBRARY"))') +lance_python_home=$(uv run --no-sync python -c 'import sys; print(sys.base_prefix)') lance_test_artifacts=$(mktemp) cargo rustc --manifest-path ../crates/lance-graph-python/Cargo.toml --tests --message-format=json -- -C "link-arg=$lance_python_library" > "$lance_test_artifacts" lance_test_binary=$(uv run --no-sync python -c 'import json, sys; artifacts = [json.loads(line) for line in open(sys.argv[1])]; print(next(a["executable"] for a in artifacts if a.get("reason") == "compiler-artifact" and a.get("profile", {}).get("test") and a.get("executable")))' "$lance_test_artifacts") -"$lance_test_binary" +PYTHONHOME="$lance_python_home" "$lance_test_binary" rm "$lance_test_artifacts" ``` -This uses Cargo's emitted executable path rather than selecting an older binary from `target/`. It does not +`PYTHONHOME` lets tests that initialize the interpreter find its standard library, including with a +relocatable Python installation. This uses Cargo's emitted executable path rather than selecting an older binary from `target/`. It does not change extension packaging or replace the Python tests. Use a separate `CARGO_TARGET_DIR` for each worktree when validating multiple revisions, so each test binary is built from that worktree's sources. From 28566db2f8af07e3b6eb3ba406290deddc84145a Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 14:25:11 -0700 Subject: [PATCH 46/66] docs(cypher): clarify unsigned math conversion Limit the Int64 overflow restriction to integer abs. The remaining math functions convert numeric values, including UInt64, to Float64; cover those Arrow inputs with real execution tests. Submitted by David Hughes (@stugorf), with AI Collaboration. --- docs/cypher.md | 4 +++- python/python/tests/test_math_functions.py | 13 +++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/docs/cypher.md b/docs/cypher.md index 210f248..3342df6 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -419,7 +419,9 @@ raise an error. Unsigned Arrow arguments outside Int64 range also raise an error converted to null. Null arguments return null. `abs` preserves integer results; the other math functions return floating point values. Integer -absolute-value overflow and unsigned inputs outside Int64 range raise an error. `round` accepts one argument, with ties toward positive +`abs` raises an error on absolute-value overflow or unsigned inputs outside Int64 range. The other +math functions convert numeric inputs, including UInt64, to Float64. `round` accepts one argument, +with ties toward positive infinity (`round(-1.5)` is `-1.0`), as specified in the [Neo4j 4.4 numeric reference](https://neo4j.com/docs/cypher-manual/4.4/functions/mathematical-numeric/). `sqrt` of a negative number returns NaN. The math functions return null for null input and reject diff --git a/python/python/tests/test_math_functions.py b/python/python/tests/test_math_functions.py index 4cb066c..49d90a2 100644 --- a/python/python/tests/test_math_functions.py +++ b/python/python/tests/test_math_functions.py @@ -117,6 +117,19 @@ def test_math_out_of_range_unsigned_integer(): engine.execute("MATCH (p:Person) RETURN abs(p.large) AS value") +@pytest.mark.parametrize("function", ["round", "floor", "ceil", "sqrt"]) +def test_math_unsigned_integer_converts_to_float(function): + config = GraphConfig.builder().node_table("Person", "id").build() + engine = CypherEngine( + config, + {"Person": pa.table({"id": [1], "large": pa.array([2**63], type=pa.uint64())})}, + ) + result = engine.execute(f"MATCH (p:Person) RETURN {function}(p.large) AS value") + expected = math.sqrt(2**63) if function == "sqrt" else float(2**63) + assert result.to_pylist() == [{"value": expected}] + assert result.schema.field("value").type == pa.float64() + + def test_math_non_finite_arrow_values(): config = GraphConfig.builder().node_table("Person", "id").build() engine = CypherEngine( From 063e15672cbe23298118f831f7bea5981726ed24 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 14:27:37 -0700 Subject: [PATCH 47/66] fix(cypher): preserve typed nulls in integrated SQL export Retain exact Arrow types with the existing DataFusion arrow_cast function after constant folding while keeping native list size/range lowering. Bare null stays untyped; other dialect typed-null exports fail explicitly. Include real execution regressions and public documentation. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-graph-python/src/graph.rs | 3 +- crates/lance-graph/src/query.rs | 61 ++++++++++++++++++++++++-- docs/api-reference.md | 2 +- docs/sql.md | 4 ++ python/python/tests/test_to_sql.py | 24 +++++++++- 5 files changed, 88 insertions(+), 6 deletions(-) diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index f53246e..66c260d 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -612,7 +612,8 @@ impl CypherQuery { /// Raises /// ------ /// NotImplementedError - /// If a runtime-only Cypher function has no faithful SQL translation + /// If a runtime-only Cypher function has no faithful SQL translation, or a + /// typed-null literal is exported outside the Default dialect /// RuntimeError /// If SQL generation fails #[pyo3(signature = (datasets, dialect=None))] diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 17dbd91..632b7ae 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -477,17 +477,18 @@ impl CypherQuery { /// /// # Errors /// Returns `GraphError::UnsupportedFeature` for runtime-only Cypher functions - /// without a faithful SQL translation. Constant expressions may be folded first. + /// without a faithful SQL translation or typed nulls outside the Default dialect. + /// Constant expressions may be folded first. pub async fn to_sql( &self, datasets: HashMap, dialect: Option, ) -> Result { use datafusion::common::tree_node::{Transformed, TreeNode}; - use datafusion::common::{not_impl_err, DFSchema}; + use datafusion::common::{not_impl_err, DFSchema, ScalarValue}; use datafusion::logical_expr::expr_rewriter::NamePreserver; use datafusion::logical_expr::registry::FunctionRegistry; - use datafusion::logical_expr::{utils::merge_schema, Expr, LogicalPlan}; + use datafusion::logical_expr::{lit, utils::merge_schema, Expr, LogicalPlan}; use std::sync::Arc; let dialect = dialect.unwrap_or_default(); @@ -534,6 +535,21 @@ impl CypherQuery { ); } } + if let Expr::Literal(value, _) = &expr { + let data_type = value.data_type(); + if value.is_null() && !data_type.is_null() { + let type_name = data_type.to_string(); + if dialect != SqlDialect::Default + || !type_name.parse::().is_ok_and(|parsed| parsed == data_type) + { + return not_impl_err!("SQL export of typed null {data_type} requires a faithful Default-dialect Arrow cast"); + } + // The unparser emits bare NULL for typed nulls, losing their Arrow type. + return Ok(Transformed::yes(datafusion::functions::core::arrow_cast().call(vec![ + lit(ScalarValue::Null), lit(type_name), + ]))); + } + } Ok(Transformed::no(expr)) })?; Ok(rewritten.update_data(|expr| name.restore(expr))) @@ -2584,6 +2600,45 @@ mod tests { .contains("SQL export of Cypher function 'round'")); } + #[tokio::test] + async fn test_to_sql_preserves_constant_null_types() { + let batch = RecordBatch::try_from_iter(vec![( + "id", + Arc::new(arrow_array::Int64Array::from(vec![1])) as arrow_array::ArrayRef, + )]) + .unwrap(); + let datasets = HashMap::from([("Person".to_string(), batch)]); + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + for expression in [ + "round(null)", + "abs(null)", + "sqrt(null)", + "floor(null)", + "ceil(null)", + "toUpper(null)", + "null", + ] { + let query = CypherQuery::new(&format!("MATCH (p:Person) RETURN {expression} AS value")) + .unwrap() + .with_config(config.clone()); + let cypher = query.execute(datasets.clone(), None).await.unwrap(); + let sql_text = query.to_sql(datasets.clone(), None).await.unwrap(); + let sql = crate::sql_query::SqlQuery::new(&sql_text) + .execute(datasets.clone()) + .await + .unwrap(); + assert_eq!( + sql.schema().field(0).data_type(), + cypher.schema().field(0).data_type(), + "{expression}: {sql_text}" + ); + assert_eq!(sql.column(0).logical_null_count(), 1, "{expression}"); + } + } + async fn write_lance_dataset(path: &std::path::Path, batch: arrow_array::RecordBatch) { use arrow_array::{RecordBatch, RecordBatchIterator}; use lance::dataset::{Dataset, WriteParams}; diff --git a/docs/api-reference.md b/docs/api-reference.md index 968ffb1..a2afa3d 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -215,7 +215,7 @@ A parsed Cypher query. Raises `ValueError` if `query_text` doesn't parse. Builde `#!python to_sql(datasets: dict, dialect: SqlDialect | None = None) -> str` : Exports supported expressions as SQL joins. Raises `NotImplementedError` for runtime-only Cypher - functions without a faithful SQL translation. See [From Cypher to SQL](sql.md#from-cypher-to-sql). + functions without a faithful SQL translation, and for typed-null literals outside the Default dialect. See [From Cypher to SQL](sql.md#from-cypher-to-sql). `#!python execute_with_vector_rerank(datasets: dict, vector_search: VectorSearch) -> pyarrow.Table` : Runs the query, then reranks by vector similarity. diff --git a/docs/sql.md b/docs/sql.md index e2cd70e..85a5b1c 100644 --- a/docs/sql.md +++ b/docs/sql.md @@ -89,6 +89,10 @@ table values. Cypher `round(-1.5)` produces `-1.0`; DataFusion SQL `round(-1.5)` Integer `abs` also normalizes its result to Int64 in Cypher, while SQL preserves the input width. Use `CypherQuery.execute` or `CypherEngine.execute` for these expressions. +Constant-folded typed nulls use DataFusion `arrow_cast` in the Default dialect so their Arrow +result types are preserved. Other dialects raise `NotImplementedError` for typed-null literals; +a bare Cypher `null` remains an untyped SQL `NULL`. + List `size` and inclusive `range` retain SQL export support through native expressions, including null arguments, integer casts and the optional range step. Constant expressions that the optimizer evaluates completely can also export their resulting literals. These checks establish DataFusion diff --git a/python/python/tests/test_to_sql.py b/python/python/tests/test_to_sql.py index db86e81..78b9614 100644 --- a/python/python/tests/test_to_sql.py +++ b/python/python/tests/test_to_sql.py @@ -5,7 +5,7 @@ import pyarrow as pa import pytest -from lance_csr import CypherQuery, GraphConfig, SqlEngine +from lance_csr import CypherQuery, GraphConfig, SqlDialect, SqlEngine @pytest.fixture @@ -39,6 +39,13 @@ def scalar_export_env(): "size(range(1, p.end))", "toUpper(p.text)", "round(-1.5)", + "round(null)", + "abs(null)", + "sqrt(null)", + "floor(null)", + "ceil(null)", + "toUpper(null)", + "null", "left('a猫😀z', 3)", ], ) @@ -89,6 +96,21 @@ def test_sql_export_rejects_runtime_only_cypher_functions( assert result.schema.field("value").type == pa.int64() +@pytest.mark.parametrize( + "dialect", + [SqlDialect.Spark, SqlDialect.PostgreSql, SqlDialect.MySql, SqlDialect.Sqlite], +) +def test_sql_export_typed_null_requires_default_dialect(scalar_export_env, dialect): + config, datasets = scalar_export_env + query = CypherQuery("MATCH (p:Person) RETURN round(null) AS value").with_config( + config + ) + with pytest.raises(NotImplementedError, match="SQL export of typed null Float64"): + query.to_sql(datasets, dialect) + bare_null = CypherQuery("MATCH (p:Person) RETURN null AS value").with_config(config) + assert "NULL" in bare_null.to_sql(datasets, dialect) + + def test_sql_export_checks_filter_functions(scalar_export_env): config, datasets = scalar_export_env query = CypherQuery( From bfbbb24db2fdd2753b959ca775d5c85eaea75cae Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 14:29:45 -0700 Subject: [PATCH 48/66] docs(cypher): explain path and CALL AST migrations --- crates/lance-graph/README.md | 18 ++++++++++++++++++ crates/lance-graph/src/logical_plan.rs | 11 +++++------ 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/crates/lance-graph/README.md b/crates/lance-graph/README.md index 8a7c8cb..5124844 100644 --- a/crates/lance-graph/README.md +++ b/crates/lance-graph/README.md @@ -144,6 +144,24 @@ flag; each branch has its own variable scope. Exhaustive matches on `LogicalOper new `SingleRow` and `QueryUnion` variants. These are Rust source compatibility changes; Python method signatures are unchanged by UNION. +Named ordinary paths add `PathPattern.variable`; set it to `None` in existing unnamed-path struct +literals, or `Some(name)` for `p = (a)-[:R]->(b)`. Shortest-path names still use +`ShortestPathSpec.variable`. `LogicalOperator::PathMatch` retains ordered node and relationship +records, so visitors must traverse its input and path pattern instead of treating a path value as a +length. The public `nodes(path)` result remains a list of integer keys; the Cypher guide describes the +separate Arrow record returned by a path value. + +Read-only subqueries add `ReadingClause::Call(Box)` and +`LogicalOperator::{Apply, CallInput}`. Update exhaustive matches and visitors to recurse into the CALL +query, including its UNION branches. A leading importing `WITH` declares the outer variables available +in that query; each CALL body otherwise has its own scope. + +List folds add `ValueExpression::{Reduce, ListPredicate, LocalVariable}`. Update exhaustive matches +and visitors to traverse the input and body while respecting iterator and accumulator scope. Prefer +parsing Cypher text to construct these expressions: the parser marks local variable references so they +cannot accidentally bind to graph variables with the same name. These path, CALL and fold additions +also require updates from Rust consumers that construct ASTs or exhaustively match public enums. + ## Testing ```bash diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index 75e1097..a15fbce 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -143,14 +143,14 @@ pub enum LogicalOperator { /// same variables, possibly with different labels; columns a branch lacks are null. Union { inputs: Vec }, - /// `shortestPath` / `allShortestPaths` between two variables bound in `input`; adds - /// `__length` and `__nodes` (node ids, source first). /// A named path retaining actual relationship occurrences and typed node records. PathMatch { input: Box, path: PathPattern, }, + /// `shortestPath` / `allShortestPaths` between two variables bound in `input`; adds + /// `__length` and `__nodes` (node ids, source first). ShortestPath { input: Box, source_variable: String, @@ -684,10 +684,9 @@ impl<'a> LogicalPlanner<'a> { }) } - /// 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. + /// Resolve node variables to ids in scalar expressions while retaining path values. + /// Path functions and null checks are lowered according to the path representation. + /// Bare WITH and RETURN items, which carry whole nodes, are not resolved. fn resolve_value(&self, expr: &mut ValueExpression) { match expr { ValueExpression::Variable(var) => { From 3a6039b4d841be3f3e51730b9f40461231715972 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 18:59:17 -0700 Subject: [PATCH 49/66] perf(csr): inline neighbor cursor lookups --- crates/lance-graph/src/csr_expand.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index 0e8a6b3..dede8ef 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -230,6 +230,7 @@ impl Adjacency { self.neighbors_from(key, &mut 0).for_each(f); } + #[inline] fn base_neighbors(&self, key: i64) -> &[u64] { self.base .dense @@ -238,6 +239,7 @@ impl Adjacency { } /// Resume at an occurrence offset, borrowing all neighbour and deletion storage. + #[inline] fn neighbors_from<'a>( &'a self, key: i64, From ad3b35ab990cc71a1e17a15c770ca0687dc1a88d Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 19:41:42 -0700 Subject: [PATCH 50/66] fix(query): stop Cartesian work at deadline checkpoints --- crates/lance-graph-python/src/executor.rs | 4 + crates/lance-graph/src/csr_expand.rs | 94 +++++++++------- crates/lance-graph/src/query.rs | 102 +++++++++++++++++ crates/lance-graph/src/query/deadline.rs | 127 ++++++++++++++++++++++ crates/lance-graph/src/shortest_path.rs | 2 +- docs/api-reference.md | 4 +- docs/csr.md | 4 +- docs/cypher.md | 3 +- 8 files changed, 299 insertions(+), 41 deletions(-) create mode 100644 crates/lance-graph/src/query/deadline.rs diff --git a/crates/lance-graph-python/src/executor.rs b/crates/lance-graph-python/src/executor.rs index 6e619ad..1fe65ac 100644 --- a/crates/lance-graph-python/src/executor.rs +++ b/crates/lance-graph-python/src/executor.rs @@ -117,6 +117,10 @@ impl BackgroundExecutor { F::Output: Send, { deadline.check()?; + let future = lance_graph::query::with_execution_deadline( + deadline.instant().map(tokio::time::Instant::into_std), + future, + ); let future = Self::result_or_interrupt(future, deadline); if let Some(py) = py { py.detach(move || self.runtime.block_on(future)) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index dede8ef..fa48783 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -245,6 +245,15 @@ impl Adjacency { key: i64, offset: &'a mut usize, ) -> impl Iterator + 'a { + self.candidates_from(key, offset).flatten() + } + + /// Keep tombstones visible to cancellation without repeating lookups for every edge. + fn candidates_from<'a>( + &'a self, + key: i64, + offset: &'a mut usize, + ) -> impl Iterator> + 'a { let base = self.base_neighbors(key); let added = self.added.get(&key).map_or(&[][..], Vec::as_slice); let removed = self.removed.get(&key).map_or(&[][..], Vec::as_slice); @@ -253,7 +262,7 @@ impl Adjacency { .chain(added) .enumerate() .skip(*offset) - .filter_map(move |(position, &neighbor)| { + .map(move |(position, &neighbor)| { *offset = position + 1; if removed.first() == Some(&position) { removed = &removed[1..]; @@ -264,23 +273,6 @@ impl Adjacency { }) } - /// Inspect one occurrence so cancelled work also stops through long tombstone runs. - fn next_candidate(&self, key: i64, offset: &mut usize) -> Option> { - let base = self.base_neighbors(key); - let position = *offset; - let neighbor = if position < base.len() { - base[position] - } else { - *self.added.get(&key)?.get(position - base.len())? - }; - *offset += 1; - let removed = self - .removed - .get(&key) - .is_some_and(|r| r.binary_search(&position).is_ok()); - Some((!removed).then_some(neighbor)) - } - /// Live edges: base + appended - removed. pub fn num_edges(&self) -> u64 { self.base.csr.num_edges() + self.added_count - self.removed_count @@ -848,34 +840,28 @@ impl Neighbors { key: i64, position: &'a mut NeighborPosition, ) -> impl Iterator + 'a { - let forward = self.forward.neighbors_from(key, &mut position.forward); + self.candidates_from(key, position).flatten() + } + + /// Borrow each adjacency once while exposing tombstones and suppressed self-loops as work. + pub(crate) fn candidates_from<'a>( + &'a self, + key: i64, + position: &'a mut NeighborPosition, + ) -> impl Iterator> + 'a { + let forward = self.forward.candidates_from(key, &mut position.forward); let reverse = self .reverse .as_ref() - .map(|r| r.neighbors_from(key, &mut position.reverse)); + .map(|r| r.candidates_from(key, &mut position.reverse)); forward.chain( reverse .into_iter() .flatten() - .filter(move |&n| n as i64 != key), + .map(move |candidate| candidate.filter(|&n| n as i64 != key)), ) } - /// Inspect one raw edge, including tombstones and suppressed reverse self-loops. - pub(crate) fn next_candidate( - &self, - key: i64, - position: &mut NeighborPosition, - ) -> Option> { - if let Some(candidate) = self.forward.next_candidate(key, &mut position.forward) { - return Some(candidate); - } - self.reverse - .as_ref()? - .next_candidate(key, &mut position.reverse) - .map(|candidate| candidate.filter(|&n| n as i64 != key)) - } - /// Edges visited over all keys (self-loops count once per direction in this estimate). pub fn num_edges(&self) -> u64 { self.forward.num_edges() + self.reverse.as_ref().map_or(0, |r| r.num_edges()) @@ -1249,7 +1235,7 @@ impl BatchExpansion { if self.keys.is_valid(i) && self.bound.as_ref().is_none_or(|b| b.is_valid(i)) { let key = self.keys.value(i); let required = self.bound.as_ref().map(|b| b.value(i)); - while let Some(candidate) = adjacency.next_candidate(key, &mut self.position) { + for candidate in adjacency.candidates_from(key, &mut self.position) { work += 1; if work.is_multiple_of(1024) { cooperative_yield().await; @@ -2090,6 +2076,40 @@ mod tests { assert!(expansion.position.forward > 0); } + #[tokio::test] + async fn expansion_yields_and_resumes_through_a_long_tombstone_run() { + let n = 5_000; + let base = Adjacency::build( + edge_stream(vec![Some(1); n], (2..n as i64 + 2).map(Some).collect()), + "src", + "dst", + ) + .await + .unwrap(); + let removed = (2..n as i64 + 1).map(|id| (1, id)).collect::>(); + let adjacency = Neighbors::directed(Arc::new(base.with_changes(&[], &removed))); + let input = + RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(vec![1])) as ArrayRef)]) + .unwrap(); + let mut expansion = BatchExpansion::try_new(input, 0, None).unwrap(); + let mut next = Box::pin(expansion.next_chunk(&adjacency, None, 1)); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert_eq!(expansion.position.forward, 1024); + assert_eq!( + expansion + .next_chunk(&adjacency, None, 1) + .await + .neighbor_values, + vec![n as i64 + 1] + ); + assert!(expansion + .next_chunk(&adjacency, None, 1) + .await + .rows + .is_empty()); + } + #[tokio::test] async fn high_degree_expansion_visits_only_the_first_chunk() { let n = 100_000; diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 632b7ae..63e95e9 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -3,6 +3,11 @@ //! High-level Cypher query interface for Lance datasets +mod deadline; + +#[doc(hidden)] +pub use deadline::with_execution_deadline; + use crate::ast::CypherQuery as CypherAST; use crate::ast::ReadingClause; use crate::config::GraphConfig; @@ -791,6 +796,7 @@ impl CypherQuery { ctx: datafusion::execution::context::SessionContext, ) -> Result { let (_logical_plan, df_logical_plan) = self.create_logical_plans(catalog)?; + let ctx = deadline::context(ctx); ctx.execute_logical_plan(df_logical_plan) .await .map_err(|e| GraphError::ExecutionError { @@ -1140,6 +1146,7 @@ impl CypherQuery { let (logical_plan, df_logical_plan) = self.create_logical_plans(catalog)?; // Phase 4: DataFusion Physical Plan + let ctx = deadline::context(ctx.clone()); let df = ctx .execute_logical_plan(df_logical_plan.clone()) .await @@ -1621,6 +1628,101 @@ mod tests { } } + #[tokio::test] + async fn cartesian_stream_yields_and_drops_before_full_collection() { + use arrow_array::{ArrayRef, Int64Array}; + use datafusion::datasource::{DefaultTableSource, MemTable}; + use datafusion::execution::context::{SessionConfig, SessionContext}; + use futures::TryStreamExt; + use lance_graph_catalog::InMemoryCatalog; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let batch = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from_iter_values(0..1000)) as ArrayRef, + )]) + .unwrap(); + let table = Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch]]).unwrap()); + let catalog = Arc::new( + InMemoryCatalog::new() + .with_node_source("Person", Arc::new(DefaultTableSource::new(table))), + ); + let ctx = SessionContext::new_with_config( + SessionConfig::new() + .with_target_partitions(1) + .with_round_robin_repartition(false), + ); + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + let query = CypherQuery::new( + "MATCH (a:Person), (b:Person) RETURN a.id AS left_id, b.id AS right_id", + ) + .unwrap() + .with_config(config.clone()); + let stream = with_execution_deadline( + Some(std::time::Instant::now() + std::time::Duration::from_secs(60)), + query.execute_stream_with_catalog_and_context(catalog.clone(), ctx.clone()), + ) + .await + .unwrap(); + // A fresh Tokio task supplies its own budget; no inherited source yield can pass this test. + tokio::spawn(async move { + let visited = Arc::new(AtomicUsize::new(0)); + let progress = visited.clone(); + let mut read = Box::pin(stream.try_fold(0_usize, move |count, batch| { + progress.fetch_add(batch.num_rows(), Ordering::SeqCst); + futures::future::ready(Ok(count + batch.num_rows())) + })); + assert!(futures::poll!(read.as_mut()).is_pending()); + let stopped = visited.load(Ordering::SeqCst); + assert!(stopped > 0 && stopped < 1_000_000); + drop(read); + tokio::task::yield_now().await; + assert_eq!(visited.load(Ordering::SeqCst), stopped); + }) + .await + .unwrap(); + let result = query + .execute_with_catalog_and_context(catalog.clone(), ctx.clone()) + .await + .unwrap(); + assert_eq!(result.num_rows(), 1_000_000); + assert_eq!(result.num_columns(), 2); + + let filtered = CypherQuery::new( + "MATCH (a:Person), (b:Person) WHERE a.id + b.id >= 999 RETURN count(*) AS n", + ) + .unwrap() + .with_config(config); + let expired = with_execution_deadline( + Some(std::time::Instant::now() - std::time::Duration::from_secs(1)), + filtered.execute_with_catalog_and_context(catalog.clone(), ctx.clone()), + ) + .await + .unwrap_err(); + assert!(expired.to_string().contains("Query timeout exceeded")); + for timed in [false, true] { + let result = with_execution_deadline( + timed.then(|| std::time::Instant::now() + std::time::Duration::from_secs(60)), + filtered.execute_with_catalog_and_context(catalog.clone(), ctx.clone()), + ) + .await + .unwrap(); + assert_eq!(result.num_rows(), 1); + assert_eq!( + result + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 500_500 + ); + } + } + #[test] fn test_parse_simple_cypher_query() { let query = CypherQuery::new("MATCH (n:Person) RETURN n.name").unwrap(); diff --git a/crates/lance-graph/src/query/deadline.rs b/crates/lance-graph/src/query/deadline.rs new file mode 100644 index 0000000..0633eff --- /dev/null +++ b/crates/lance-graph/src/query/deadline.rs @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Internal deadline propagation from the Python bridge into physical join work. + +use std::future::Future; +use std::sync::Arc; +use std::time::Instant; + +use arrow_schema::{DataType, Field, Schema}; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::{DataFusionError, Result, ScalarValue}; +use datafusion::execution::context::SessionContext; +use datafusion::execution::session_state::SessionStateBuilder; +use datafusion::logical_expr::{create_udf, ColumnarValue, Operator, Volatility}; +use datafusion::physical_expr::expressions::BinaryExpr; +use datafusion::physical_expr::{PhysicalExpr, ScalarFunctionExpr}; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::coop::CooperativeExec; +use datafusion::physical_plan::joins::utils::JoinFilter; +use datafusion::physical_plan::joins::{ + CrossJoinExec, NestedLoopJoinExec, NestedLoopJoinExecBuilder, +}; +use datafusion::physical_plan::ExecutionPlan; + +tokio::task_local! { + static EXECUTION_DEADLINE: Option; +} + +/// Internal cross-crate bridge plumbing; no deadline is published into an engine snapshot. +#[doc(hidden)] +pub async fn with_execution_deadline(deadline: Option, future: F) -> F::Output { + EXECUTION_DEADLINE.scope(deadline, future).await +} + +pub(super) fn context(ctx: SessionContext) -> SessionContext { + let Some(deadline) = EXECUTION_DEADLINE + .try_with(|deadline| *deadline) + .ok() + .flatten() + else { + return ctx; + }; + SessionContext::new_with_state( + SessionStateBuilder::new_from_existing(ctx.state()) + .with_physical_optimizer_rule(Arc::new(DeadlineJoins { deadline })) + .build(), + ) +} + +#[derive(Debug)] +struct DeadlineJoins { + deadline: Instant, +} + +impl PhysicalOptimizerRule for DeadlineJoins { + fn name(&self) -> &str { + "query_deadline_joins" + } + + fn schema_check(&self) -> bool { + true + } + + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> Result> { + let deadline = self.deadline; + let function = create_udf( + "__query_deadline", + vec![], + DataType::Boolean, + Volatility::Volatile, + Arc::new(move |_| { + if Instant::now() >= deadline { + return Err(DataFusionError::Execution("Query timeout exceeded".into())); + } + Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(true)))) + }), + ); + let guard: Arc = Arc::new(ScalarFunctionExpr::new( + "__query_deadline", + Arc::new(function), + vec![], + Arc::new(Field::new("__query_deadline", DataType::Boolean, false)), + Arc::new(config.clone()), + )); + plan.transform_up(|plan| { + if let Some(join) = plan.as_any().downcast_ref::() { + // Output-only yielding cannot interrupt long runs of rejected join candidates. + let filter = match join.filter() { + Some(filter) => JoinFilter::new( + Arc::new(BinaryExpr::new( + guard.clone(), + Operator::And, + filter.expression().clone(), + )), + filter.column_indices().to_vec(), + filter.schema().clone(), + ), + None => JoinFilter::new(guard.clone(), vec![], Arc::new(Schema::empty())), + }; + let joined = NestedLoopJoinExecBuilder::new( + join.left().clone(), + join.right().clone(), + *join.join_type(), + ) + .with_filter(Some(filter)) + .with_projection_ref(join.projection().clone()) + .build()?; + Ok(Transformed::yes( + Arc::new(CooperativeExec::new(Arc::new(joined))) as Arc, + )) + } else if plan.as_any().is::() { + Ok(Transformed::yes( + Arc::new(CooperativeExec::new(plan)) as Arc + )) + } else { + Ok(Transformed::no(plan)) + } + }) + .map(|plan| plan.data) + } +} diff --git a/crates/lance-graph/src/shortest_path.rs b/crates/lance-graph/src/shortest_path.rs index 024889c..a5d5397 100644 --- a/crates/lance-graph/src/shortest_path.rs +++ b/crates/lance-graph/src/shortest_path.rs @@ -423,7 +423,7 @@ async fn expand_level( let mut work = 0_usize; for &node in frontier { let mut position = crate::csr_expand::NeighborPosition::default(); - while let Some(candidate) = adjacency.next_candidate(node, &mut position) { + for candidate in adjacency.candidates_from(node, &mut position) { work += 1; if work.is_multiple_of(1024) { crate::csr_expand::cooperative_yield().await; diff --git a/docs/api-reference.md b/docs/api-reference.md index a2afa3d..da27a13 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -75,7 +75,9 @@ between reads. Before returning a reader, expiry raises `TimeoutError`. Reading reader raises `pyarrow.ArrowInvalid` containing `Query timeout exceeded`. Earlier batches remain valid. Expiry or reader close releases the stream even while the caller is idle; the engine remains usable. -Cancellation is cooperative. Traversal loops yield at bounded work intervals, but synchronous parsing, +Cancellation is cooperative. Traversal loops yield at bounded work intervals, and timed cross joins +yield between output batches. Timed nested-loop joins check the deadline while evaluating candidates, +including batches that produce no matches. Synchronous parsing, Python data providers and foreign blocking calls cannot be forcibly preempted. Deadlines are checked between planning stages and before accepting completed results; this is not a hard real-time limit. diff --git a/docs/csr.md b/docs/csr.md index 97e2244..4d01a3d 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -95,7 +95,9 @@ Node index output chunk of neighbour positions before fetching properties. It yields that chunk before visiting the remaining neighbours, even for a single high-degree node. Deleted edge positions are read from the shared snapshot rather than copied for each query; refresh resolves those positions by scanning the - affected sources' neighbour lists. The chunk row limit comes from DataFusion's batch + affected sources' neighbour lists. Cancellation checkpoints count deleted edges and suppressed + reverse self-loops as work, using borrowed adjacency slices while a chunk is being collected. + The chunk row limit comes from DataFusion's batch size; wide properties and multiple active partitions still increase working memory. This does not bound the resident graph, the input data, or buffering in joins, sorts and aggregations elsewhere in the plan. diff --git a/docs/cypher.md b/docs/cypher.md index 3342df6..88abd89 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -139,7 +139,8 @@ except TimeoutError: Initial expiry raises `TimeoutError`. Reading an expired native Arrow reader raises `pyarrow.ArrowInvalid` with `Query timeout exceeded`; partial output already read remains valid. Close readers when you stop -consuming them. See the [API reference](api-reference.md) for validation and cooperative cancellation limits. +consuming them. Timed cross joins yield between batches, and nested-loop joins check the deadline +while evaluating candidates, including when no rows match. See the [API reference](api-reference.md) for validation and cooperative cancellation limits. ### `explain(query, params=None)` From 2a9357a669e90f5e47c5a2ff28abee0b64563270 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 19:46:43 -0700 Subject: [PATCH 51/66] fix(query): wire existing runtime and physical plan APIs --- crates/lance-graph-python/Cargo.lock | 1 + crates/lance-graph/Cargo.toml | 1 + crates/lance-graph/src/query/deadline.rs | 4 ++-- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/lance-graph-python/Cargo.lock b/crates/lance-graph-python/Cargo.lock index 248dee7..45e9ed2 100644 --- a/crates/lance-graph-python/Cargo.lock +++ b/crates/lance-graph-python/Cargo.lock @@ -3983,6 +3983,7 @@ dependencies = [ "serde", "serde_json", "snafu 0.8.9", + "tokio", ] [[package]] diff --git a/crates/lance-graph/Cargo.toml b/crates/lance-graph/Cargo.toml index aa40f55..f280a76 100644 --- a/crates/lance-graph/Cargo.toml +++ b/crates/lance-graph/Cargo.toml @@ -41,6 +41,7 @@ nom = "7.1" serde = { version = "1", features = ["derive"] } serde_json = "1" snafu = "0.8" +tokio = { version = "1.37", features = ["rt"] } [dev-dependencies] futures = "0.3" diff --git a/crates/lance-graph/src/query/deadline.rs b/crates/lance-graph/src/query/deadline.rs index 0633eff..d882c44 100644 --- a/crates/lance-graph/src/query/deadline.rs +++ b/crates/lance-graph/src/query/deadline.rs @@ -89,7 +89,7 @@ impl PhysicalOptimizerRule for DeadlineJoins { Arc::new(config.clone()), )); plan.transform_up(|plan| { - if let Some(join) = plan.as_any().downcast_ref::() { + if let Some(join) = plan.downcast_ref::() { // Output-only yielding cannot interrupt long runs of rejected join candidates. let filter = match join.filter() { Some(filter) => JoinFilter::new( @@ -114,7 +114,7 @@ impl PhysicalOptimizerRule for DeadlineJoins { Ok(Transformed::yes( Arc::new(CooperativeExec::new(Arc::new(joined))) as Arc, )) - } else if plan.as_any().is::() { + } else if plan.is::() { Ok(Transformed::yes( Arc::new(CooperativeExec::new(plan)) as Arc )) From 1d5d2834c24e5b21bd0973fb09bf1126c80d4bfc Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 19:49:41 -0700 Subject: [PATCH 52/66] fix(query): share one monotonic clock through deadline planning --- crates/lance-graph-python/src/executor.rs | 5 +---- crates/lance-graph/Cargo.toml | 2 +- crates/lance-graph/src/query.rs | 6 +++--- crates/lance-graph/src/query/deadline.rs | 2 +- 4 files changed, 6 insertions(+), 9 deletions(-) diff --git a/crates/lance-graph-python/src/executor.rs b/crates/lance-graph-python/src/executor.rs index 1fe65ac..08c16d7 100644 --- a/crates/lance-graph-python/src/executor.rs +++ b/crates/lance-graph-python/src/executor.rs @@ -117,10 +117,7 @@ impl BackgroundExecutor { F::Output: Send, { deadline.check()?; - let future = lance_graph::query::with_execution_deadline( - deadline.instant().map(tokio::time::Instant::into_std), - future, - ); + let future = lance_graph::query::with_execution_deadline(deadline.instant(), future); let future = Self::result_or_interrupt(future, deadline); if let Some(py) = py { py.detach(move || self.runtime.block_on(future)) diff --git a/crates/lance-graph/Cargo.toml b/crates/lance-graph/Cargo.toml index f280a76..52a93e4 100644 --- a/crates/lance-graph/Cargo.toml +++ b/crates/lance-graph/Cargo.toml @@ -41,7 +41,7 @@ nom = "7.1" serde = { version = "1", features = ["derive"] } serde_json = "1" snafu = "0.8" -tokio = { version = "1.37", features = ["rt"] } +tokio = { version = "1.37", features = ["rt", "time"] } [dev-dependencies] futures = "0.3" diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 63e95e9..38488c7 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -1662,7 +1662,7 @@ mod tests { .unwrap() .with_config(config.clone()); let stream = with_execution_deadline( - Some(std::time::Instant::now() + std::time::Duration::from_secs(60)), + Some(tokio::time::Instant::now() + std::time::Duration::from_secs(60)), query.execute_stream_with_catalog_and_context(catalog.clone(), ctx.clone()), ) .await @@ -1697,7 +1697,7 @@ mod tests { .unwrap() .with_config(config); let expired = with_execution_deadline( - Some(std::time::Instant::now() - std::time::Duration::from_secs(1)), + Some(tokio::time::Instant::now() - std::time::Duration::from_secs(1)), filtered.execute_with_catalog_and_context(catalog.clone(), ctx.clone()), ) .await @@ -1705,7 +1705,7 @@ mod tests { assert!(expired.to_string().contains("Query timeout exceeded")); for timed in [false, true] { let result = with_execution_deadline( - timed.then(|| std::time::Instant::now() + std::time::Duration::from_secs(60)), + timed.then(|| tokio::time::Instant::now() + std::time::Duration::from_secs(60)), filtered.execute_with_catalog_and_context(catalog.clone(), ctx.clone()), ) .await diff --git a/crates/lance-graph/src/query/deadline.rs b/crates/lance-graph/src/query/deadline.rs index d882c44..cb6b26c 100644 --- a/crates/lance-graph/src/query/deadline.rs +++ b/crates/lance-graph/src/query/deadline.rs @@ -5,7 +5,7 @@ use std::future::Future; use std::sync::Arc; -use std::time::Instant; +use tokio::time::Instant; use arrow_schema::{DataType, Field, Schema}; use datafusion::common::config::ConfigOptions; From 8ab8bf250a7d4f8fdc022a5e9ea6b7f127cf4f86 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 19:06:10 -0700 Subject: [PATCH 53/66] perf(paths): keep null checks on compact shortest plans --- crates/lance-graph/src/path_value.rs | 50 +++++++++++++++++++++++- docs/cypher.md | 2 +- python/python/tests/test_engine_paths.py | 29 ++++++++++++++ 3 files changed, 79 insertions(+), 2 deletions(-) diff --git a/crates/lance-graph/src/path_value.rs b/crates/lance-graph/src/path_value.rs index a1b35be..91c3259 100644 --- a/crates/lance-graph/src/path_value.rs +++ b/crates/lance-graph/src/path_value.rs @@ -4,7 +4,7 @@ //! Arrow values for paths and their entities. Properties live in a marked child so //! user property names cannot collide with identity or endpoint metadata. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use arrow::array::{Array, ArrayRef, AsArray, Int64Array, ListArray, StructArray, UInt32Array}; @@ -123,6 +123,23 @@ pub(crate) fn paths_array( ))) } +/// Null tests need only the compact path's validity, not its entity inventory. +fn mask_path_null_checks(predicate: &mut crate::ast::BooleanExpression, names: &HashSet) { + use crate::ast::{BooleanExpression as B, PropertyValue, ValueExpression as V}; + match predicate { + B::IsNull(value) | B::IsNotNull(value) if matches!(value, V::Variable(name) if names.contains(&name.to_lowercase())) => + { + *value = V::Literal(PropertyValue::Null); + } + B::And(left, right) | B::Or(left, right) => { + mask_path_null_checks(left, names); + mask_path_null_checks(right, names); + } + B::Not(inner) => mask_path_null_checks(inner, names), + _ => {} + } +} + /// Homogeneous length/nodes queries, including WITH aliases, keep their CSR plan. /// Returned path values, relationship access and containers need actual path records. pub(crate) fn needs_values( @@ -217,6 +234,22 @@ pub(crate) fn needs_values( let mut inspect = |expression: &V| { let mut value = expression.clone(); crate::ast::transform_value(&mut value, &mut |v| match v { + V::Boolean(predicate) => mask_path_null_checks(predicate, &names), + V::Case { whens, .. } => { + for (predicate, _) in whens { + mask_path_null_checks(predicate, &names); + } + } + V::ListComprehension { + filter: Some(predicate), + .. + } + | V::PatternComprehension { + filter: Some(predicate), + .. + } => { + mask_path_null_checks(predicate, &names); + } V::ScalarFunction { name, args } if matches!(name.to_lowercase().as_str(), "length" | "nodes") => { @@ -316,6 +349,21 @@ mod tests { let prefix = "MATCH p = shortestPath((a:Person)-[:R*]-(b:Person)) "; for (suffix, rich) in [ ("RETURN length(p), nodes(p)", false), + ("RETURN p IS NULL, p IS NOT NULL", false), + ( + "RETURN CASE p IS NULL WHEN true THEN -1 ELSE length(p) END", + false, + ), + ( + "RETURN CASE WHEN p IS NOT NULL THEN length(p) ELSE -1 END", + false, + ), + ( + "WITH p AS q WHERE q IS NOT NULL OR NOT q IS NULL RETURN length(q)", + false, + ), + ("RETURN p, p IS NULL", true), + ("RETURN relationships(p) IS NULL", true), ("RETURN [n IN nodes(p) | n.id]", false), ("WITH p RETURN nodes(p)", false), ( diff --git a/docs/cypher.md b/docs/cypher.md index 88abd89..2df5007 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -363,7 +363,7 @@ properties. It therefore materializes those tables, even with `use_csr=True`. Tr in batches and named ordinary paths keep a cursor per hop, but inventory memory grows with the graph. Rich path queries bypass the query-template cache: each execution binds its current parameters and creates an inventory from its own snapshot. A reader opened before `refresh()` retains the old snapshot. -Queries asking only for shortest-path length or node keys, including homogeneous node-property +Queries asking only for shortest-path length, node keys or path null checks, including homogeneous node-property comprehensions and paths carried through `WITH` aliases, retain the usual CSR execution route when the relationship joins the same node label. Cross-label shortest paths use typed identities even for length or node keys, so equal numeric keys from different labels cannot become a false zero-hop path. diff --git a/python/python/tests/test_engine_paths.py b/python/python/tests/test_engine_paths.py index 922a369..1951a92 100644 --- a/python/python/tests/test_engine_paths.py +++ b/python/python/tests/test_engine_paths.py @@ -137,3 +137,32 @@ def test_shortest_paths_shares_a_source_across_targets(use_csr): (2, 2, [5, 1, 2]), (3, 2, [5, 1, 3]), ] + + +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize(("target", "hops", "missing"), [(3, 2, False), (6, -1, True)]) +@pytest.mark.parametrize( + "projection", + [ + "RETURN CASE p IS NULL WHEN true THEN -1 ELSE length(p) END AS hops, " + "p IS NULL AS missing", + "RETURN CASE WHEN p IS NOT NULL THEN length(p) ELSE -1 END AS hops, " + "p IS NULL AS missing", + "WITH p AS q RETURN CASE q IS NULL WHEN true THEN -1 ELSE length(q) END AS hops, " + "q IS NULL AS missing", + "WITH p AS q WHERE q IS NOT NULL OR q IS NULL " + "RETURN coalesce(length(q), -1) AS hops, q IS NULL AS missing", + ], +) +def test_shortest_path_null_checks_keep_compact_plan( + use_csr, target, hops, missing, projection +): + engine = CypherEngine(CONFIG, {"Person": PEOPLE, "KNOWS": KNOWS}, use_csr=use_csr) + query = ( + f"MATCH (a:Person {{id: 5}}), (b:Person {{id: {target}}}) " + "OPTIONAL MATCH p = shortestPath((a)-[:KNOWS*1..5]-(b)) " + projection + ) + assert engine.execute(query).to_pylist() == [{"hops": hops, "missing": missing}] + plan = engine.explain(query) + assert "ShortestPathExec" in plan + assert "PathMatchExec" not in plan From 6745cbe30ff25b76e248265a99c14e0cd3f13804 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 19:16:20 -0700 Subject: [PATCH 54/66] fix(paths): borrow null-check inputs during demand analysis --- crates/lance-graph/src/path_value.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/crates/lance-graph/src/path_value.rs b/crates/lance-graph/src/path_value.rs index 91c3259..e2c677f 100644 --- a/crates/lance-graph/src/path_value.rs +++ b/crates/lance-graph/src/path_value.rs @@ -127,9 +127,10 @@ pub(crate) fn paths_array( fn mask_path_null_checks(predicate: &mut crate::ast::BooleanExpression, names: &HashSet) { use crate::ast::{BooleanExpression as B, PropertyValue, ValueExpression as V}; match predicate { - B::IsNull(value) | B::IsNotNull(value) if matches!(value, V::Variable(name) if names.contains(&name.to_lowercase())) => - { - *value = V::Literal(PropertyValue::Null); + B::IsNull(value) | B::IsNotNull(value) => { + if matches!(&*value, V::Variable(name) if names.contains(&name.to_lowercase())) { + *value = V::Literal(PropertyValue::Null); + } } B::And(left, right) | B::Or(left, right) => { mask_path_null_checks(left, names); From c89fe9e144ae8de130a6d15c7f132bd5acc0f332 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 19:45:08 -0700 Subject: [PATCH 55/66] fix(paths): bind compact null tests to the length column --- crates/lance-graph/src/logical_plan.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index a15fbce..0123ede 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -691,6 +691,7 @@ impl<'a> LogicalPlanner<'a> { match expr { ValueExpression::Variable(var) => { let property = match self.variables.get(var).map(String::as_str) { + Some("Path") if !self.rich_paths => Some("length".to_string()), Some("Path") => None, Some(label) => self .config From 1398769078ed1a08d9201f196f53ff3fb9b0bdb6 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 19:55:38 -0700 Subject: [PATCH 56/66] style(tests): wrap the compact-path null regression query --- python/python/tests/test_engine_paths.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/python/tests/test_engine_paths.py b/python/python/tests/test_engine_paths.py index 1951a92..bece713 100644 --- a/python/python/tests/test_engine_paths.py +++ b/python/python/tests/test_engine_paths.py @@ -148,8 +148,8 @@ def test_shortest_paths_shares_a_source_across_targets(use_csr): "p IS NULL AS missing", "RETURN CASE WHEN p IS NOT NULL THEN length(p) ELSE -1 END AS hops, " "p IS NULL AS missing", - "WITH p AS q RETURN CASE q IS NULL WHEN true THEN -1 ELSE length(q) END AS hops, " - "q IS NULL AS missing", + "WITH p AS q RETURN CASE q IS NULL WHEN true THEN -1 " + "ELSE length(q) END AS hops, q IS NULL AS missing", "WITH p AS q WHERE q IS NOT NULL OR q IS NULL " "RETURN coalesce(length(q), -1) AS hops, q IS NULL AS missing", ], From 6f86a1aa17083532e23f9e0ffd4490743fecda4e Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 20:42:28 -0700 Subject: [PATCH 57/66] perf(csr): inline the measured cursor construction hot path --- crates/lance-graph/src/csr_expand.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index fa48783..cd7abeb 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -230,7 +230,7 @@ impl Adjacency { self.neighbors_from(key, &mut 0).for_each(f); } - #[inline] + #[inline(always)] fn base_neighbors(&self, key: i64) -> &[u64] { self.base .dense @@ -239,7 +239,7 @@ impl Adjacency { } /// Resume at an occurrence offset, borrowing all neighbour and deletion storage. - #[inline] + #[inline(always)] fn neighbors_from<'a>( &'a self, key: i64, From 23bc699b121f3ac3e8bed4f8c2bd02474d8bc420 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 20:42:56 -0700 Subject: [PATCH 58/66] test(query): cover deadlines in correlated CALL subqueries --- python/python/tests/test_timeout.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/python/python/tests/test_timeout.py b/python/python/tests/test_timeout.py index 9717815..e931ccf 100644 --- a/python/python/tests/test_timeout.py +++ b/python/python/tests/test_timeout.py @@ -90,3 +90,14 @@ def test_expensive_query_stops_and_engine_stays_usable(engine): ) assert time.monotonic() - started < 2.0 assert engine.execute("RETURN 2 AS n").to_pylist() == [{"n": 2}] + + +def test_correlated_call_inherits_deadline_and_leaves_engine_reusable(engine): + query = ( + "UNWIND [1, 2] AS minimum CALL { WITH minimum " + "MATCH (a:Person), (b:Person), (c:Person) " + "WHERE a.id + b.id + c.id > minimum RETURN count(*) AS n } RETURN n" + ) + with pytest.raises(TimeoutError, match="Query timeout exceeded"): + engine.execute(query, timeout=0.05) + assert engine.execute("RETURN 4 AS n", timeout=10).to_pylist() == [{"n": 4}] From 7d2ef3e11b3858cc4571db315df4e848a37bdb82 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 23:10:59 -0700 Subject: [PATCH 59/66] fix(csr): integrate bounded raw work across traversal operators --- crates/lance-graph/src/csr_expand.rs | 186 +++++++++++++++++++++--- crates/lance-graph/src/csr_traversal.rs | 58 +++++++- crates/lance-graph/src/shortest_path.rs | 31 ++++ docs/csr.md | 4 +- 4 files changed, 252 insertions(+), 27 deletions(-) diff --git a/crates/lance-graph/src/csr_expand.rs b/crates/lance-graph/src/csr_expand.rs index cd7abeb..8dedce6 100644 --- a/crates/lance-graph/src/csr_expand.rs +++ b/crates/lance-graph/src/csr_expand.rs @@ -249,6 +249,7 @@ impl Adjacency { } /// Keep tombstones visible to cancellation without repeating lookups for every edge. + #[inline(always)] fn candidates_from<'a>( &'a self, key: i64, @@ -258,12 +259,14 @@ impl Adjacency { let added = self.added.get(&key).map_or(&[][..], Vec::as_slice); let removed = self.removed.get(&key).map_or(&[][..], Vec::as_slice); let mut removed = &removed[removed.partition_point(|&position| position < *offset)..]; - base.iter() - .chain(added) - .enumerate() - .skip(*offset) - .map(move |(position, &neighbor)| { - *offset = position + 1; + let added_offset = offset.saturating_sub(base.len()); + base.get(*offset..) + .unwrap_or_default() + .iter() + .chain(added.get(added_offset..).unwrap_or_default()) + .map(move |&neighbor| { + let position = *offset; + *offset += 1; if removed.first() == Some(&position) { removed = &removed[1..]; None @@ -809,7 +812,7 @@ pub struct Neighbors { } /// Occurrence offsets are sufficient to pause even a high-degree source row. -#[derive(Default)] +#[derive(Clone, Copy, Default)] pub(crate) struct NeighborPosition { forward: usize, reverse: usize, @@ -835,6 +838,7 @@ impl Neighbors { .for_each(f); } + #[inline(always)] pub(crate) fn neighbors_from<'a>( &'a self, key: i64, @@ -844,6 +848,7 @@ impl Neighbors { } /// Borrow each adjacency once while exposing tombstones and suppressed self-loops as work. + #[inline(always)] pub(crate) fn candidates_from<'a>( &'a self, key: i64, @@ -1229,17 +1234,38 @@ impl BatchExpansion { batch_size: usize, ) -> ExpandedRows { let mut chunk = ExpandedRows::default(); + loop { + if self.fill_chunk(adjacency, target, batch_size, &mut chunk) { + cooperative_yield().await; + } + if self.row == self.keys.len() || chunk.rows.len() == batch_size.max(1) { + return chunk; + } + } + } + + /// Keep coroutine state outside the scan while counting deleted occurrences as work. + fn fill_chunk( + &mut self, + adjacency: &Neighbors, + target: Option<&TargetLookup>, + batch_size: usize, + chunk: &mut ExpandedRows, + ) -> bool { + let mut position = self.position; + let mut row = self.row; let mut work = 0_usize; - while self.row < self.keys.len() { - let i = self.row; + 'rows: while row < self.keys.len() { + let i = row; + let previous_work = work; if self.keys.is_valid(i) && self.bound.as_ref().is_none_or(|b| b.is_valid(i)) { let key = self.keys.value(i); let required = self.bound.as_ref().map(|b| b.value(i)); - for candidate in adjacency.candidates_from(key, &mut self.position) { + for candidate in adjacency + .candidates_from(key, &mut position) + .take(1024 - work) + { work += 1; - if work.is_multiple_of(1024) { - cooperative_yield().await; - } let Some(neighbor) = candidate else { continue; }; @@ -1248,24 +1274,34 @@ impl BatchExpansion { continue; } if let Some(target) = target { - // No target row: the inner join would drop this edge too. - let Some(position) = target.rows.position(neighbor) else { + let Some(target_row) = target.rows.position(neighbor) else { continue; }; - chunk.target_positions.push(position); + chunk.target_positions.push(target_row); } chunk.rows.push(i as u32); chunk.key_values.push(key); chunk.neighbor_values.push(neighbor); if chunk.rows.len() == batch_size.max(1) { - return chunk; + break 'rows; } } + if work == 1024 { + break; + } + } + if work == previous_work { + work += 1; + } + row += 1; + position = NeighborPosition::default(); + if work == 1024 { + break; } - self.row += 1; - self.position = NeighborPosition::default(); } - chunk + self.row = row; + self.position = position; + work == 1024 } async fn next_path_chunk( @@ -1275,7 +1311,12 @@ impl BatchExpansion { batch_size: usize, ) -> ExpandedRows { let mut chunk = ExpandedRows::default(); + let mut source_work = 0_usize; while self.row < self.keys.len() { + source_work += 1; + if source_work.is_multiple_of(1024) { + cooperative_yield().await; + } let i = self.row; if self.keys.is_valid(i) && self.bound.as_ref().is_none_or(|b| b.is_valid(i)) { let key = self.keys.value(i); @@ -1937,6 +1978,33 @@ mod tests { assert!(Arc::ptr_eq(&base.base, &later.base)); } + #[tokio::test] + async fn resumed_neighbors_cross_the_base_and_delta_boundary() { + let base = Adjacency::build( + edge_stream(vec![Some(1), Some(1)], vec![Some(10), Some(11)]), + "src", + "dst", + ) + .await + .unwrap(); + let next = base.with_changes(&[(1, 12), (1, 12), (1, 13)], &[(1, 11), (1, 12)]); + for (mut offset, expected) in [ + (0, vec![10, 12, 13]), + (1, vec![12, 13]), + (2, vec![12, 13]), + (3, vec![12, 13]), + (4, vec![13]), + (5, vec![]), + (6, vec![]), + ] { + assert_eq!( + next.neighbors_from(1, &mut offset).collect::>(), + expected + ); + } + assert_eq!(neighbors(&base, 1), vec![10, 11]); + } + #[tokio::test] async fn debug_output_does_not_grow_with_edges() { let n = 10_000; @@ -2045,6 +2113,84 @@ mod tests { assert!(expansion.next_chunk(&adj, None, 1).await.rows.is_empty()); } + #[tokio::test] + async fn cancellation_without_edges_in_variable_paths_yields_between_source_rows() { + let config = crate::config::GraphConfig::builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + .unwrap(); + let nodes = + RecordBatch::try_from_iter([("id", Arc::new(Int64Array::from(vec![1])) as ArrayRef)]) + .unwrap(); + let mut catalog = AdjacencyCatalog::default(); + catalog.insert_nodes( + "Person", + Arc::new(NodeRows::from_batch(nodes, "id").unwrap().unwrap()), + ); + catalog.insert( + "KNOWS", + true, + Arc::new( + Adjacency::build(edge_stream(vec![], vec![]), "src", "dst") + .await + .unwrap(), + ), + ); + let walk = crate::csr_traversal::Traversal::new( + &config, + &catalog, + &["KNOWS".into()], + &crate::ast::RelationshipDirection::Outgoing, + "Person", + "Person", + 1, + 1, + ) + .unwrap(); + let input = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from(vec![1; 100_000])) as ArrayRef, + )]) + .unwrap(); + let mut expansion = BatchExpansion::try_new(input, 0, None).unwrap(); + let mut next = Box::pin(expansion.next_path_chunk(&walk, None, 1)); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert!(expansion.row > 0 && expansion.row <= 1024); + assert!(expansion + .next_path_chunk(&walk, None, 1) + .await + .rows + .is_empty()); + } + + #[tokio::test] + async fn cancellation_without_edges_yields_while_skipping_source_rows() { + let adjacency = Neighbors::directed(Arc::new( + Adjacency::build(edge_stream(vec![Some(1)], vec![Some(2)]), "src", "dst") + .await + .unwrap(), + )); + for missing in [Some(99), None] { + let input = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from(vec![missing; 100_000])) as ArrayRef, + )]) + .unwrap(); + let mut expansion = BatchExpansion::try_new(input, 0, None).unwrap(); + let mut next = Box::pin(expansion.next_chunk(&adjacency, None, 1)); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert!(expansion.row > 0 && expansion.row <= 1024); + assert!(expansion + .next_chunk(&adjacency, None, 1) + .await + .rows + .is_empty()); + } + } + #[tokio::test] async fn expansion_yields_before_scanning_all_rejected_neighbors() { let adj = Neighbors::directed(Arc::new( diff --git a/crates/lance-graph/src/csr_traversal.rs b/crates/lance-graph/src/csr_traversal.rs index a19eb36..cdf11cb 100644 --- a/crates/lance-graph/src/csr_traversal.rs +++ b/crates/lance-graph/src/csr_traversal.rs @@ -161,14 +161,18 @@ impl TraversalCursor { let next = (step.from == frame.label) .then(|| { step.adjacency - .neighbors_from(frame.key, &mut frame.position) + .candidates_from(frame.key, &mut frame.position) .next() }) .flatten(); - let Some(key) = next.map(|n| n as i64) else { - frame.step += 1; - frame.position = NeighborPosition::default(); - continue; + let key = match next { + Some(Some(key)) => key as i64, + Some(None) => continue, + None => { + frame.step += 1; + frame.position = NeighborPosition::default(); + continue; + } }; if graph.nodes[step.to].position(key).is_none() { continue; @@ -271,6 +275,50 @@ mod tests { assert!(cursor.stack.is_empty()); } + #[tokio::test] + async fn cancellation_visits_raw_tombstones_before_returning_the_live_path() { + let (config, mut catalog) = graph().await; + let schema = Arc::new(Schema::new(vec![ + Field::new("src", DataType::Int64, false), + Field::new("dst", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1; 5000])), + Arc::new(Int64Array::from(vec![1; 5000])), + ], + ) + .unwrap(); + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(vec![Ok(batch)]), + )); + let base = Adjacency::build(stream, "src", "dst").await.unwrap(); + catalog.insert( + "hasType", + true, + Arc::new(base.with_changes(&[], &vec![(1, 1); 4999])), + ); + let walk = Traversal::new( + &config, + &catalog, + &["hasType".into()], + &RelationshipDirection::Outgoing, + "Tag", + "Class", + 1, + 1, + ) + .unwrap(); + let mut cursor = TraversalCursor::new(&walk, 1); + let mut next = Box::pin(cursor.next(&walk)); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert_eq!(cursor.next(&walk).await, Some(1)); + assert_eq!(cursor.next(&walk).await, None); + } + #[tokio::test] async fn zero_hop_is_identity_and_depth_bound_survives_cycles() { let (config, catalog) = graph().await; diff --git a/crates/lance-graph/src/shortest_path.rs b/crates/lance-graph/src/shortest_path.rs index a5d5397..f730638 100644 --- a/crates/lance-graph/src/shortest_path.rs +++ b/crates/lance-graph/src/shortest_path.rs @@ -422,6 +422,7 @@ async fn expand_level( let mut next = Vec::new(); let mut work = 0_usize; for &node in frontier { + let previous_work = work; let mut position = crate::csr_expand::NeighborPosition::default(); for candidate in adjacency.candidates_from(node, &mut position) { work += 1; @@ -442,6 +443,12 @@ async fn expand_level( } } } + if work == previous_work { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } + } } next } @@ -699,6 +706,30 @@ mod tests { assert!(!depth.contains_key(&4)); } + #[tokio::test] + async fn cancellation_without_edges_yields_while_scanning_frontier() { + let (out, _) = adjacencies(&[]).await; + let adjacency = Neighbors::directed(out); + let frontier = vec![99; 100_000]; + let mut depth = FastMap::default(); + let mut parents = FastMap::default(); + let mut next = Box::pin(expand_level( + &adjacency, + &frontier, + 1, + &mut depth, + &mut parents, + )); + assert!(futures::poll!(next.as_mut()).is_pending()); + drop(next); + assert!(depth.is_empty() && parents.is_empty()); + assert!( + expand_level(&adjacency, &frontier, 1, &mut depth, &mut parents) + .await + .is_empty() + ); + } + #[tokio::test] async fn bidirectional_keeps_every_meeting_node_and_parallel_edge() { // A wide first layer makes the target's side expand. Two meeting nodes, diff --git a/docs/csr.md b/docs/csr.md index 4d01a3d..4bcfb21 100644 --- a/docs/csr.md +++ b/docs/csr.md @@ -95,8 +95,8 @@ Node index output chunk of neighbour positions before fetching properties. It yields that chunk before visiting the remaining neighbours, even for a single high-degree node. Deleted edge positions are read from the shared snapshot rather than copied for each query; refresh resolves those positions by scanning the - affected sources' neighbour lists. Cancellation checkpoints count deleted edges and suppressed - reverse self-loops as work, using borrowed adjacency slices while a chunk is being collected. + affected sources' neighbour lists. Cancellation checkpoints count source rows and frontiers with no surviving edges, + deleted edges and suppressed reverse self-loops as work, using borrowed adjacency slices while a chunk is being collected. The chunk row limit comes from DataFusion's batch size; wide properties and multiple active partitions still increase working memory. This does not bound the resident graph, the input data, or buffering in joins, sorts and aggregations elsewhere in the plan. From bd11b767fe7f7964817eebb0d863261154f6e113 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 26 Oct 2026 10:19:36 -0700 Subject: [PATCH 60/66] fix(cypher): keep string-number comparisons coercible The expression type check rejected comparisons such as p.age = '30', which DataFusion has always run by casting the string side. Let DataFusion's coercion decide scalar compatibility again, so only genuinely incompatible operands (string arithmetic, booleans against numbers, entity arithmetic) fail before planning. Update the tests that asserted the old rejection and document the behaviour. --- crates/lance-csr/src/query.rs | 2 +- crates/lance-csr/src/semantic/types.rs | 35 ++++++++++-------------- docs/cypher.md | 8 ++++-- python/python/tests/test_engine_types.py | 32 ++++++++++++---------- python/python/tests/test_plan_cache.py | 2 +- 5 files changed, 39 insertions(+), 40 deletions(-) diff --git a/crates/lance-csr/src/query.rs b/crates/lance-csr/src/query.rs index 5e9ca14..676b661 100644 --- a/crates/lance-csr/src/query.rs +++ b/crates/lance-csr/src/query.rs @@ -1862,7 +1862,7 @@ mod tests { assert_eq!(query.execute(data(2), None).await.unwrap().num_rows(), 0); assert!(query .clone() - .with_parameter("value", "wrong") + .with_parameter("value", true) .execute(data(1), None) .await .unwrap_err() diff --git a/crates/lance-csr/src/semantic/types.rs b/crates/lance-csr/src/semantic/types.rs index 910bfb5..5718e4d 100644 --- a/crates/lance-csr/src/semantic/types.rs +++ b/crates/lance-csr/src/semantic/types.rs @@ -227,22 +227,11 @@ impl SemanticAnalyzer { | Operator::GtEq ); let compatible = match (&left_type, &right_type) { - (Scalar(left), Scalar(right)) => { - let string_numeric = (left.is_numeric() - && matches!( - right, - DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View - )) - || (right.is_numeric() - && matches!( - left, - DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View - )); - !string_numeric - && BinaryTypeCoercer::new(left, &operator, right) - .get_result_type() - .is_ok() - } + // DataFusion decides scalar compatibility; it casts the string side of a + // string-number comparison, as queries did before this check existed. + (Scalar(left), Scalar(right)) => BinaryTypeCoercer::new(left, &operator, right) + .get_result_type() + .is_ok(), (Entity(left, _), Entity(right, _)) => comparison && left == right, // Collection comparison rules remain with the planner; arithmetic on entities does not. (Entity(..), _) | (_, Entity(..)) => false, @@ -367,7 +356,6 @@ mod tests { #[test] fn rejects_known_operand_types_in_their_scope() { for (query, expression, types) in [ - ("MATCH (p:Person) WHERE p.name = 1 RETURN p.id", "p.name = 1", ["Utf8", "Int64"]), ("MATCH (p:Person) RETURN p + 1", "p + 1", ["Node", "Int64"]), ("MATCH (p:Person) WITH p.name AS value RETURN value * 2", "value * 2", ["Utf8", "Int64"]), ("MATCH (p:Person) WITH collect(p.name) AS names UNWIND names AS name RETURN name / 2", "name / 2", ["Utf8", "Int64"]), @@ -383,16 +371,23 @@ mod tests { fn scalar_return_types_follow_with_and_unwind_scopes() { for query in [ "RETURN left('12', 1) + 1", - "RETURN abs(1) = '1'", - "WITH sqrt(4) AS value RETURN value = '2'", "WITH split('1,2', ',') AS values UNWIND values AS value RETURN value + 1", "RETURN last(reverse(['1'])) + 1", "WITH tail(['1','2']) AS values UNWIND values AS value RETURN value + 1", + ] { + let errors = analyze(query).errors.join("\n"); + assert!(errors.contains("Type mismatch"), "{query}: {errors}"); + } + // String-number comparisons are left to DataFusion, which casts the string side. + for query in [ + "RETURN abs(1) = '1'", + "WITH sqrt(4) AS value RETURN value = '2'", "UNWIND keys({one: 1}) AS key RETURN key = 1", "MATCH (p:Person) UNWIND labels(p) AS label RETURN label = 1", + "MATCH (p:Person) WHERE p.name = 1 RETURN p.id", ] { let errors = analyze(query).errors.join("\n"); - assert!(errors.contains("Type mismatch"), "{query}: {errors}"); + assert!(!errors.contains("Type mismatch"), "{query}: {errors}"); } } diff --git a/docs/cypher.md b/docs/cypher.md index 2df5007..e18c6b6 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -538,6 +538,8 @@ aggregation and `OPTIONAL MATCH`. | `RuntimeError` | Execution fails, for example when a Lance dataset can't be read. | Known comparison and arithmetic types are checked against the current dataset schemas before planning. -For example, comparing a string property with a number reports the expression and both types in a -`ValueError`. Mixed numeric types and supported date/timestamp coercions remain valid. Types that cannot -be determined in the current scope are left to the planner, so these checks are not a complete type system. +For example, adding a number to a string property reports the expression and both types in a +`ValueError`. Mixed numeric types and supported date/timestamp coercions remain valid. Comparing a string +with a number is allowed: DataFusion casts the string side, so `p.age = '40'` matches an age of 40, and a +string that isn't a number raises a `RuntimeError` when the query runs. Types that cannot be determined in +the current scope are left to the planner, so these checks are not a complete type system. diff --git a/python/python/tests/test_engine_types.py b/python/python/tests/test_engine_types.py index bcd1e22..f2efc5e 100644 --- a/python/python/tests/test_engine_types.py +++ b/python/python/tests/test_engine_types.py @@ -38,11 +38,6 @@ def engine(request): @pytest.mark.parametrize( ("query", "expression", "types"), [ - ( - "MATCH (p:Person) WHERE p.name = 1 RETURN p.id", - "p.name = 1", - ("Utf8", "Int64"), - ), ("MATCH (p:Person) RETURN p + 1", "p + 1", ("Node", "Int64")), ( "MATCH (p:Person) WITH p.name AS value RETURN value * 2", @@ -60,11 +55,6 @@ def engine(request): "p + 1", ("Utf8", "Int64"), ), - ( - "MATCH (p:Person)-[r:KNOWS]->(q:Person) WHERE r.label = 1 RETURN q.id", - "r.label = 1", - ("Utf8", "Int64"), - ), ], ) def test_type_errors_name_expression_and_operands(engine, query, expression, types): @@ -101,20 +91,32 @@ def test_comprehension_shadow_does_not_change_outer_type(engine): "query", [ "RETURN left('12', 1) + 1", - "RETURN abs(1) = '1'", - "WITH sqrt(4) AS value RETURN value = '2'", "WITH split('1,2', ',') AS values UNWIND values AS value RETURN value + 1", "RETURN last(reverse(['1'])) + 1", "WITH tail(['1', '2']) AS values UNWIND values AS value RETURN value + 1", - "UNWIND keys({one: 1}) AS key RETURN key = 1", - "MATCH (p:Person) UNWIND labels(p) AS label RETURN label = 1", ], ) -def test_scalar_return_types_reject_string_numeric_coercion(engine, query): +def test_scalar_return_types_reject_string_numeric_arithmetic(engine, query): with pytest.raises(ValueError, match="Type mismatch"): engine.execute(query) +@pytest.mark.parametrize( + ("query", "expected"), + [ + ("RETURN abs(1) = '1' AS same", [{"same": True}]), + ("WITH sqrt(4) AS value RETURN value = '2' AS same", [{"same": True}]), + ( + "MATCH (p:Person) WHERE p.age = '40' RETURN p.name AS name", + [{"name": "Bob"}], + ), + ], +) +def test_string_numeric_comparisons_coerce_as_before(engine, query, expected): + # DataFusion casts the string side of a string-number comparison. + assert engine.execute(query).to_pylist() == expected + + @pytest.mark.parametrize( ("query", "expected"), [ diff --git a/python/python/tests/test_plan_cache.py b/python/python/tests/test_plan_cache.py index 30aa404..700dc25 100644 --- a/python/python/tests/test_plan_cache.py +++ b/python/python/tests/test_plan_cache.py @@ -80,7 +80,7 @@ def test_cached_query_still_rejects_missing_or_wrong_types(engine): with pytest.raises(ValueError, match="[Pp]arameter"): engine.execute(query) with pytest.raises(ValueError, match="Type mismatch"): - engine.execute(query, {"id": "wrong"}) + engine.execute(query, {"id": True}) @pytest.mark.parametrize( From dd1909788963b83025b044bba5eca90e0720e26d Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 26 Oct 2026 10:21:23 -0700 Subject: [PATCH 61/66] fix(csr): give unlabeled variable-length targets the source label Variable-length expansion started reading the target label from the logical plan, which defaults an unlabeled node to "Node". With relationships that record no endpoint labels (the deprecated builder), MATCH (a:Person)-[:KNOWS*1..2]->(b) then failed with "No mapping found for node label: Node". Fall back to the source variable's label in that case, as the planner did before, in both the analysis and label collection. --- .../src/datafusion_planner/analysis.rs | 22 ++++++++++++--- .../src/datafusion_planner/builder/helpers.rs | 7 ++++- python/python/tests/test_engine_paths.py | 27 ++++++++++++++++++- 3 files changed, 51 insertions(+), 5 deletions(-) diff --git a/crates/lance-csr/src/datafusion_planner/analysis.rs b/crates/lance-csr/src/datafusion_planner/analysis.rs index 477231b..76ab3fe 100644 --- a/crates/lance-csr/src/datafusion_planner/analysis.rs +++ b/crates/lance-csr/src/datafusion_planner/analysis.rs @@ -77,6 +77,20 @@ impl<'a> PlanningContext<'a> { } /// Analyze the logical plan to extract metadata +/// The label a variable-length target resolves to: its own label, or the source's +/// when the pattern left it unlabeled (the planner's `"Node"` default), so +/// `(a:Person)-[:KNOWS*1..2]->(b)` makes `b` a Person, as single-label graphs expect. +// shortcut: an explicit `(b:Node)` is indistinguishable from the default; make the default an Option if a graph ever needs a real `Node` label on a variable-length target. +pub(crate) fn variable_length_target_label( + target_label: &str, + source_label: Option<&String>, +) -> String { + match source_label { + Some(source) if target_label == "Node" => source.clone(), + _ => target_label.to_string(), + } +} + pub fn analyze(logical_plan: &LogicalOperator) -> Result { let mut analysis = QueryAnalysis::default(); let mut rel_counter: HashMap = HashMap::new(); @@ -159,9 +173,11 @@ fn analyze_operator( // Recursively analyze input first analyze_operator(input, analysis, rel_counter)?; - analysis - .var_to_label - .insert(target_variable.clone(), target_label.clone()); + let label = variable_length_target_label( + target_label, + analysis.var_to_label.get(source_variable), + ); + analysis.var_to_label.insert(target_variable.clone(), label); // For variable-length paths, register multiple instances (one per hop) // We need to register instances for all possible hop counts diff --git a/crates/lance-csr/src/datafusion_planner/builder/helpers.rs b/crates/lance-csr/src/datafusion_planner/builder/helpers.rs index d2d3e38..f2e5721 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/helpers.rs @@ -153,12 +153,17 @@ impl DataFusionPlanner { } LogicalOperator::VariableLengthExpand { input, + source_variable, target_variable, target_label, .. } => { Self::collect_labels(input, labels); - labels.insert(target_variable.clone(), target_label.clone()); + let label = crate::datafusion_planner::analysis::variable_length_target_label( + target_label, + labels.get(source_variable), + ); + labels.insert(target_variable.clone(), label); } LogicalOperator::Filter { input, .. } | LogicalOperator::Project { input, .. } diff --git a/python/python/tests/test_engine_paths.py b/python/python/tests/test_engine_paths.py index bece713..b548f91 100644 --- a/python/python/tests/test_engine_paths.py +++ b/python/python/tests/test_engine_paths.py @@ -7,9 +7,11 @@ (twice), 1->3, 2->3, 3->3 and 5->1; 4->99 points at a missing person. """ +import warnings + import pyarrow as pa import pytest -from lance_csr import CypherEngine +from lance_csr import CypherEngine, GraphConfig from test_engine_csr import CONFIG, KNOWS, PEOPLE CASES = [ @@ -166,3 +168,26 @@ def test_shortest_path_null_checks_keep_compact_plan( plan = engine.explain(query) assert "ShortestPathExec" in plan assert "PathMatchExec" not in plan + + +@pytest.mark.parametrize("use_csr", [False, True]) +def test_variable_length_target_inherits_source_label_without_endpoints(use_csr): + # The deprecated builder records no endpoint labels; an unlabeled + # variable-length target then takes the source's label, as before. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + config = ( + GraphConfig.builder() + .with_node_label("Person", "id") + .with_relationship("KNOWS", "src", "dst") + .build() + ) + data = { + "Person": pa.table({"id": [1, 2, 3, 4], "name": ["A", "B", "C", "D"]}), + "KNOWS": pa.table({"src": [1, 2, 2], "dst": [2, 3, 4]}), + } + rows = CypherEngine(config, data, use_csr).execute( + "MATCH (a:Person)-[:KNOWS*1..2]->(b) WHERE a.id = 1 " + "RETURN b.name AS name ORDER BY name" + ) + assert rows.to_pylist() == [{"name": "B"}, {"name": "C"}, {"name": "D"}] From 3c1809e2360a16f80184eacb8c56e6f362ea2ecc Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 26 Oct 2026 10:23:58 -0700 Subject: [PATCH 62/66] fix(cypher): export typed nulls as NULL outside the Default dialect SQL export rejected every typed null outside the Default dialect, so ordinary queries such as coalesce(p.name, null) or CASE ... ELSE null stopped exporting for Spark, PostgreSQL, MySQL and SQLite. Those dialects have no arrow_cast; emit a bare NULL there, as exports did before, and keep the arrow_cast wrapper for the Default dialect. --- crates/lance-csr-python/src/graph.rs | 3 +-- crates/lance-csr/src/query.rs | 20 ++++++++++---------- docs/api-reference.md | 2 +- docs/sql.md | 4 ++-- python/python/tests/test_to_sql.py | 20 ++++++++++++-------- 5 files changed, 26 insertions(+), 23 deletions(-) diff --git a/crates/lance-csr-python/src/graph.rs b/crates/lance-csr-python/src/graph.rs index 7a3a55f..9bb3d44 100644 --- a/crates/lance-csr-python/src/graph.rs +++ b/crates/lance-csr-python/src/graph.rs @@ -604,8 +604,7 @@ impl CypherQuery { /// Raises /// ------ /// NotImplementedError - /// If a runtime-only Cypher function has no faithful SQL translation, or a - /// typed-null literal is exported outside the Default dialect + /// If a runtime-only Cypher function has no faithful SQL translation /// RuntimeError /// If SQL generation fails #[pyo3(signature = (datasets, dialect=None))] diff --git a/crates/lance-csr/src/query.rs b/crates/lance-csr/src/query.rs index 676b661..37481d8 100644 --- a/crates/lance-csr/src/query.rs +++ b/crates/lance-csr/src/query.rs @@ -482,8 +482,7 @@ impl CypherQuery { /// /// # Errors /// Returns `GraphError::UnsupportedFeature` for runtime-only Cypher functions - /// without a faithful SQL translation or typed nulls outside the Default dialect. - /// Constant expressions may be folded first. + /// without a faithful SQL translation. Constant expressions may be folded first. pub async fn to_sql( &self, datasets: HashMap, @@ -542,14 +541,15 @@ impl CypherQuery { } if let Expr::Literal(value, _) = &expr { let data_type = value.data_type(); - if value.is_null() && !data_type.is_null() { - let type_name = data_type.to_string(); - if dialect != SqlDialect::Default - || !type_name.parse::().is_ok_and(|parsed| parsed == data_type) - { - return not_impl_err!("SQL export of typed null {data_type} requires a faithful Default-dialect Arrow cast"); - } - // The unparser emits bare NULL for typed nulls, losing their Arrow type. + let type_name = data_type.to_string(); + // The unparser emits bare NULL for typed nulls, losing their Arrow type. + // Only the Default dialect has arrow_cast; elsewhere bare NULL is the + // closest SQL and yields the same null value, as exports always did. + if value.is_null() + && !data_type.is_null() + && dialect == SqlDialect::Default + && type_name.parse::().is_ok_and(|parsed| parsed == data_type) + { return Ok(Transformed::yes(datafusion::functions::core::arrow_cast().call(vec![ lit(ScalarValue::Null), lit(type_name), ]))); diff --git a/docs/api-reference.md b/docs/api-reference.md index 278fca6..400d26e 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -219,7 +219,7 @@ A parsed Cypher query. Raises `ValueError` if `query_text` doesn't parse. Builde `#!python to_sql(datasets: dict, dialect: SqlDialect | None = None) -> str` : Exports supported expressions as SQL joins. Raises `NotImplementedError` for runtime-only Cypher - functions without a faithful SQL translation, and for typed-null literals outside the Default dialect. See [From Cypher to SQL](sql.md#from-cypher-to-sql). + functions without a faithful SQL translation. See [From Cypher to SQL](sql.md#from-cypher-to-sql). `#!python execute_with_vector_rerank(datasets: dict, vector_search: VectorSearch) -> pyarrow.Table` : Runs the query, then reranks by vector similarity. diff --git a/docs/sql.md b/docs/sql.md index 85a5b1c..fd5bf8d 100644 --- a/docs/sql.md +++ b/docs/sql.md @@ -90,8 +90,8 @@ Integer `abs` also normalizes its result to Int64 in Cypher, while SQL preserves Use `CypherQuery.execute` or `CypherEngine.execute` for these expressions. Constant-folded typed nulls use DataFusion `arrow_cast` in the Default dialect so their Arrow -result types are preserved. Other dialects raise `NotImplementedError` for typed-null literals; -a bare Cypher `null` remains an untyped SQL `NULL`. +result types are preserved. Other dialects have no `arrow_cast`, so typed nulls export as a bare SQL +`NULL`: the value is the same, but its column type is left to the target database. List `size` and inclusive `range` retain SQL export support through native expressions, including null arguments, integer casts and the optional range step. Constant expressions that the optimizer diff --git a/python/python/tests/test_to_sql.py b/python/python/tests/test_to_sql.py index 78b9614..f0ac77e 100644 --- a/python/python/tests/test_to_sql.py +++ b/python/python/tests/test_to_sql.py @@ -100,15 +100,19 @@ def test_sql_export_rejects_runtime_only_cypher_functions( "dialect", [SqlDialect.Spark, SqlDialect.PostgreSql, SqlDialect.MySql, SqlDialect.Sqlite], ) -def test_sql_export_typed_null_requires_default_dialect(scalar_export_env, dialect): +def test_sql_export_typed_null_is_bare_null_outside_default(scalar_export_env, dialect): + # Only the Default dialect has arrow_cast; elsewhere typed nulls export as NULL, + # so ordinary queries such as coalesce(..., null) keep exporting. config, datasets = scalar_export_env - query = CypherQuery("MATCH (p:Person) RETURN round(null) AS value").with_config( - config - ) - with pytest.raises(NotImplementedError, match="SQL export of typed null Float64"): - query.to_sql(datasets, dialect) - bare_null = CypherQuery("MATCH (p:Person) RETURN null AS value").with_config(config) - assert "NULL" in bare_null.to_sql(datasets, dialect) + for text in [ + "MATCH (p:Person) RETURN round(null) AS value", + "MATCH (p:Person) RETURN coalesce(p.text, null) AS value", + "MATCH (p:Person) RETURN CASE WHEN p.id > 1 THEN p.text ELSE null END AS value", + "MATCH (p:Person) RETURN null AS value", + ]: + sql = CypherQuery(text).with_config(config).to_sql(datasets, dialect) + assert "NULL" in sql.upper(), sql + assert "arrow_cast" not in sql, sql def test_sql_export_checks_filter_functions(scalar_export_env): From ff2123fab3d438116d1f83ac6ddf6e110c843564 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 26 Oct 2026 10:26:30 -0700 Subject: [PATCH 63/66] fix(cypher): reject queries that mix UNION and UNION ALL openCypher allows one union mode per query; the planner folded a mix left to right instead. Report it during semantic analysis, update the tests that expected the folded result and document the rule. --- crates/lance-csr/src/query.rs | 5 +---- crates/lance-csr/src/semantic.rs | 11 +++++++++++ docs/cypher.md | 4 ++-- python/python/tests/test_union.py | 4 ++-- 4 files changed, 16 insertions(+), 8 deletions(-) diff --git a/crates/lance-csr/src/query.rs b/crates/lance-csr/src/query.rs index 37481d8..524dde1 100644 --- a/crates/lance-csr/src/query.rs +++ b/crates/lance-csr/src/query.rs @@ -1601,10 +1601,6 @@ mod tests { "MATCH (p:Person) RETURN p.id AS n UNION ALL RETURN 2 AS n", 4, ), - ( - "RETURN 1 AS n UNION ALL RETURN 1 AS n UNION RETURN 2 AS n", - 2, - ), ("RETURN NULL AS n UNION RETURN NULL AS n", 1), ] { let result = CypherQuery::new(query) @@ -1618,6 +1614,7 @@ mod tests { for query in [ "RETURN 1 AS x UNION RETURN 2 AS y", "RETURN 1 AS x UNION RETURN 'a' AS x", + "RETURN 1 AS n UNION ALL RETURN 1 AS n UNION RETURN 2 AS n", ] { assert!(CypherQuery::new(query) .unwrap() diff --git a/crates/lance-csr/src/semantic.rs b/crates/lance-csr/src/semantic.rs index 223cd5f..f9909a0 100644 --- a/crates/lance-csr/src/semantic.rs +++ b/crates/lance-csr/src/semantic.rs @@ -241,6 +241,17 @@ impl SemanticAnalyzer { // Phase 8: Type checking self.validate_types(&mut errors); + // openCypher allows UNION or UNION ALL throughout a query, not a mix of both. + if let Some(first) = analyzed_query.union_branches.first() { + if analyzed_query + .union_branches + .iter() + .any(|branch| branch.all != first.all) + { + errors.push("Invalid combination of UNION and UNION ALL".into()); + } + } + for branch in &mut analyzed_query.union_branches { let mut analyzer = SemanticAnalyzer::new(self.config.clone()); analyzer.catalog = self.catalog.clone(); diff --git a/docs/cypher.md b/docs/cypher.md index e18c6b6..7d9d613 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -227,8 +227,8 @@ vectors; conversion to float32 happens inside the vector operation. A configured for standalone `RETURN` queries. `UNION` removes duplicate rows; `UNION ALL` preserves them. Branches have independent variable scopes and -keep their own `ORDER BY`, `SKIP` and `LIMIT`. Combined output has no guaranteed order. Mixed UNION modes -are evaluated from left to right. Numeric columns can promote to a common numeric type; incompatible +keep their own `ORDER BY`, `SKIP` and `LIMIT`. Combined output has no guaranteed order. As in openCypher, +one query uses either `UNION` or `UNION ALL` throughout; mixing them raises `ValueError`. Numeric columns can promote to a common numeric type; incompatible column counts, names or types raise `ValueError`. This follows the existing homogeneous Arrow value model. ```python diff --git a/python/python/tests/test_union.py b/python/python/tests/test_union.py index acdbce9..ae8a700 100644 --- a/python/python/tests/test_union.py +++ b/python/python/tests/test_union.py @@ -23,8 +23,6 @@ def engine(request): [ ("RETURN 1 AS n UNION RETURN 1 AS n", [1]), ("RETURN 1 AS n UNION ALL RETURN 1 AS n", [1, 1]), - ("RETURN 1 AS n UNION ALL RETURN 1 AS n UNION RETURN 2 AS n", [1, 2]), - ("RETURN 1 AS n UNION RETURN 1 AS n UNION ALL RETURN 1 AS n", [1, 1]), ( "MATCH (p:Person) RETURN p.id AS n UNION MATCH (p:City) RETURN p.id AS n", [1, 2, 3], @@ -76,6 +74,8 @@ def test_numeric_promotion_and_parameters(engine): "RETURN 1 AS x UNION RETURN 'a' AS x", "WITH 1 AS x RETURN x UNION RETURN x", "RETURN 1 AS x UNION ALL", + "RETURN 1 AS n UNION ALL RETURN 1 AS n UNION RETURN 2 AS n", + "RETURN 1 AS n UNION RETURN 1 AS n UNION ALL RETURN 1 AS n", ], ) def test_invalid_union(engine, query): From e895aaa860401c82743043536bd4408f4c9817ed Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 26 Oct 2026 10:28:07 -0700 Subject: [PATCH 64/66] fix(paths): keep start-node filters above the path operator DataFusion copies a predicate pushed past an extension node into every input. PathMatch's extra inputs (node and edge inventories) lack the start node's columns, so MATCH p = (a)-[...]->(b) WHERE a.id = 1 failed to plan with "No field named a__id". Only allow pushing predicates on columns every input carries. --- crates/lance-csr/src/path_match.rs | 14 +++++++-- python/python/tests/test_path_values.py | 42 +++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/crates/lance-csr/src/path_match.rs b/crates/lance-csr/src/path_match.rs index 7d99585..07895b1 100644 --- a/crates/lance-csr/src/path_match.rs +++ b/crates/lance-csr/src/path_match.rs @@ -182,11 +182,21 @@ impl UserDefinedLogicalNodeCore for PathMatchNode { vec![] } fn prevent_predicate_push_down_columns(&self) -> HashSet { + // DataFusion copies a pushed predicate into every input, so only columns that + // every input carries may be pushed; others (start-node columns are only in the + // first input) stay above this node. + // shortcut: a WHERE on start-node columns is not pushed into the first input's scan; push it there directly if path queries with selective start filters get slow. self.schema .fields() .iter() - .skip(self.inputs[0].schema().fields().len()) - .map(|f| f.name().clone()) + .map(|f| f.name()) + .filter(|name| { + !self + .inputs + .iter() + .all(|input| input.schema().has_column_with_unqualified_name(name)) + }) + .cloned() .collect() } fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { diff --git a/python/python/tests/test_path_values.py b/python/python/tests/test_path_values.py index 512e46f..2d5c329 100644 --- a/python/python/tests/test_path_values.py +++ b/python/python/tests/test_path_values.py @@ -605,3 +605,45 @@ def test_cross_label_shortest_inventory_isolates_refresh(cross_label_graph, tmp_ assert engine.execute(query, {"start": start}).to_pylist() == ( [{"ids": [1, 2]}] if start == 1 else [] ) + + +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize( + ("query", "expected"), + [ + ( + "MATCH p = (a:Person)-[:KNOWS]->(b:Person) WHERE a.id = 1 " + "RETURN b.name AS name, size(nodes(p)) AS n", + [{"name": "B", "n": 2}], + ), + ( + "MATCH p = (a:Person)-[:KNOWS]-(b:Person) WHERE a.id = 2 " + "RETURN b.name AS name, size(nodes(p)) AS n ORDER BY name", + [{"name": "A", "n": 2}, {"name": "C", "n": 2}], + ), + ( + "MATCH p = (a:Person)-[:KNOWS*1..2]->(b:Person) WHERE a.name = 'A' " + "RETURN b.name AS name, size(nodes(p)) AS n ORDER BY name", + [{"name": "B", "n": 2}, {"name": "C", "n": 3}], + ), + ( + "MATCH p = (a:Person)-[:KNOWS]->(b:Person) WHERE a.id = 1 AND b.name = 'B' " + "RETURN b.name AS name, size(nodes(p)) AS n", + [{"name": "B", "n": 2}], + ), + ], +) +def test_path_where_on_start_node(use_csr, query, expected): + # The path operator has several inputs; a start-node filter must not be pushed + # into the inventory inputs, which lack the start node's columns. + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + data = { + "Person": pa.table({"id": [1, 2, 3], "name": ["A", "B", "C"]}), + "KNOWS": pa.table({"src": [1, 2], "dst": [2, 3]}), + } + assert CypherEngine(config, data, use_csr).execute(query).to_pylist() == expected From 4d0ed860242ea8f89bd3ffbaea61019fa151117a Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 26 Oct 2026 10:28:50 -0700 Subject: [PATCH 65/66] docs: document the target_partitions range and Rust constructor target_partitions is now validated to 1..=1024 in the engine. Say so where the option is documented, and show the Rust CypherEngine::with_target_partitions constructor in the Rust guide. --- docs/api-reference.md | 4 ++-- docs/cypher.md | 2 +- docs/lance.md | 2 +- docs/rust.md | 4 +++- 4 files changed, 7 insertions(+), 5 deletions(-) diff --git a/docs/api-reference.md b/docs/api-reference.md index 400d26e..905be47 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -143,12 +143,12 @@ See [Repeated queries](lance.md#repeated-queries) for the scope and limits. | `config` | The graph mapping. | | `datasets` | See [Datasets](#datasets). Lance datasets are pinned to the version passed in and not read here. | | `use_csr` | Build the experimental [CSR index](csr.md) for relationship types with integer endpoints, and node indexes for labels with unique integer keys. Default `False`. | -| `target_partitions` | Optional positive integer target for DataFusion query partitioning. `None` keeps DataFusion's default. Keyword-only; preserved by `refresh()`. It is not a thread or memory limit. See [Query parallelism](lance.md#query-parallelism). | +| `target_partitions` | Optional integer target from 1 to 1024 for DataFusion query partitioning. `None` keeps DataFusion's default. Keyword-only; preserved by `refresh()`. It is not a thread or memory limit. See [Query parallelism](lance.md#query-parallelism). | Raises `ValueError` if `datasets` is empty, a configured key or endpoint column is missing, or a Lance dataset was opened through a namespace client and can't be reopened by URI. Raises `RuntimeError` if a dataset can't be opened. For `target_partitions`, booleans and non-integers other than `None` raise `TypeError`; -zero, negative values and integers outside the platform's unsigned integer range raise `ValueError`. +values outside 1 to 1024 raise `ValueError`. `#!python execute(query: str, params: dict | None = None, *, timeout: float | None = None) -> pyarrow.Table` : Runs `query` and returns the whole result. `params` binds `$name` parameters: `int`, `float`, `str`, diff --git a/docs/cypher.md b/docs/cypher.md index 7d9d613..791609a 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -45,7 +45,7 @@ CypherEngine(config, datasets, use_csr=False, target_partitions=None) The constructor validates the datasets against the config, registers them with a DataFusion session, and, with `use_csr=True`, builds the [CSR index](csr.md). After that, every query reuses the same catalog and -session. The optional keyword-only `target_partitions` sets a positive target for query partitioning; +session. The optional keyword-only `target_partitions` sets a target from 1 to 1024 for query partitioning; `None` keeps the current default. See [Query parallelism](lance.md#query-parallelism) before tuning it. !!! tip "Share one engine" diff --git a/docs/lance.md b/docs/lance.md index 5a6a83b..7989c7c 100644 --- a/docs/lance.md +++ b/docs/lance.md @@ -98,7 +98,7 @@ These rules apply with and without CSR. They do not promise a speedup for a part ## Query parallelism `CypherEngine(..., target_partitions=None)` keeps DataFusion's default query partitioning. If you have -measured your workload, you can choose a positive target explicitly: +measured your workload, you can choose a target from 1 to 1024 explicitly: ```python serial_engine = CypherEngine( diff --git a/docs/rust.md b/docs/rust.md index eeede83..f4b45ff 100644 --- a/docs/rust.md +++ b/docs/rust.md @@ -52,7 +52,9 @@ Each input is an `EngineSource`: The third argument of `CypherEngine::new` is `use_csr`. With `true`, relationship types with integer endpoints get [CSR adjacencies](csr.md) and node labels with unique integer ids get a node index, exactly as with -`use_csr=True` in Python. +`use_csr=True` in Python. `CypherEngine::with_target_partitions(config, sources, use_csr, Some(n))` also sets +DataFusion's query partition target (1 to 1024), like `target_partitions=` in Python; see +[Query parallelism](lance.md#query-parallelism). ## Run queries From bec6758b48b580a1bc4f23e35a584c9ba06a8387 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Mon, 26 Oct 2026 10:59:39 -0700 Subject: [PATCH 66/66] feat(rust): document cooperative query deadlines for Rust callers with_execution_deadline was hidden plumbing for the Python bridge. Make it a documented API: nested-loop joins inside it check the deadline and fail with "Query timeout exceeded", and cross joins yield so an outer tokio::time::timeout_at can stop them. Add a test, a Timeouts section to the Rust guide with a compiled example, and a link from the Python deadline docs. --- crates/lance-csr/examples/embedded_query.rs | 29 ++++++++++++++-- crates/lance-csr/src/engine.rs | 27 +++++++++++++++ crates/lance-csr/src/query/deadline.rs | 37 +++++++++++++++++++-- docs/cypher.md | 1 + docs/rust.md | 20 +++++++++-- 5 files changed, 106 insertions(+), 8 deletions(-) diff --git a/crates/lance-csr/examples/embedded_query.rs b/crates/lance-csr/examples/embedded_query.rs index d3a0f8d..ba9e86d 100644 --- a/crates/lance-csr/examples/embedded_query.rs +++ b/crates/lance-csr/examples/embedded_query.rs @@ -1,9 +1,10 @@ //! Run Cypher over Lance datasets from Rust, with no Python. //! //! Writes a small social graph as two Lance datasets in a temporary directory, opens -//! them in a CSR-backed engine, runs a two-hop query, then appends an edge and picks -//! it up with `refresh()`. The documentation includes this file section by section -//! (the `--8<--` markers), so keep it compiling and its printed output in step. +//! them in a CSR-backed engine, runs a two-hop query, appends an edge and picks it up +//! with `refresh()`, and runs a query under a deadline. The documentation includes this +//! file section by section (the `--8<--` markers), so keep it compiling and its printed +//! output in step. //! //! ```bash //! cargo run -p lance-csr --example embedded_query @@ -11,12 +12,15 @@ use std::collections::HashMap; use std::sync::Arc; +use std::time::Duration; use arrow::util::pretty::print_batches; use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator, StringArray}; use arrow_schema::{DataType, Field, Schema}; use lance::dataset::{Dataset, WriteMode, WriteParams}; +use lance_csr::query::with_execution_deadline; use lance_csr::{CypherEngine, EngineSource, GraphConfig}; +use tokio::time::{timeout_at, Instant}; fn knows(src: Vec, dst: Vec) -> lance_csr::Result { let schema = Arc::new(Schema::new(vec![ @@ -117,5 +121,24 @@ async fn main() -> lance_csr::Result<()> { ); // 5 edges, update=delta, version=Some(2) // --8<-- [end:refresh] + + // --8<-- [start:timeout] + // One deadline for both: checks inside nested-loop joins, and a timer that stops + // everything else (cross joins yield to it between batches). + let deadline = Instant::now() + Duration::from_secs(5); + let pairs = engine.execute( + "MATCH (a:Person), (b:Person) RETURN count(*) AS pairs", + HashMap::new(), + ); + match timeout_at(deadline, with_execution_deadline(Some(deadline), pairs)).await { + Ok(result) => print_batches(&[result?])?, + Err(_) => println!("query timed out"), + } + // +-------+ + // | pairs | + // +-------+ + // | 16 | + // +-------+ + // --8<-- [end:timeout] Ok(()) } diff --git a/crates/lance-csr/src/engine.rs b/crates/lance-csr/src/engine.rs index 655b74a..ef1d047 100644 --- a/crates/lance-csr/src/engine.rs +++ b/crates/lance-csr/src/engine.rs @@ -1200,6 +1200,33 @@ mod tests { assert_eq!(cached(), PLAN_CACHE_ENTRIES); } + #[tokio::test] + async fn execution_deadline_stops_nested_loop_joins() { + let engine = CypherEngine::new(person_config(), people_engine_sources(), false) + .await + .unwrap(); + // A non-equality join plans as a nested-loop join, whose filter checks the deadline. + let query = "MATCH (a:Person), (b:Person) WHERE a.id < b.id RETURN count(*) AS n"; + let expired = crate::query::with_execution_deadline( + Some(tokio::time::Instant::now()), + engine.execute(query, HashMap::new()), + ) + .await; + let message = expired.unwrap_err().to_string(); + assert!(message.contains("Query timeout exceeded"), "{message}"); + // No deadline, or a distant one, leaves the same query untouched. + let later = tokio::time::Instant::now() + std::time::Duration::from_secs(60); + for deadline in [None, Some(later)] { + let batch = crate::query::with_execution_deadline( + deadline, + engine.execute(query, HashMap::new()), + ) + .await + .unwrap(); + assert_eq!(batch.num_rows(), 1); + } + } + async fn write_people(uri: &str, ids: Vec, mode: WriteMode) { let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); let batch = diff --git a/crates/lance-csr/src/query/deadline.rs b/crates/lance-csr/src/query/deadline.rs index cb6b26c..47ae82c 100644 --- a/crates/lance-csr/src/query/deadline.rs +++ b/crates/lance-csr/src/query/deadline.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The lance-csr Authors -//! Internal deadline propagation from the Python bridge into physical join work. +//! Cooperative query deadlines: checkpoints in join work that stop a query once its deadline passes. use std::future::Future; use std::sync::Arc; @@ -28,8 +28,39 @@ tokio::task_local! { static EXECUTION_DEADLINE: Option; } -/// Internal cross-crate bridge plumbing; no deadline is published into an engine snapshot. -#[doc(hidden)] +/// Run `future` with a cooperative query deadline. +/// +/// Queries planned inside `future` (for example [`CypherEngine::execute`] or +/// [`CypherEngine::execute_stream`]) get deadline checks in their joins: nested-loop +/// joins (joins with non-equality conditions) check `deadline` while evaluating +/// candidates and fail with a `"Query timeout exceeded"` execution error, even when no +/// rows match; cross joins yield to the runtime between batches. `None` sets no +/// deadline. The deadline is never stored in the engine; it applies only to work inside +/// `future`. +/// +/// Yielding only lets a timer stop the work, and other operators are not checked, so +/// pair this with [`tokio::time::timeout_at`] on the same deadline: +/// +/// ```no_run +/// # use std::collections::HashMap; +/// # async fn example(engine: &lance_csr::CypherEngine) -> lance_csr::Result<()> { +/// use std::time::Duration; +/// use lance_csr::query::with_execution_deadline; +/// use tokio::time::{timeout_at, Instant}; +/// +/// let deadline = Instant::now() + Duration::from_secs(2); +/// let text = "MATCH (a), (b) WHERE a.id < b.id RETURN count(*) AS n"; +/// let query = engine.execute(text, HashMap::new()); +/// match timeout_at(deadline, with_execution_deadline(Some(deadline), query)).await { +/// Ok(result) => println!("{:?}", result?), +/// Err(_) => println!("timed out"), +/// } +/// # Ok(()) +/// # } +/// ``` +/// +/// [`CypherEngine::execute`]: crate::CypherEngine::execute +/// [`CypherEngine::execute_stream`]: crate::CypherEngine::execute_stream pub async fn with_execution_deadline(deadline: Option, future: F) -> F::Output { EXECUTION_DEADLINE.scope(deadline, future).await } diff --git a/docs/cypher.md b/docs/cypher.md index 791609a..7aed051 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -141,6 +141,7 @@ Initial expiry raises `TimeoutError`. Reading an expired native Arrow reader rai with `Query timeout exceeded`; partial output already read remains valid. Close readers when you stop consuming them. Timed cross joins yield between batches, and nested-loop joins check the deadline while evaluating candidates, including when no rows match. See the [API reference](api-reference.md) for validation and cooperative cancellation limits. +From Rust, use `with_execution_deadline`; see [Timeouts](rust.md#timeouts). ### `explain(query, params=None)` diff --git a/docs/rust.md b/docs/rust.md index f4b45ff..e279e9a 100644 --- a/docs/rust.md +++ b/docs/rust.md @@ -107,8 +107,22 @@ different exception depending on the variant: - Building the CSR in `new()` and `refresh()` is CPU-heavy work inside an async function, so it holds a runtime worker thread until it finishes. In an application with its own runtime (a desktop or web backend), run those two calls with `tokio::task::spawn_blocking` so other tasks keep running. -- Dropping a query's future, or the stream from `execute_stream`, cancels the rest of that query. There are no - built-in timeouts; wrap a call in `tokio::time::timeout` if you need one. +- Dropping a query's future, or the stream from `execute_stream`, cancels the rest of that query. + +## Timeouts + +`lance_csr::query::with_execution_deadline` runs a query with a cooperative deadline, the same mechanism behind +Python's `timeout=`. Inside it, nested-loop joins (joins with non-equality conditions) check the deadline while +evaluating candidates and fail with a `Query timeout exceeded` error, even when no rows match, and cross joins +yield to the runtime between batches. Yielding only lets a timer stop the work, and other operators aren't +checked, so also race the query against the same deadline with `tokio::time::timeout_at`: + +```rust +--8<-- "crates/lance-csr/examples/embedded_query.rs:timeout" +``` + +The deadline applies only to work inside the call; it isn't stored in the engine. A timed-out query surfaces as +`Err(Elapsed)` from `timeout_at`, or as an execution error from a join checkpoint; treat both as a timeout. ## Python and Rust side by side @@ -120,6 +134,8 @@ different exception depending on the variant: | `engine.execute(query, params)` → `pyarrow.Table` | `engine.execute(query, params).await` → `RecordBatch` | | `engine.execute_batches(...)` → `RecordBatchReader` | `engine.execute_stream(...).await` → stream | | `engine.refresh()` | `engine.refresh().await` | +| `engine.execute(query, timeout=2.0)` | `timeout_at(deadline, with_execution_deadline(Some(deadline), engine.execute(...)))` | +| `CypherEngine(..., target_partitions=4)` | `CypherEngine::with_target_partitions(config, sources, use_csr, Some(4)).await` | | `engine.csr_stats()` → `dict` | `engine.csr_stats()` → `HashMap` | | `ValueError`, `RuntimeError`, ... | `GraphError` variants |