diff --git a/Cargo.lock b/Cargo.lock index d92f820..c11b0c0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3981,6 +3981,7 @@ dependencies = [ "arrow-array", "arrow-schema", "criterion", + "datafusion", "futures", "lance", "lance-csr", diff --git a/crates/lance-csr-benches/Cargo.toml b/crates/lance-csr-benches/Cargo.toml index f0f4e94..6ebe260 100644 --- a/crates/lance-csr-benches/Cargo.toml +++ b/crates/lance-csr-benches/Cargo.toml @@ -10,6 +10,7 @@ lance-csr = { path = "../lance-csr", version = "0.3.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-csr-benches/README.md b/crates/lance-csr-benches/README.md index 2e83cf8..16f2d1b 100644 --- a/crates/lance-csr-benches/README.md +++ b/crates/lance-csr-benches/README.md @@ -26,3 +26,41 @@ cargo bench -p lance-csr-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-csr-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-csr 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. + + +## 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-csr-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-csr-benches/benches/graph_execution.rs b/crates/lance-csr-benches/benches/graph_execution.rs index b1197c1..06595c6 100644 --- a/crates/lance-csr-benches/benches/graph_execution.rs +++ b/crates/lance-csr-benches/benches/graph_execution.rs @@ -336,5 +336,119 @@ 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_csr::{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(); +} + +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-csr-python/Cargo.lock b/crates/lance-csr-python/Cargo.lock index 693777c..8645b15 100644 --- a/crates/lance-csr-python/Cargo.lock +++ b/crates/lance-csr-python/Cargo.lock @@ -3854,6 +3854,7 @@ dependencies = [ "serde", "serde_json", "snafu 0.8.9", + "tokio", ] [[package]] diff --git a/crates/lance-csr-python/Cargo.toml b/crates/lance-csr-python/Cargo.toml index b67264b..8767642 100644 --- a/crates/lance-csr-python/Cargo.toml +++ b/crates/lance-csr-python/Cargo.toml @@ -21,4 +21,4 @@ lance-csr = { path = "../lance-csr" } 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-csr-python/src/executor.rs b/crates/lance-csr-python/src/executor.rs index cf95aaf..2028862 100644 --- a/crates/lance-csr-python/src/executor.rs +++ b/crates/lance-csr-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,22 @@ 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 = lance_csr::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)) } else { @@ -65,7 +127,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 +146,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-csr-python/src/graph.rs b/crates/lance-csr-python/src/graph.rs index 3c93708..9bb3d44 100644 --- a/crates/lance-csr-python/src/graph.rs +++ b/crates/lance-csr-python/src/graph.rs @@ -26,19 +26,20 @@ use datafusion::execution::context::SessionContext; use futures::StreamExt; use lance_csr::{ ast::{DistanceMetric as RustDistanceMetric, GraphPattern, ReadingClause}, - engine::{CypherEngine as RustCypherEngine, EngineSource}, + engine::{CypherEngine as RustCypherEngine, EngineSource, MAX_TARGET_PARTITIONS}, CypherQuery as RustCypherQuery, ExecutionStrategy as RustExecutionStrategy, GraphConfig as RustGraphConfig, GraphError as RustGraphError, SqlDialect as RustSqlDialect, SqlQuery as RustSqlQuery, VectorSearch as RustVectorSearch, }; use pyo3::{ - exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, + exceptions::{PyNotImplementedError, PyRuntimeError, PyTypeError, PyValueError}, prelude::*, - types::PyDict, + types::{PyBool, PyDict, PyInt, PyList}, IntoPyObject, }; use serde_json::Value as JsonValue; +use crate::executor::{QueryDeadline, QUERY_TIMEOUT_MESSAGE}; use crate::namespace::PyDirNamespace; use crate::RT; @@ -586,7 +587,7 @@ impl CypherQuery { Ok(dict.unbind()) } - /// Convert query to SQL + /// Export supported Cypher expressions as SQL /// /// Parameters /// ---------- @@ -602,6 +603,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))] @@ -654,13 +657,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)?; @@ -672,7 +678,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) @@ -847,10 +857,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::() { @@ -860,6 +885,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. @@ -888,9 +928,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()) } } } @@ -1231,6 +1281,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.as_ref().is_none_or(|result| result.is_err()) { + 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 @@ -1281,23 +1445,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 + /// Target for DataFusion query partitioning, from 1 to 1024, 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 + /// If no datasets are given, a configured id/endpoint column is missing, a /// Lance dataset cannot be reopened by URI (for example, namespace-vended - /// credentials) + /// credentials), or target_partitions is outside 1 to 1024 + /// 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", + )); + } + // Out-of-range values that still fit a usize are rejected by the engine. + value + .extract::() + .ok() + .filter(|value| *value > 0) + .ok_or_else(|| { + PyValueError::new_err(format!( + "target_partitions must be between 1 and {MAX_TARGET_PARTITIONS}" + )) + }) + }) + .transpose()?; let mut sources = HashMap::new(); for (key, value) in datasets.iter() { let name: String = key.extract()?; @@ -1307,7 +1496,12 @@ impl CypherEngine { let inner = RT .block_on( Some(py), - RustCypherEngine::new(config.inner.clone(), sources, use_csr), + RustCypherEngine::with_target_partitions( + config.inner.clone(), + sources, + use_csr, + target_partitions, + ), )? .map_err(graph_error_to_pyerr)?; Ok(Self { inner }) @@ -1323,25 +1517,33 @@ 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). - /// ``execute_batches`` and ``explain`` take the same argument. + /// ``None``, lists, string-keyed maps, and ``datetime``/``date`` values (compared + /// as timestamps and dates). ``execute_batches`` and ``explain`` take the same argument. + /// timeout : float, optional + /// Nonnegative number of seconds for the whole query; expiry raises + /// ``TimeoutError``. Cancellation is cooperative. /// /// Raises /// ------ /// ValueError /// If the query is invalid + /// TimeoutError + /// If the query runs past ``timeout`` /// 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 params = python_params(params)?; let result_batch = RT - .block_on(Some(py), self.inner.execute(query, params))? + .block_on_until(Some(py), deadline, self.inner.execute(query, params))? .map_err(graph_error_to_pyerr)?; record_batch_to_python_table(py, &result_batch) } @@ -1352,23 +1554,26 @@ 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 params = python_params(params)?; let stream = RT - .block_on(Some(py), self.inner.execute_stream(query, params))? + .block_on_until(Some(py), deadline, self.inner.execute_stream(query, params))? .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()) } @@ -1729,6 +1934,97 @@ pub fn register_graph_module(py: Python, parent_module: &Bound<'_, PyModule>) -> mod tests { use super::*; + 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 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/crates/lance-csr/Cargo.toml b/crates/lance-csr/Cargo.toml index 7c70898..12b0f61 100644 --- a/crates/lance-csr/Cargo.toml +++ b/crates/lance-csr/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", @@ -40,6 +41,7 @@ nom = "7.1" serde = { version = "1", features = ["derive"] } serde_json = "1" snafu = "0.8" +tokio = { version = "1.37", features = ["rt", "time"] } [dev-dependencies] futures = "0.3" diff --git a/crates/lance-csr/README.md b/crates/lance-csr/README.md index 3225f65..2df4bbb 100644 --- a/crates/lance-csr/README.md +++ b/crates/lance-csr/README.md @@ -167,6 +167,34 @@ 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. + +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-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/apply.rs b/crates/lance-csr/src/apply.rs new file mode 100644 index 0000000..33d0843 --- /dev/null +++ b/crates/lance-csr/src/apply.rs @@ -0,0 +1,527 @@ +// 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, + iterations_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 { + // 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?; + 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; + } + } + 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, + iterations_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_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::{Future, 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 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(); + 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()); + // 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!(batches.is_empty()); + } + + #[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}"); + } + } + + #[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 (config, _) = graph(); + let catalog = Arc::new(lance_csr_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", + "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().with_config(config.clone()); + 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-csr/src/ast.rs b/crates/lance-csr/src/ast.rs index 52d40c7..c3930b7 100644 --- a/crates/lance-csr/src/ast.rs +++ b/crates/lance-csr/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,19 +51,60 @@ 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, + /// 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 { + 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), + ) + } + + /// 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(); // 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 { @@ -71,19 +133,40 @@ impl CypherQuery { } } } + for branch in &self.union_branches { + for label in branch.query.get_node_labels() { + if !labels.contains(&label) { + labels.push(label); + } + } + } labels } /// 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); } } } + 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 } @@ -109,6 +192,49 @@ impl CypherQuery { pub enum ReadingClause { Match(MatchClause), Unwind(UnwindClause), + /// Read-only subquery, with an optional explicit importing WITH. + Call(Box), +} + +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 + .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 @@ -158,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 @@ -422,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), @@ -476,7 +622,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, @@ -641,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(_) @@ -650,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-csr/src/csr_expand.rs b/crates/lance-csr/src/csr_expand.rs index 41d811a..8dedce6 100644 --- a/crates/lance-csr/src/csr_expand.rs +++ b/crates/lance-csr/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}; @@ -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"; @@ -230,6 +245,16 @@ 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. + #[inline(always)] + 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); @@ -239,7 +264,7 @@ impl Adjacency { .unwrap_or_default() .iter() .chain(added.get(added_offset..).unwrap_or_default()) - .filter_map(move |&neighbor| { + .map(move |&neighbor| { let position = *offset; *offset += 1; if removed.first() == Some(&position) { @@ -491,6 +516,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 { @@ -533,6 +560,8 @@ impl NodeRows { data: NodeData::Lance(dataset), schema, id_index, + #[cfg(test)] + fetched_cells: std::sync::atomic::AtomicUsize::new(0), })) } @@ -551,6 +580,8 @@ impl NodeRows { schema: batch.schema(), data: NodeData::Memory(batch), id_index, + #[cfg(test)] + fetched_cells: std::sync::atomic::AtomicUsize::new(0), })) } @@ -558,7 +589,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() } @@ -574,6 +605,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()); @@ -777,7 +813,7 @@ pub struct Neighbors { /// Occurrence offsets are sufficient to pause even a high-degree source row. #[derive(Clone, Copy, Default)] -struct NeighborPosition { +pub(crate) struct NeighborPosition { forward: usize, reverse: usize, } @@ -803,21 +839,31 @@ impl Neighbors { } #[inline(always)] - fn neighbors_from<'a>( + pub(crate) 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); + self.candidates_from(key, position).flatten() + } + + /// 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, + 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)), ) } @@ -864,6 +910,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, } @@ -893,6 +940,7 @@ impl CsrExpandNode { target, target_filter: None, closing_column: None, + traversal: None, schema: Arc::new(schema), }) } @@ -903,6 +951,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( @@ -915,6 +968,7 @@ impl CsrExpandNode { )?; node.target_filter = self.target_filter.clone(); node.closing_column = self.closing_column.clone(); + node.traversal = self.traversal.clone(); Ok(node) } @@ -970,6 +1024,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 {} @@ -1072,6 +1131,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(()) } @@ -1096,6 +1158,7 @@ pub struct CsrExpandExec { /// The moved target filter, if any. target_filter: Option, closing_column: Option, + traversal: Option>, schema: SchemaRef, properties: Arc, } @@ -1109,6 +1172,7 @@ impl CsrExpandExec { target: Option, target_filter: Option, closing_column: Option, + traversal: Option>, schema: SchemaRef, ) -> Self { let properties = Arc::new(PlanProperties::new( @@ -1124,6 +1188,7 @@ impl CsrExpandExec { target, target_filter, closing_column, + traversal, schema, properties, } @@ -1137,6 +1202,7 @@ struct BatchExpansion { bound: Option, row: usize, position: NeighborPosition, + walk: Option, } impl BatchExpansion { @@ -1156,35 +1222,62 @@ impl BatchExpansion { bound, row: 0, position: NeighborPosition::default(), + walk: None, }) } /// 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(); + 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; '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 neighbor in adjacency.neighbors_from(key, &mut position) { + for candidate in adjacency + .candidates_from(key, &mut position) + .take(1024 - work) + { + work += 1; + let Some(neighbor) = candidate else { + continue; + }; 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 { + 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); @@ -1193,12 +1286,69 @@ impl BatchExpansion { break 'rows; } } + if work == 1024 { + break; + } + } + if work == previous_work { + work += 1; } row += 1; position = NeighborPosition::default(); + if work == 1024 { + break; + } } self.row = row; self.position = position; + work == 1024 + } + + async fn next_path_chunk( + &mut self, + graph: &crate::csr_traversal::Traversal, + target: Option<&TargetLookup>, + 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); + 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 } } @@ -1227,15 +1377,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)? @@ -1244,6 +1419,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); @@ -1268,7 +1448,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 @@ -1280,6 +1460,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") }); @@ -1314,6 +1496,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(()) } } @@ -1332,6 +1517,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) => { @@ -1365,6 +1553,7 @@ impl ExecutionPlan for CsrExpandExec { self.target.clone(), self.target_filter.clone(), self.closing_column.clone(), + self.traversal.clone(), self.schema.clone(), ))) } @@ -1388,6 +1577,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)| { @@ -1397,6 +1587,7 @@ impl ExecutionPlan for CsrExpandExec { target_filter.clone(), schema.clone(), ); + let traversal = traversal.clone(); async move { loop { if pending.is_none() { @@ -1406,7 +1597,18 @@ 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) + .await + } + }; if chunk.rows.is_empty() { pending = None; continue; @@ -1441,6 +1643,20 @@ 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::() + { + return crate::path_match::plan_path_match(node, physical_inputs, session_state) + .map(Some); + } if let Some(node) = node .as_any() .downcast_ref::() @@ -1464,13 +1680,14 @@ impl ExtensionPlanner for CsrExtensionPlanner { node.target.clone(), target_filter, node.closing_column.clone(), + node.traversal.clone(), Arc::new(node.schema.as_arrow().clone()), )))) } } /// 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; @@ -1873,21 +2090,170 @@ 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 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( + 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] + 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] @@ -1906,7 +2272,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::>() @@ -1915,7 +2281,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); } @@ -1950,13 +2319,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()); } @@ -1984,4 +2354,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/crates/lance-csr/src/csr_traversal.rs b/crates/lance-csr/src/csr_traversal.rs new file mode 100644 index 0000000..cdf11cb --- /dev/null +++ b/crates/lance-csr/src/csr_traversal.rs @@ -0,0 +1,376 @@ +// 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 + .candidates_from(frame.key, &mut frame.position) + .next() + }) + .flatten(); + 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; + } + // 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 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; + 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-csr/src/datafusion_planner/analysis.rs b/crates/lance-csr/src/datafusion_planner/analysis.rs index 3329517..76ab3fe 100644 --- a/crates/lance-csr/src/datafusion_planner/analysis.rs +++ b/crates/lance-csr/src/datafusion_planner/analysis.rs @@ -37,8 +37,10 @@ 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, pub(crate) comprehension_elements: Vec, } @@ -46,8 +48,10 @@ 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, } } @@ -73,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(); @@ -144,6 +162,7 @@ fn analyze_operator( input, source_variable, target_variable, + target_label, relationship_types, direction, relationship_variable, @@ -154,13 +173,11 @@ 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); - } + 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 @@ -205,7 +222,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)?; } @@ -223,6 +241,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)?; } @@ -256,6 +286,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(()) @@ -371,6 +416,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-csr/src/datafusion_planner/builder/aggregate_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/aggregate_ops.rs index 0cf71fc..e18d4fe 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/aggregate_ops.rs +++ b/crates/lance-csr/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 = super::super::expression::to_df_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(super::super::expression::to_df_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(super::super::expression::to_df_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 = super::super::expression::to_df_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-csr/src/datafusion_planner/builder/basic_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/basic_ops.rs index c10b116..cd4b6f5 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-csr/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 @@ -28,6 +56,9 @@ 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.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. ctx.relationship_instance_idx = branch.relationship_instance_idx; @@ -79,6 +110,79 @@ 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(); + 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; + } + 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!()), + }); + } + } + // 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 { + 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, @@ -86,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 = super::super::expression::to_df_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))? @@ -119,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)?; @@ -174,8 +286,7 @@ impl DataFusionPlanner { exprs.extend(columns); 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(&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 { @@ -219,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)) } @@ -234,7 +347,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.unresolved_value_expr(&bound), asc: matches!(item.direction, crate::ast::SortDirection::Ascending), nulls_first: true, } @@ -288,8 +401,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(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-csr/src/datafusion_planner/builder/call_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/call_ops.rs new file mode 100644 index 0000000..d93e67d --- /dev/null +++ b/crates/lance-csr/src/datafusion_planner/builder/call_ops.rs @@ -0,0 +1,79 @@ +// 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); + 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))?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(node), + })) + } +} diff --git a/crates/lance-csr/src/datafusion_planner/builder/comprehension_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/comprehension_ops.rs index 43b77d2..51c99bb 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/comprehension_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/comprehension_ops.rs @@ -9,14 +9,14 @@ use arrow_schema::DataType; use datafusion::common::Column; use datafusion::functions::core::expr_fn::coalesce; use datafusion::functions_aggregate::array_agg::array_agg; -use datafusion::functions_nested::expr_fn::{array_element, cardinality, gen_series, make_array}; +use datafusion::functions_aggregate::count::count; +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, to_df_boolean_expr, to_df_value_expr}; use crate::ast::{BooleanExpression, ValueExpression}; use crate::case_insensitive::qualify_column; use crate::datafusion_planner::DataFusionPlanner; @@ -42,6 +42,7 @@ impl DataFusionPlanner { keep: Option<&BooleanExpression>, map: &ValueExpression, output: &str, + count_only: bool, ) -> Result { let error = |what: &str| { let what = what.to_string(); @@ -54,8 +55,17 @@ 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 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()) @@ -130,14 +140,20 @@ 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(keep, elements.schema())?); } - let value = to_df_value_expr(&bind_value(map, elements.schema())); - let collected = array_agg(value) - .filter(kept) - .order_by(vec![col(&pos).sort(true, false)]) - .build() - .map_err(error("Failed to collect comprehension values"))? + 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(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 +176,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 +189,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-csr/src/datafusion_planner/builder/expand_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs index 4f4f943..ed30c57 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/expand_ops.rs @@ -420,11 +420,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, @@ -438,7 +434,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 @@ -479,6 +475,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(); @@ -558,6 +578,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()]; @@ -618,6 +691,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. @@ -726,6 +923,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()), @@ -761,6 +959,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, @@ -796,6 +995,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, @@ -845,6 +1045,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, @@ -893,6 +1094,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, @@ -942,6 +1144,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, @@ -991,6 +1194,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, @@ -1040,6 +1244,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-csr/src/datafusion_planner/builder/helpers.rs b/crates/lance-csr/src/datafusion_planner/builder/helpers.rs index 34647d5..f2e5721 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/helpers.rs +++ b/crates/lance-csr/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); } @@ -95,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, @@ -112,6 +122,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 => {} } } @@ -140,12 +155,15 @@ impl DataFusionPlanner { 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); - } + 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, .. } @@ -153,7 +171,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); } @@ -173,9 +192,28 @@ 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), + 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 => {} } } @@ -303,6 +341,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-csr/src/datafusion_planner/builder/join_builder.rs b/crates/lance-csr/src/datafusion_planner/builder/join_builder.rs index 4c74348..427932b 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/join_builder.rs +++ b/crates/lance-csr/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-csr/src/datafusion_planner/builder/mod.rs b/crates/lance-csr/src/datafusion_planner/builder/mod.rs index c2ca6a4..bf4f379 100644 --- a/crates/lance-csr/src/datafusion_planner/builder/mod.rs +++ b/crates/lance-csr/src/datafusion_planner/builder/mod.rs @@ -14,10 +14,12 @@ mod aggregate_ops; mod basic_ops; +mod call_ops; mod comprehension_ops; mod expand_ops; mod helpers; mod join_builder; +mod path_ops; use super::DataFusionPlanner; use crate::error::Result; @@ -34,6 +36,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)), @@ -108,6 +120,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, @@ -117,6 +132,7 @@ impl DataFusionPlanner { keep, map, output, + count_only, } => self.build_list_comprehension( ctx, input, @@ -126,6 +142,7 @@ impl DataFusionPlanner { keep.as_ref(), map, output, + *count_only, ), LogicalOperator::ComprehensionElement => Ok(ctx .comprehension_elements @@ -137,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-csr/src/datafusion_planner/builder/path_ops.rs b/crates/lance-csr/src/datafusion_planner/builder/path_ops.rs new file mode 100644 index 0000000..7cf18c2 --- /dev/null +++ b/crates/lance-csr/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-csr/src/datafusion_planner/cypher_udf.rs b/crates/lance-csr/src/datafusion_planner/cypher_udf.rs index e4f47a2..171e787 100644 --- a/crates/lance-csr/src/datafusion_planner/cypher_udf.rs +++ b/crates/lance-csr/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-csr/src/datafusion_planner/expression.rs b/crates/lance-csr/src/datafusion_planner/expression.rs index 8b2bd0a..7eba94a 100644 --- a/crates/lance-csr/src/datafusion_planner/expression.rs +++ b/crates/lance-csr/src/datafusion_planner/expression.rs @@ -19,6 +19,65 @@ 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 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, + 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 { + 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 { @@ -122,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), @@ -131,10 +216,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) @@ -189,22 +271,26 @@ 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), }, - "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" => { + 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" => { + 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]), @@ -411,12 +497,12 @@ 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) 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() @@ -458,6 +544,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) @@ -469,28 +576,46 @@ 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); + 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); } @@ -513,34 +638,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-csr/src/datafusion_planner/join_ops.rs b/crates/lance-csr/src/datafusion_planner/join_ops.rs index f96c187..94c1aaf 100644 --- a/crates/lance-csr/src/datafusion_planner/join_ops.rs +++ b/crates/lance-csr/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, @@ -59,9 +59,8 @@ 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.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-csr/src/datafusion_planner/list_folds.rs b/crates/lance-csr/src/datafusion_planner/list_folds.rs new file mode 100644 index 0000000..9f3a8b7 --- /dev/null +++ b/crates/lance-csr/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-csr/src/datafusion_planner/math_functions.rs b/crates/lance-csr/src/datafusion_planner/math_functions.rs new file mode 100644 index 0000000..e82f427 --- /dev/null +++ b/crates/lance-csr/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-csr/src/datafusion_planner/mod.rs b/crates/lance-csr/src/datafusion_planner/mod.rs index 15bd354..cb910d1 100644 --- a/crates/lance-csr/src/datafusion_planner/mod.rs +++ b/crates/lance-csr/src/datafusion_planner/mod.rs @@ -18,8 +18,11 @@ pub mod analysis; mod builder; mod config_helpers; mod cypher_udf; -mod expression; +pub(crate) mod expression; mod join_ops; +mod list_folds; +mod math_functions; +mod scalar_functions; mod scan_ops; mod udf; pub mod vector_ops; @@ -29,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; @@ -47,6 +51,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 +60,7 @@ impl DataFusionPlanner { config, catalog: None, adjacency: None, + parameter_types: Default::default(), } } @@ -63,6 +69,7 @@ impl DataFusionPlanner { config, catalog: Some(catalog), adjacency: None, + parameter_types: Default::default(), } } diff --git a/crates/lance-csr/src/datafusion_planner/scalar_functions.rs b/crates/lance-csr/src/datafusion_planner/scalar_functions.rs new file mode 100644 index 0000000..d877529 --- /dev/null +++ b/crates/lance-csr/src/datafusion_planner/scalar_functions.rs @@ -0,0 +1,628 @@ +// 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, AsArray, Int64Array, ListBuilder, StringArray, StringBuilder, +}; +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; +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() { + "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" => { + 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(), + _ => 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() { + "startnode" | "endnode" => match &types[0] { + DataType::Struct(fields) => fields[2].data_type().clone(), + _ => DataType::Null, + }, + "__predicate_all" | "__predicate_any" | "__predicate_none" | "__predicate_single" => { + DataType::Boolean + } + "size" => DataType::Int64, + "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))) + } + "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 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 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(); + 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" || 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)); + 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 = 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); + } + 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) + }; + } + 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) +} + +/// 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::*; + use crate::{config::GraphConfig, query::CypherQuery}; + 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)]) + .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_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 [ + ("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-csr/src/datafusion_planner/scan_ops.rs b/crates/lance-csr/src/datafusion_planner/scan_ops.rs index 83de637..aae8d8b 100644 --- a/crates/lance-csr/src/datafusion_planner/scan_ops.rs +++ b/crates/lance-csr/src/datafusion_planner/scan_ops.rs @@ -44,7 +44,7 @@ impl DataFusionPlanner { let filter_exprs: Vec = properties .iter() .map(|(k, v)| { - let lit_expr = super::expression::to_df_value_expr( + let lit_expr = self.unresolved_value_expr( &crate::ast::ValueExpression::Literal(v.clone()), ); Expr::BinaryExpr(BinaryExpr { @@ -140,9 +140,8 @@ 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.unresolved_value_expr(&crate::ast::ValueExpression::Literal(v.clone())); let filter_expr = Expr::BinaryExpr(BinaryExpr { left: Box::new(col(k)), op: Operator::Eq, @@ -358,6 +357,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-csr/src/engine.rs b/crates/lance-csr/src/engine.rs index a1aa970..ef1d047 100644 --- a/crates/lance-csr/src/engine.rs +++ b/crates/lance-csr/src/engine.rs @@ -30,13 +30,13 @@ //! # } //! ``` -use std::collections::{HashMap, HashSet}; -use std::sync::{Arc, RwLock}; +use std::collections::{HashMap, HashSet, VecDeque}; +use std::sync::{Arc, Mutex, RwLock}; use arrow_array::RecordBatch; use arrow_schema::SchemaRef; use datafusion::datasource::{DefaultTableSource, MemTable, TableProvider}; -use datafusion::execution::context::SessionContext; +use datafusion::execution::context::{SessionConfig, SessionContext}; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::execution::SendableRecordBatchStream; use futures::StreamExt; @@ -165,6 +165,8 @@ pub struct NodeIndexStats { /// keep the `Arc` they started with, so they finish on their original snapshot. struct EngineState { sources: HashMap, + /// Unoptimized plan templates for this snapshot; a refresh starts an empty cache. + plans: Mutex>, catalog: Arc, context: SessionContext, /// CSR adjacencies by relationship type; empty unless the engine uses CSR. @@ -175,6 +177,36 @@ struct EngineState { key_stats: HashMap, } +/// Plan templates kept per snapshot. +const PLAN_CACHE_ENTRIES: usize = 128; + +/// A small LRU of unoptimized plans. Values are bound after lookup, before optimization. +struct CachedQuery { + text: String, + parameter_types: Vec<(String, u8)>, + query: CypherQuery, +} + +// 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 { @@ -217,6 +249,10 @@ struct CsrEntry { // shortcut: fixed threshold, chosen from python/benchmarks/csr_delta_accumulated.py; tune per workload if deltas grow large. const CSR_REBUILD_FRACTION: f64 = 0.25; +/// Upper bound for `target_partitions`: far above any core count, below values whose +/// per-partition setup would exhaust memory instead of returning an error. +pub const MAX_TARGET_PARTITIONS: usize = 1024; + /// Apply the edge changes between `old`'s snapshot and `source` as a delta. /// /// `None` means rebuild: not both Lance snapshots, the diff is not meaningful @@ -378,12 +414,18 @@ async fn build_engine_state( config: &GraphConfig, sources: HashMap, use_csr: bool, + target_partitions: Option, previous: Option<&EngineState>, ) -> Result { + 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. @@ -445,6 +487,7 @@ async fn build_engine_state( } Ok(EngineState { sources, + plans: Mutex::default(), catalog: Arc::new(catalog), context: ctx, csr, @@ -750,6 +793,7 @@ pub struct CypherEngine { config: GraphConfig, state: RwLock>, use_csr: bool, + target_partitions: Option, /// Held for a whole refresh, so a slower refresh can't swap in older data over a newer one. refresh_lock: futures::lock::Mutex<()>, } @@ -770,15 +814,42 @@ impl CypherEngine { config: GraphConfig, sources: HashMap, use_csr: bool, + ) -> Result { + Self::with_target_partitions(config, sources, use_csr, None).await + } + + /// Like [`new`](Self::new), with a target for DataFusion query partitioning. + /// + /// `Some(n)` sets DataFusion's `target_partitions` for every query, kept across + /// [`refresh`](Self::refresh); `None` keeps DataFusion's default (the number of + /// cores). It is not a thread or memory limit. In-memory Arrow inputs are still + /// sliced into one partition per core. + /// + /// # Errors + /// As [`new`](Self::new), plus [`GraphError::ConfigError`] if `target_partitions` + /// is outside `1..=1024`. + pub async fn with_target_partitions( + config: GraphConfig, + sources: HashMap, + use_csr: bool, + target_partitions: Option, ) -> Result { if sources.is_empty() { return Err(config_error("No input datasets provided".to_string())); } - let state = build_engine_state(&config, sources, use_csr, None).await?; + if let Some(partitions) = target_partitions { + if !(1..=MAX_TARGET_PARTITIONS).contains(&partitions) { + return Err(config_error(format!( + "target_partitions must be between 1 and {MAX_TARGET_PARTITIONS}" + ))); + } + } + let state = build_engine_state(&config, sources, use_csr, target_partitions, None).await?; Ok(Self { config, state: RwLock::new(Arc::new(state)), use_csr, + target_partitions, refresh_lock: Default::default(), }) } @@ -791,21 +862,62 @@ impl CypherEngine { .clone() } + /// Parse and configure a query, reusing this snapshot's plan template when + /// `use_cache` is set and every parameter is a scalar. fn parse( &self, - query: &str, + text: &str, params: HashMap, state: &EngineState, + use_cache: bool, ) -> Result { - let query = CypherQuery::new(query)? - .with_parameters(params) + let parameters: HashMap = params + .into_iter() + .map(|(name, value)| (name.to_lowercase(), value)) + .collect(); + let key = use_cache.then(|| parameter_types(¶meters)).flatten(); + if let Some(key) = &key { + let mut plans = state.plans.lock().expect("plan cache lock poisoned"); + if let Some(index) = plans + .iter() + .position(|entry| entry.text == text && entry.parameter_types == *key) + { + let entry = plans.remove(index).expect("index from position"); + let query = entry.query.clone().with_parameters(parameters); + plans.push_back(entry); + return Ok(query); + } + } + let query = CypherQuery::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().expect("plan cache lock poisoned"); + plans + .retain(|entry| entry.text != text || entry.parameter_types != parameter_types); + if plans.len() == PLAN_CACHE_ENTRIES { + plans.pop_front(); + } + plans.push_back(CachedQuery { + text: text.to_string(), + parameter_types, + query: query.clone(), + }); + } + } + Ok(query) } /// Execute a Cypher query and return the whole result as one batch. @@ -818,7 +930,7 @@ impl CypherEngine { params: HashMap, ) -> Result { let state = self.snapshot(); - self.parse(query, params, &state)? + self.parse(query, params, &state, true)? .execute_with_catalog_and_context(state.catalog.clone(), state.context.clone()) .await } @@ -836,7 +948,7 @@ impl CypherEngine { params: HashMap, ) -> Result { let state = self.snapshot(); - self.parse(query, params, &state)? + self.parse(query, params, &state, true)? .execute_stream_with_catalog_and_context(state.catalog.clone(), state.context.clone()) .await } @@ -844,7 +956,7 @@ impl CypherEngine { /// Return the query plan, including the physical plan with pushed-down scans. pub async fn explain(&self, query: &str, params: HashMap) -> Result { let state = self.snapshot(); - self.parse(query, params, &state)? + self.parse(query, params, &state, false)? .explain_with_catalog_and_context(state.catalog.clone(), state.context.clone()) .await } @@ -855,7 +967,11 @@ impl CypherEngine { query: &str, vector_search: &VectorSearch, ) -> Result { - let candidates = self.execute(query, HashMap::new()).await?; + let state = self.snapshot(); + let candidates = self + .parse(query, HashMap::new(), &state, false)? + .execute_with_catalog_and_context(state.catalog.clone(), state.context.clone()) + .await?; vector_search.search(&candidates).await } @@ -883,7 +999,14 @@ impl CypherEngine { } } } - let state = build_engine_state(&self.config, sources, self.use_csr, Some(¤t)).await?; + let state = build_engine_state( + &self.config, + sources, + self.use_csr, + self.target_partitions, + Some(¤t), + ) + .await?; *self.state.write().expect("engine state lock poisoned") = Arc::new(state); Ok(()) } @@ -950,6 +1073,160 @@ mod tests { use arrow_schema::{DataType, Field, Schema}; use lance::dataset::{WriteMode, WriteParams}; + fn people_engine_sources() -> HashMap { + let batch = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from(vec![1, 2, 3])) as arrow_array::ArrayRef, + )]) + .unwrap(); + HashMap::from([("Person".to_string(), EngineSource::Arrow(batch))]) + } + + fn person_config() -> GraphConfig { + GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap() + } + + #[tokio::test] + async fn target_partitions_survive_refresh_without_changing_other_defaults() { + let defaults = SessionConfig::new(); + for target_partitions in [None, Some(1), Some(2)] { + let engine = CypherEngine::with_target_partitions( + person_config(), + people_engine_sources(), + false, + target_partitions, + ) + .await + .unwrap(); + let before = engine.snapshot(); + engine.refresh().await.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() + ); + } + } + for out_of_range in [0, MAX_TARGET_PARTITIONS + 1] { + let result = CypherEngine::with_target_partitions( + person_config(), + people_engine_sources(), + false, + Some(out_of_range), + ) + .await; + assert!(matches!(result, Err(GraphError::ConfigError { .. }))); + } + } + + #[tokio::test] + async fn plan_cache_reuses_bounds_and_resets_on_refresh() { + let engine = CypherEngine::new(person_config(), people_engine_sources(), false) + .await + .unwrap(); + let snapshot = engine.snapshot(); + let cached = || snapshot.plans.lock().unwrap().len(); + let text = "MATCH (p:Person) WHERE p.id = $value RETURN p.id"; + for value in [1, 2, 3] { + let params = HashMap::from([("value".to_string(), JsonValue::from(value))]); + engine.parse(text, params, &snapshot, true).unwrap(); + assert_eq!(cached(), 1); + // The retained template keeps its first parameters: later calls looked it + // up and rebound values rather than replacing the entry. + assert_eq!( + snapshot.plans.lock().unwrap()[0].query.parameters()["value"], + JsonValue::from(1) + ); + } + // A different scalar type is a different template. + for value in [JsonValue::Null, JsonValue::from(1.5)] { + let params = HashMap::from([("value".to_string(), value)]); + engine.parse(text, params, &snapshot, true).unwrap(); + } + assert_eq!(cached(), 3); + // Container parameters stay on the uncached path. + let ids = HashMap::from([("ids".to_string(), serde_json::json!([1, 2]))]); + engine + .parse( + "MATCH (p:Person) WHERE p.id IN $ids RETURN p.id", + ids, + &snapshot, + true, + ) + .unwrap(); + assert_eq!(cached(), 3); + // Explain-style parses never cache. + engine + .parse( + "MATCH (p:Person) RETURN p.id", + HashMap::new(), + &snapshot, + false, + ) + .unwrap(); + assert_eq!(cached(), 3); + for limit in 0..130 { + let text = format!("MATCH (p:Person) RETURN p.id AS id LIMIT {limit}"); + engine + .parse(&text, HashMap::new(), &snapshot, true) + .unwrap(); + } + assert_eq!(cached(), PLAN_CACHE_ENTRIES); + assert!(snapshot + .plans + .lock() + .unwrap() + .iter() + .all(|e| e.text != text)); + + // Cached and uncached plans return the same rows. + let params = || HashMap::from([("value".to_string(), JsonValue::from(2))]); + let first = engine.execute(text, params()).await.unwrap(); + let again = engine.execute(text, params()).await.unwrap(); + assert_eq!(first, again); + assert_eq!(first.num_rows(), 1); + + // A refresh publishes a new snapshot with an empty cache; the old one is untouched. + engine.refresh().await.unwrap(); + assert!(engine.snapshot().plans.lock().unwrap().is_empty()); + 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/lib.rs b/crates/lance-csr/src/lib.rs index 2481444..1516756 100644 --- a/crates/lance-csr/src/lib.rs +++ b/crates/lance-csr/src/lib.rs @@ -43,11 +43,13 @@ //! # } //! ``` +mod apply; pub mod ast; pub mod case_insensitive; pub mod config; pub mod csr_expand; pub mod csr_index; +mod csr_traversal; pub mod datafusion_planner; pub mod engine; pub mod error; @@ -56,6 +58,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-csr/src/logical_plan.rs b/crates/lance-csr/src/logical_plan.rs index fd66b89..0123ede 100644 --- a/crates/lance-csr/src/logical_plan.rs +++ b/crates/lance-csr/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, @@ -74,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, @@ -82,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 @@ -135,6 +143,12 @@ pub enum LogicalOperator { /// same variables, possibly with different labels; columns a branch lacks are null. Union { inputs: Vec }, + /// 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 { @@ -172,7 +186,21 @@ pub enum LogicalOperator { keep: Option, map: ValueExpression, output: String, + /// Count retained elements without building an ordered list. + #[serde(default)] + 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, @@ -212,14 +240,17 @@ 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. 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, + 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. @@ -232,10 +263,13 @@ impl<'a> LogicalPlanner<'a> { variables: HashMap::new(), config, fan_out: None, + adjacency: None, exists_marks: 0, node_lists: HashMap::new(), comprehensions: 0, + calls: 0, path_labels: HashMap::new(), + rich_paths: false, seed_source: None, } } @@ -246,6 +280,79 @@ 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 { + 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; + }; + 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 @@ -253,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() @@ -308,14 +418,37 @@ 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.config); + self.plan_query(query, None, false) + } + + fn plan_query( + &mut self, + query: &CypherQuery, + 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(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 { @@ -349,7 +482,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 { @@ -360,21 +506,57 @@ 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()?, }; } + for branch in &query.union_branches { + let mut planner = LogicalPlanner::new(self.config) + .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(); + 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), + all: branch.all, + }; + } Ok(plan) } @@ -408,6 +590,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) } @@ -424,6 +607,59 @@ impl<'a> LogicalPlanner<'a> { } } + fn plan_call( + &mut self, + input: Option, + query: &CypherQuery, + ) -> Result { + 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) + .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()); + } + 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, @@ -448,15 +684,15 @@ 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) => { let property = match self.variables.get(var).map(String::as_str) { - Some("Path") => Some("length".to_string()), + Some("Path") if !self.rich_paths => Some("length".to_string()), + Some("Path") => None, Some(label) => self .config .get_node_mapping(label) @@ -470,9 +706,45 @@ 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") => {} + 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)) @@ -490,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 @@ -574,11 +862,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); } @@ -628,9 +918,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)?), } } @@ -744,6 +1048,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)?; @@ -754,6 +1067,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() { @@ -765,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() { @@ -859,6 +1264,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 } @@ -901,6 +1313,7 @@ impl<'a> LogicalPlanner<'a> { keep, map: value, output: output.clone(), + count_only: false, }) })(); self.variables = before; @@ -964,6 +1377,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(), @@ -1007,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]; @@ -1026,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())), @@ -1234,6 +1667,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(), @@ -1261,14 +1695,65 @@ 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 { 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 { @@ -1283,7 +1768,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)) @@ -1298,6 +1784,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()), @@ -1339,7 +1833,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), @@ -1386,6 +1880,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, @@ -1432,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()), @@ -1510,10 +2018,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()?, }; } @@ -1530,6 +2038,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/crates/lance-csr/src/parameter_substitution.rs b/crates/lance-csr/src/parameter_substitution.rs index f66765a..e4676d6 100644 --- a/crates/lance-csr/src/parameter_substitution.rs +++ b/crates/lance-csr/src/parameter_substitution.rs @@ -7,6 +7,12 @@ 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)?; + // Substitute in READING clauses for reading_clause in &mut query.reading_clauses { substitute_in_reading_clause(reading_clause, parameters)?; @@ -53,11 +59,28 @@ 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, ) -> 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)?; @@ -143,6 +166,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)?; } @@ -240,29 +264,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, .. } => { @@ -301,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, .. } => { @@ -317,6 +335,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 +362,78 @@ 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 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(); + 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/crates/lance-csr/src/parser.rs b/crates/lance-csr/src/parser.rs index 7f215d6..716b351 100644 --- a/crates/lance-csr/src/parser.rs +++ b/crates/lance-csr/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}, @@ -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,15 @@ 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) + 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 @@ -106,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)?; @@ -140,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, @@ -153,9 +181,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)?; @@ -204,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) @@ -232,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)?; @@ -251,6 +300,7 @@ fn path_pattern(input: &str) -> IResult<&str, PathPattern> { start_node, segments, shortest: None, + variable: None, }, )) } @@ -492,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) } @@ -669,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 @@ -733,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(( @@ -1309,35 +1422,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; @@ -1380,14 +1495,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())) } @@ -1702,7 +1817,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 +1825,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 +1834,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 +1843,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 +1852,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 +2408,27 @@ 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(); + 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]", "[]"] { @@ -2488,7 +2624,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); @@ -2518,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-csr/src/path_inventory.rs b/crates/lance-csr/src/path_inventory.rs new file mode 100644 index 0000000..713a1db --- /dev/null +++ b/crates/lance-csr/src/path_inventory.rs @@ -0,0 +1,317 @@ +// 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).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()); + 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, + }) + } +} + +#[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/crates/lance-csr/src/path_match.rs b/crates/lance-csr/src/path_match.rs new file mode 100644 index 0000000..07895b1 --- /dev/null +++ b/crates/lance-csr/src/path_match.rs @@ -0,0 +1,1009 @@ +// 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(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 }, +} + +// 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()))?; + 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 { + // 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() + .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 { + 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)>, +} + +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 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) + } +} +impl ExecutionPlan for PathMatchExec { + fn name(&self) -> &str { + "PathMatchExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + 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, + 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-csr/src/path_value.rs b/crates/lance-csr/src/path_value.rs new file mode 100644 index 0000000..6f0ee44 --- /dev/null +++ b/crates/lance-csr/src/path_value.rs @@ -0,0 +1,422 @@ +// 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, HashSet}; +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_csr.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), + ], + ))) +} + +/// 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( + query: &crate::ast::CypherQuery, + config: &crate::config::GraphConfig, +) -> 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 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()); + } + } + } + } + } + } + // 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() + .any(|branch| needs_values(&branch.query, config)) + || clauses.iter().any(|clause| match clause { + ReadingClause::Call(inner) => { + needs_values(inner, config) + || 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::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") => + { + 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 { + 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 { + 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 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), + ("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), + ( + "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 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, + ), + ("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, &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-csr/src/pattern_alternatives.rs b/crates/lance-csr/src/pattern_alternatives.rs index f2d1356..64b37df 100644 --- a/crates/lance-csr/src/pattern_alternatives.rs +++ b/crates/lance-csr/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-csr/src/pattern_order.rs b/crates/lance-csr/src/pattern_order.rs index b5333e4..818b9dd 100644 --- a/crates/lance-csr/src/pattern_order.rs +++ b/crates/lance-csr/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; @@ -258,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(_) => {} } @@ -415,6 +433,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-csr/src/query.rs b/crates/lance-csr/src/query.rs index 4d95cd5..524dde1 100644 --- a/crates/lance-csr/src/query.rs +++ b/crates/lance-csr/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; @@ -129,7 +134,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 +164,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 +181,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 +189,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 +233,82 @@ 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 { + // 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 = 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() + || crate::path_value::needs_values(&self.ast, self.require_config()?) + { + // CALL binds its own child templates; row counts alter plan shape. + 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) + } + + 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 { @@ -377,11 +479,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, ScalarValue}; + use datafusion::logical_expr::expr_rewriter::NamePreserver; + use datafusion::logical_expr::registry::FunctionRegistry; + use datafusion::logical_expr::{lit, utils::merge_schema, Expr, LogicalPlan}; use std::sync::Arc; let dialect = dialect.unwrap_or_default(); @@ -405,6 +516,56 @@ 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() + ); + } + } + if let Expr::Literal(value, _) = &expr { + let data_type = value.data_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), + ]))); + } + } + 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(); @@ -565,14 +726,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 @@ -580,9 +751,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!()), }) @@ -625,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 { @@ -685,7 +857,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 @@ -765,7 +942,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(); @@ -876,6 +1058,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; @@ -883,7 +1096,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 { @@ -893,13 +1106,27 @@ 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 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 { + &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)) @@ -919,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 @@ -1043,13 +1271,15 @@ impl CypherQuery { pub fn referenced_node_labels(&self) -> Vec { let mut labels = Vec::new(); - for clause in &self.ast.reading_clauses { - 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(); @@ -1060,13 +1290,15 @@ impl CypherQuery { pub fn referenced_relationship_types(&self) -> Vec { let mut types = Vec::new(); - for clause in &self.ast.reading_clauses { - 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(); @@ -1077,18 +1309,29 @@ impl CypherQuery { pub fn variables(&self) -> Vec { let mut variables = Vec::new(); - for clause in &self.ast.reading_clauses { - match clause { - 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(); @@ -1300,6 +1543,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, @@ -1311,8 +1555,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) @@ -1325,6 +1569,7 @@ impl CypherQueryBuilder { parameters: self.parameters, adjacency: None, fan_out: None, + cached_plan: None, }; Ok(query) @@ -1336,6 +1581,145 @@ 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 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", + "RETURN 1 AS n UNION ALL RETURN 1 AS n UNION RETURN 2 AS n", + ] { + assert!(CypherQuery::new(query) + .unwrap() + .with_config(config.clone()) + .execute(HashMap::from([("Person".to_string(), batch.clone())]), None) + .await + .is_err()); + } + } + + #[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_csr_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(tokio::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(tokio::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(|| tokio::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(); @@ -1360,6 +1744,290 @@ 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); + 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) + .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", true) + .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()); + } + + #[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()); + } + } + + #[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(). + #[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() @@ -2000,6 +2668,76 @@ 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'")); + } + + #[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/crates/lance-csr/src/query/deadline.rs b/crates/lance-csr/src/query/deadline.rs new file mode 100644 index 0000000..47ae82c --- /dev/null +++ b/crates/lance-csr/src/query/deadline.rs @@ -0,0 +1,158 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The lance-csr Authors + +//! Cooperative query deadlines: checkpoints in join work that stop a query once its deadline passes. + +use std::future::Future; +use std::sync::Arc; +use tokio::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; +} + +/// 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 +} + +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.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.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-csr/src/semantic.rs b/crates/lance-csr/src/semantic.rs index b8cf649..f9909a0 100644 --- a/crates/lance-csr/src/semantic.rs +++ b/crates/lance-csr/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,9 @@ 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, + call_body: bool, } /// Information about a variable in the query @@ -70,15 +77,38 @@ impl SemanticAnalyzer { variables: HashMap::new(), retired: Vec::new(), current_scope: ScopeType::Match, + catalog: None, + expression_types: HashMap::new(), + call_body: false, } } + /// 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, 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(); @@ -92,6 +122,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)); @@ -125,6 +160,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)); @@ -155,6 +195,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) @@ -186,12 +227,48 @@ 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); // 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(); + 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( + analyzed + .errors + .into_iter() + .map(|e| format!("UNION branch: {e}")), + ); + warnings.extend(analyzed.warnings); + } + Ok(SemanticResult { ast: analyzed_query, variables: self.variables.clone(), @@ -200,6 +277,61 @@ 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(); + inner.call_body = true; + for name in query.call_union_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 { @@ -214,6 +346,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 +369,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(()) } @@ -263,7 +397,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; @@ -308,6 +446,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 +507,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 +530,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)?; @@ -425,7 +589,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 { @@ -438,7 +602,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" | "keys" | "labels" | "relationships" | "startnode" | "endnode" => { if args.len() != 1 { return Err(GraphError::PlanError { message: format!( @@ -450,10 +616,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 { @@ -468,6 +635,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 { @@ -506,7 +685,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, 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!()), @@ -514,6 +693,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)?; @@ -623,26 +830,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 +901,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,8 +910,46 @@ impl SemanticAnalyzer { self.analyze_value_expression(map) })(); self.variables = before; + 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, self.expression_type(initial)), + (variable, self.expression_type(list).element()), + ], + expression, + false, + )?; + } + ValueExpression::ListPredicate { + variable, + list, + predicate, + .. + } => { + self.analyze_value_expression(list)?; + self.analyze_lambda( + &[(variable, self.expression_type(list).element())], + predicate, + true, + )?; + } ValueExpression::ListComprehension { variable, list, @@ -724,7 +959,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,12 +974,45 @@ impl SemanticAnalyzer { Ok(()) })(); self.variables = before; + self.expression_types = before_types; result?; } } Ok(()) } + fn analyze_lambda( + &mut self, + names: &[(&String, ExpressionType)], + expression: &ValueExpression, + allow_patterns: bool, + ) -> Result<()> { + let before = self.variables.clone(); + 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); + 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; + self.expression_types = before_types; + result + } + fn register_projection_alias(&mut self, alias: &str) { // Use case-insensitive lookup and store normalized key if self.variables.contains_key_ci(alias) { @@ -774,7 +1046,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 +1063,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 +1077,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 +1098,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 +1147,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) { @@ -978,6 +1257,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 { @@ -987,6 +1278,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 { @@ -1020,6 +1312,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 { @@ -1056,6 +1349,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![], @@ -1096,6 +1390,7 @@ mod tests { end_node: end, }], shortest: None, + variable: None, }; let query = CypherQuery { @@ -1109,6 +1404,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![], @@ -1144,6 +1440,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![], @@ -1177,6 +1474,7 @@ mod tests { end_node: end, }], shortest: None, + variable: None, }; let query = CypherQuery { @@ -1190,6 +1488,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![], @@ -1220,6 +1519,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 { @@ -1263,6 +1563,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 { @@ -1301,6 +1602,7 @@ mod tests { end_node: end, }], shortest: None, + variable: None, }; let query = CypherQuery { @@ -1312,6 +1614,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 { @@ -1364,6 +1667,7 @@ mod tests { }, ], shortest: None, + variable: None, }; // Custom config that knows both relationship types to avoid warnings muddying the assertion @@ -1383,6 +1687,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 { @@ -1427,6 +1732,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 { @@ -1594,7 +1900,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(); @@ -1815,7 +2121,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-csr/src/semantic/types.rs b/crates/lance-csr/src/semantic/types.rs new file mode 100644 index 0000000..5718e4d --- /dev/null +++ b/crates/lance-csr/src/semantic/types.rs @@ -0,0 +1,410 @@ +// 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) | ValueExpression::LocalVariable(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) + } + "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), + "datetime" => Scalar(DataType::Timestamp( + TimeUnit::Millisecond, + Some("UTC".into()), + )), + "collect" => List(Box::new(first())), + "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( + 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(_) | ValueExpression::ListPredicate { .. } => { + 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::Reduce { .. } + | 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) { + // 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, + _ => 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_csr_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) 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 scalar_return_types_follow_with_and_unwind_scopes() { + for query in [ + "RETURN left('12', 1) + 1", + "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}"); + } + } + + #[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/crates/lance-csr/src/shortest_path.rs b/crates/lance-csr/src/shortest_path.rs index 9c37dec..f730638 100644 --- a/crates/lance-csr/src/shortest_path.rs +++ b/crates/lance-csr/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,18 @@ 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 previous_work = work; + let mut position = crate::csr_expand::NeighborPosition::default(); + for candidate in adjacency.candidates_from(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 +442,20 @@ fn expand_level( next.push(n); } } - }); + } + if work == previous_work { + work += 1; + if work.is_multiple_of(1024) { + crate::csr_expand::cooperative_yield().await; + } + } } 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 +476,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 +485,8 @@ fn search_bidirectional( reverse_level, &mut reverse_depth, &mut reverse_parents, - ); + ) + .await; } let mut pending: Vec<_> = frontier .iter() @@ -480,7 +498,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 +522,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 +554,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 +609,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 +657,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 +681,57 @@ 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 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, // including parallel edges on both sides of the cut, yield six paths. let (out, inc) = adjacencies(&[ @@ -700,11 +746,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 +760,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 +779,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 +788,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 +814,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 +823,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/crates/lance-csr/tests/test_datafusion_pipeline.rs b/crates/lance-csr/tests/test_datafusion_pipeline.rs index 554b1c9..56570d5 100644 --- a/crates/lance-csr/tests/test_datafusion_pipeline.rs +++ b/crates/lance-csr/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}" diff --git a/docs/api-reference.md b/docs/api-reference.md index 8ffd132..905be47 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -64,13 +64,26 @@ 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` | --- +`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, 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. + ## `GraphConfig` An immutable mapping from tables to a property graph. Create one with `GraphConfig.builder()`. @@ -115,27 +128,36 @@ 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 -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, 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. | Parameter | Description | | --- | --- | | `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 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. +be opened. For `target_partitions`, booleans and non-integers other than `None` raise `TypeError`; +values outside 1 to 1024 raise `ValueError`. -`#!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`, `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`. + 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()`. @@ -144,8 +166,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, and concurrent calls run one at a time. See - [Lance datasets](lance.md#refresh). + All-or-nothing; running queries keep their snapshot, and concurrent calls run one at a time. 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. @@ -165,6 +187,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` @@ -180,7 +208,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` @@ -190,7 +218,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. @@ -199,7 +228,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/csr.md b/docs/csr.md index 0d3ae5a..4bcfb21 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 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. @@ -121,7 +123,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,12 +134,37 @@ 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 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 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 +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 @@ -157,6 +185,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/docs/cypher.md b/docs/cypher.md index 0645739..7aed051 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 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" @@ -72,7 +73,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 +84,23 @@ 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'}] +``` + +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 +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})`. @@ -111,6 +125,24 @@ 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. 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)` Returns the plan as text: the graph logical plan, DataFusion's logical plan and the physical plan. @@ -168,7 +200,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, including named path values and the relationship-path operations +used by IC14. The tables below specify its supported surface. ### Clauses @@ -179,12 +212,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 integer literals, not parameters. | +| `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`), 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 @@ -192,6 +227,58 @@ 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. 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 +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: + +```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`. + +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 +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 @@ -214,8 +301,14 @@ 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. + 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. @@ -224,7 +317,58 @@ 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. 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 +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, 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. +This also applies inside `CALL` and `UNION`, and when equal-label endpoints pass through another label. ### Expressions and predicates @@ -251,16 +395,112 @@ 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) | +| Properties | `keys(map)`, `keys(node)`, `keys(relationship)`, `labels(node)` | +| 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)` 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) | -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. 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 +`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 +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}] +``` + +`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. + +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. + +### 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 @@ -285,15 +525,22 @@ 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 | 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. 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. | + +Known comparison and arithmetic types are checked against the current dataset schemas before planning. +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/docs/execution.md b/docs/execution.md index 7039694..c4df547 100644 --- a/docs/execution.md +++ b/docs/execution.md @@ -23,15 +23,25 @@ Source paths below are in `crates/lance-csr/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/docs/lance.md b/docs/lance.md index d18a1cd..7989c7c 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). @@ -79,6 +82,41 @@ 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, maps and parameterized `SKIP`/`LIMIT` counts use ordinary planning because their shape or value can +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. + +## Query parallelism + +`CypherEngine(..., target_partitions=None)` keeps DataFusion's default query partitioning. If you have +measured your workload, you can choose a target from 1 to 1024 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/docs/rust.md b/docs/rust.md index eeede83..e279e9a 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 @@ -105,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 @@ -118,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 | diff --git a/docs/sql.md b/docs/sql.md index 57f60c3..fd5bf8d 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,22 @@ 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. + +Constant-folded typed nulls use DataFusion `arrow_cast` in the Default dialect so their Arrow +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 +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/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/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/DEVELOPMENT.md b/python/DEVELOPMENT.md index fa10669..232458d 100644 --- a/python/DEVELOPMENT.md +++ b/python/DEVELOPMENT.md @@ -52,6 +52,25 @@ 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_python_home=$(uv run --no-sync python -c 'import sys; print(sys.base_prefix)') +lance_test_artifacts=$(mktemp) +cargo rustc --manifest-path ../crates/lance-csr-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") +PYTHONHOME="$lance_python_home" "$lance_test_binary" +rm "$lance_test_artifacts" +``` + +`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. + ## Formatting and linting ```shell 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/lance_csr/_internal/graph.pyi b/python/python/lance_csr/_internal/graph.pyi index 1c6b13a..0c921a6 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 @@ -103,13 +115,26 @@ 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, 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/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_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 diff --git a/python/python/tests/test_call_subqueries.py b/python/python/tests/test_call_subqueries.py new file mode 100644 index 0000000..87668e5 --- /dev/null +++ b/python/python/tests/test_call_subqueries.py @@ -0,0 +1,475 @@ +# 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 + + +@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 + + +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" + ) + + +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]}, + ] + + +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]}, + ] 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() 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}) 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..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", @@ -119,11 +124,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( 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}) diff --git a/python/python/tests/test_engine_paths.py b/python/python/tests/test_engine_paths.py index 922a369..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 = [ @@ -137,3 +139,55 @@ 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 + + +@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"}] diff --git a/python/python/tests/test_engine_types.py b/python/python/tests/test_engine_types.py new file mode 100644 index 0000000..f2efc5e --- /dev/null +++ b/python/python/tests/test_engine_types.py @@ -0,0 +1,143 @@ +# 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) 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"), + ), + ], +) +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"}, + ] + + +@pytest.mark.parametrize( + "query", + [ + "RETURN left('12', 1) + 1", + "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", + ], +) +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"), + [ + ("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 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) 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}] diff --git a/python/python/tests/test_math_functions.py b/python/python/tests/test_math_functions.py new file mode 100644 index 0000000..49d90a2 --- /dev/null +++ b/python/python/tests/test_math_functions.py @@ -0,0 +1,145 @@ +# 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") + + +@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( + 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]) 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) diff --git a/python/python/tests/test_path_values.py b/python/python/tests/test_path_values.py new file mode 100644 index 0000000..2d5c329 --- /dev/null +++ b/python/python/tests/test_path_values.py @@ -0,0 +1,649 @@ +# 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] + + +@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 + + +@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) + + +@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") + + +@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 [] + ) + + +@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 diff --git a/python/python/tests/test_plan_cache.py b/python/python/tests/test_plan_cache.py new file mode 100644 index 0000000..700dc25 --- /dev/null +++ b/python/python/tests/test_plan_cache.py @@ -0,0 +1,101 @@ +# 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_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 == ["expr"] + assert second.column_names == ["expr"] + 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": True}) + + +@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 + ) diff --git a/python/python/tests/test_scalar_functions.py b/python/python/tests/test_scalar_functions.py index 83eb1b5..b52b308 100644 --- a/python/python/tests/test_scalar_functions.py +++ b/python/python/tests/test_scalar_functions.py @@ -1,22 +1,221 @@ # 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.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() + 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]}, + ] diff --git a/python/python/tests/test_target_partitions.py b/python/python/tests/test_target_partitions.py new file mode 100644 index 0000000..2bc0f1e --- /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, 1025, 10**9, -(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/test_timeout.py b/python/python/tests/test_timeout.py new file mode 100644 index 0000000..e931ccf --- /dev/null +++ b/python/python/tests/test_timeout.py @@ -0,0 +1,103 @@ +# 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_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"): + 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}] + + +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}] diff --git a/python/python/tests/test_to_sql.py b/python/python/tests/test_to_sql.py index 8806d8f..f0ac77e 100644 --- a/python/python/tests/test_to_sql.py +++ b/python/python/tests/test_to_sql.py @@ -5,7 +5,125 @@ import pyarrow as pa import pytest -from lance_csr import CypherQuery, GraphConfig +from lance_csr import CypherQuery, GraphConfig, SqlDialect, 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)", + "round(null)", + "abs(null)", + "sqrt(null)", + "floor(null)", + "ceil(null)", + "toUpper(null)", + "null", + "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() + + +@pytest.mark.parametrize( + "dialect", + [SqlDialect.Spark, SqlDialect.PostgreSql, SqlDialect.MySql, SqlDialect.Sqlite], +) +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 + 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): + 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 diff --git a/python/python/tests/test_union.py b/python/python/tests/test_union.py new file mode 100644 index 0000000..ae8a700 --- /dev/null +++ b/python/python/tests/test_union.py @@ -0,0 +1,117 @@ +# 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]), + ( + "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", + "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): + 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"} + + +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, + ] diff --git a/python/python/tests/typing/public_api.py b/python/python/tests/typing/public_api.py index 1edc328..27f2965 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 @@ -53,12 +66,19 @@ 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) + 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]) @@ -89,7 +109,9 @@ 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] + one.with_parameter("ids", {1, 2}) # pyright: ignore[reportArgumentType] namespace.base_uri = "elsewhere" # pyright: ignore[reportAttributeAccessIssue]