diff --git a/crates/lance-csr-catalog/Cargo.toml b/crates/lance-csr-catalog/Cargo.toml index bc8b89e..a05507f 100644 --- a/crates/lance-csr-catalog/Cargo.toml +++ b/crates/lance-csr-catalog/Cargo.toml @@ -7,7 +7,7 @@ authors = ["The lance-csr Authors"] repository = "https://github.com/hyperdimensionalcomputingai/lance-csr" homepage = "https://lance-csr.hdclabs.ai/" readme = "README.md" -description = "Catalog and namespace utilities for Lance graph" +description = "Table catalog and directory namespace for the lance-csr graph query engine" keywords = ["lance", "graph", "catalog", "namespace"] categories = ["database", "data-structures", "science"] diff --git a/crates/lance-csr-catalog/README.md b/crates/lance-csr-catalog/README.md index 2e7e68d..ae7d607 100644 --- a/crates/lance-csr-catalog/README.md +++ b/crates/lance-csr-catalog/README.md @@ -1,6 +1,34 @@ -# Lance Graph Catalog +# lance-csr-catalog -Catalog and namespace utilities shared by the Lance graph query engine. +The table lookups behind [lance-csr](https://crates.io/crates/lance-csr), a Cypher graph query engine for +[Lance](https://lancedb.github.io/lance/) datasets built on Apache DataFusion. lance-csr reads a graph from +ordinary tables: one per node label and one per relationship type. This crate answers the two questions that +raises: which DataFusion table holds a given label or type, and where a named Lance dataset lives. -> **Experimental.** The Rust API is not stable yet: breaking changes are possible over the next few -> weeks while the CSR work lands. Pin an exact version (`lance-csr-catalog = "=x.y.z"`) if you depend on it. +- `GraphSourceCatalog` is the trait the lance-csr planner asks for a node label's or relationship type's + DataFusion `TableSource`, so that it can build scans while planning. Lookups ignore case. + `InMemoryCatalog` is a map-backed implementation, and `SimpleTableSource` is a table source with a fixed + schema and no data, for planning and tests. +- `DirNamespace` implements Lance's `LanceNamespace` for a directory of datasets: the table `person` under + `s3://bucket/graph` is `s3://bucket/graph/person.lance`. It works with any URI Lance can open. + +lance-csr re-exports all four types, so most programs only depend on `lance-csr`. Depend on this crate directly to +implement your own catalog, for example over a metadata service, without pulling in the query engine. + +```rust +use std::sync::Arc; + +use arrow_schema::{DataType, Field, Schema}; +use lance_csr_catalog::{GraphSourceCatalog, InMemoryCatalog, SimpleTableSource}; + +let person = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); +let catalog = InMemoryCatalog::new() + .with_node_source("Person", Arc::new(SimpleTableSource::new(person))); +assert!(catalog.node_source("person").is_some()); +``` + +The [Rust guide](https://lance-csr.hdclabs.ai/rust/) covers the engine that uses these types. + +## License + +Apache-2.0. See `LICENSE` and `NOTICE`. diff --git a/crates/lance-csr-catalog/src/lib.rs b/crates/lance-csr-catalog/src/lib.rs index 3c51bb7..d543687 100644 --- a/crates/lance-csr-catalog/src/lib.rs +++ b/crates/lance-csr-catalog/src/lib.rs @@ -1,10 +1,34 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -//! Catalog and namespace utilities for Lance Graph. +//! The table lookups behind [lance-csr](https://docs.rs/lance-csr), a Cypher graph query +//! engine for Lance datasets built on Apache DataFusion. //! -//! **Experimental.** The Rust API is not stable yet: breaking changes are possible over the -//! next few weeks while the CSR work lands. Pin an exact version if you depend on it. +//! lance-csr reads a graph from ordinary tables: one per node label and one per relationship +//! type. This crate answers the two questions that raises: +//! +//! - [`GraphSourceCatalog`]: which DataFusion `TableSource` holds a node label or relationship +//! type, so the planner can build scans. Lookups ignore case. [`InMemoryCatalog`] is a +//! map-backed implementation, and [`SimpleTableSource`] is a table source with a fixed schema +//! and no data, for planning and tests. +//! - [`DirNamespace`]: where a named Lance dataset lives, as Lance's `LanceNamespace` for a +//! directory of datasets. The table `person` under `s3://bucket/graph` is +//! `s3://bucket/graph/person.lance`. +//! +//! lance-csr re-exports all four types, so most programs only depend on `lance-csr`. Depend on +//! this crate directly to implement your own catalog without pulling in the query engine. +//! +//! ``` +//! use std::sync::Arc; +//! +//! use arrow_schema::{DataType, Field, Schema}; +//! use lance_csr_catalog::{GraphSourceCatalog, InMemoryCatalog, SimpleTableSource}; +//! +//! let person = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); +//! let catalog = InMemoryCatalog::new() +//! .with_node_source("Person", Arc::new(SimpleTableSource::new(person))); +//! assert!(catalog.node_source("person").is_some()); +//! ``` pub mod namespace; pub mod source_catalog; diff --git a/crates/lance-csr-python/src/graph.rs b/crates/lance-csr-python/src/graph.rs index d57be91..45a1026 100644 --- a/crates/lance-csr-python/src/graph.rs +++ b/crates/lance-csr-python/src/graph.rs @@ -21,15 +21,13 @@ use arrow::compute::concat_batches; use arrow::ffi_stream::ArrowArrayStreamReader; use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::Schema; -use datafusion::datasource::MemTable; -use datafusion::execution::context::SessionContext; use futures::StreamExt; use lance_csr::{ ast::{DistanceMetric as RustDistanceMetric, GraphPattern, ReadingClause}, engine::{CypherEngine as RustCypherEngine, EngineOptions, 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, + SqlEngine as RustSqlEngine, SqlQuery as RustSqlQuery, VectorSearch as RustVectorSearch, }; use pyo3::{ exceptions::{PyNotImplementedError, PyRuntimeError, PyTypeError, PyValueError}, @@ -1009,7 +1007,7 @@ fn reader_with_deadline( #[pyclass(name = "CypherEngine", module = "lance_csr.graph")] pub struct CypherEngine { - inner: RustCypherEngine, + inner: Arc, } #[pymethods] @@ -1066,7 +1064,9 @@ impl CypherEngine { ), )? .map_err(graph_error_to_pyerr)?; - Ok(Self { inner }) + Ok(Self { + inner: Arc::new(inner), + }) } #[pyo3(signature = (query, params = None, *, timeout = None))] @@ -1251,54 +1251,87 @@ impl SqlQuery { #[pyclass(name = "SqlEngine", module = "lance_csr.graph")] pub struct SqlEngine { - context: Arc, + inner: RustSqlEngine, } #[pymethods] impl SqlEngine { #[new] - fn new(datasets: &Bound<'_, PyDict>) -> PyResult { - let arrow_datasets = python_datasets_to_batches(datasets)?; - - if arrow_datasets.is_empty() { - return Err(PyValueError::new_err("No input datasets provided")); + #[pyo3(signature = (datasets, config = None, use_csr = false))] + fn new( + py: Python, + datasets: &Bound<'_, PyDict>, + config: Option<&GraphConfig>, + use_csr: bool, + ) -> PyResult { + if use_csr && config.is_none() { + return Err(PyValueError::new_err("use_csr requires a config")); } - - let ctx = SessionContext::new(); - - for (name, batch) in &arrow_datasets { - let mem_table = Arc::new( - MemTable::try_new(batch.schema(), vec![vec![batch.clone()]]).map_err(|e| { - PyRuntimeError::new_err(format!( - "Failed to create MemTable for {}: {}", - name, e - )) - })?, - ); - - let normalized_name = name.to_lowercase(); - ctx.register_table(&normalized_name, mem_table) - .map_err(|e| { - PyRuntimeError::new_err(format!("Failed to register table {}: {}", name, e)) - })?; + let mut sources = HashMap::new(); + for (key, value) in datasets.iter() { + let name: String = key.extract()?; + // A Lance dataset opened through a namespace can't be reopened by URI, so + // without a graph it is read into memory, as SqlEngine always did. + let source = match python_value_to_source(py, &name, &value) { + Err(_) if config.is_none() && is_lance_dataset(&value)? => { + EngineSource::Arrow(python_dataset_to_batch(&value)?) + } + source => source?, + }; + sources.insert(name, source); } - - Ok(Self { - context: Arc::new(ctx), - }) + let config = config.map(|c| c.inner.clone()).unwrap_or_default(); + let options = EngineOptions { + use_csr, + ..Default::default() + }; + let inner = RT + .block_on(Some(py), RustSqlEngine::with_graph(config, sources, options))? + .map_err(graph_error_to_pyerr)?; + Ok(Self { inner }) } - fn execute(&self, py: Python, sql: &str) -> PyResult> { - let query = RustSqlQuery::new(sql); - let context = self.context.as_ref().clone(); + #[staticmethod] + fn from_engine(engine: &CypherEngine) -> Self { + Self { + inner: RustSqlEngine::from_engine(engine.inner.clone()), + } + } + #[pyo3(signature = (sql, params = None, *, timeout = None))] + fn execute( + &self, + py: Python, + sql: &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), query.execute_with_context(context))? + .block_on_until(Some(py), deadline, self.inner.execute(sql, params))? .map_err(graph_error_to_pyerr)?; - record_batch_to_python_table(py, &result_batch) } + #[pyo3(signature = (sql, params = None))] + fn explain( + &self, + py: Python, + sql: &str, + params: Option<&Bound<'_, PyDict>>, + ) -> PyResult { + let params = python_params(params)?; + RT.block_on(Some(py), self.inner.explain(sql, params))? + .map_err(graph_error_to_pyerr) + } + + fn checkout_latest(&self, py: Python) -> PyResult<()> { + RT.block_on(Some(py), self.inner.checkout_latest())? + .map_err(graph_error_to_pyerr) + } + fn __repr__(&self) -> String { "SqlEngine(...)".to_string() } diff --git a/crates/lance-csr/README.md b/crates/lance-csr/README.md index 167ddd8..633db59 100644 --- a/crates/lance-csr/README.md +++ b/crates/lance-csr/README.md @@ -5,9 +5,6 @@ with a CSR (compressed sparse row) adjacency index for fast multi-hop traversal. Arrow tables as property graphs. The [lance-csr Python package](https://pypi.org/project/lance-csr/) wraps this crate. -> **Experimental.** The Rust API is not stable yet: breaking changes are possible over the next few -> weeks while the CSR work lands. Pin an exact version (`lance-csr = "=x.y.z"`) if you depend on it. - ## Features - Cypher query parsing and AST construction @@ -15,6 +12,7 @@ crate. - Semantic validation with typed `GraphError` diagnostics - Pluggable execution strategies (DataFusion planner by default, Lance Native placeholder) - `CypherEngine`: register datasets once, optionally build the CSR, and serve many queries +- `SqlEngine`: SQL over the same datasets, with Cypher patterns inside a `GRAPH_TABLE` clause - Lance datasets read lazily and pinned to a version until `checkout_latest()` - Async query execution that returns Arrow `RecordBatch` results or a stream - JSON-serializable parameter binding for reusable query templates @@ -154,7 +152,8 @@ clauses, patterns, expressions and functions the engine runs. - `logical_plan` – Builders for graph logical plans. - `datafusion_planner` – DataFusion-based execution planning. - `config` – Graph configuration types and builders. -- `engine` – `CypherEngine`: cached snapshot, CSR adjacencies, `checkout_latest()`. +- `engine` – `CypherEngine` and `SqlEngine`: cached snapshot, CSR adjacencies, `checkout_latest()`. +- `graph_table` – The `GRAPH_TABLE` clause and the `graph_table` table function. - `query` – High level `CypherQuery` API and runtime. - `csr_index`, `csr_expand` – The CSR adjacency index and the traversal operators that use it. - `error` – `GraphError` and result helpers. diff --git a/crates/lance-csr/src/engine.rs b/crates/lance-csr/src/engine.rs index d30d9ec..1638f03 100644 --- a/crates/lance-csr/src/engine.rs +++ b/crates/lance-csr/src/engine.rs @@ -96,6 +96,22 @@ impl EngineSource { }) } + /// A DataFusion table over this source. Lance is scanned lazily, with projection and + /// filter pushdown; Arrow is sliced into one partition per core. + pub fn table_provider(&self, name: &str) -> Result> { + Ok(match self { + // No _rowid/_rowaddr columns, so Lance and Arrow inputs expose the same schema. + Self::Lance { dataset, .. } => { + Arc::new(LanceTableProvider::new(dataset.clone(), false, false)) + } + Self::Arrow(batch) => Arc::new( + MemTable::try_new(batch.schema(), memtable_partitions(batch)).map_err(|e| { + exec_error(format!("Failed to create MemTable for {}: {}", name, e)) + })?, + ), + }) + } + /// Pinned Lance version, or `None` for Arrow input. pub fn lance_version(&self) -> Option { match self { @@ -204,6 +220,8 @@ struct EngineState { sources: HashMap, /// Unoptimized plan templates for this snapshot; `checkout_latest()` starts an empty cache. plans: Mutex>, + /// SQL plan templates for this snapshot, under the same rules. + sql_plans: Mutex>, catalog: Arc, context: SessionContext, /// CSR adjacencies by relationship type; empty unless the engine uses CSR. @@ -224,6 +242,13 @@ struct CachedQuery { query: CypherQuery, } +/// An unoptimized SQL plan whose `$name` parameters are still placeholders. +struct CachedSql { + text: String, + parameter_types: Vec<(String, u8)>, + plan: datafusion::logical_expr::LogicalPlan, +} + // Containers change expression shape; keep them on the uncached path. fn parameter_types(parameters: &HashMap) -> Option> { let mut types = parameters @@ -244,6 +269,23 @@ fn parameter_types(parameters: &HashMap) -> Option Result { + let query = CypherQuery::new(text)? + .with_config(config.clone()) + .with_fan_out(state.fan_out()); + Ok(if use_csr { + query.with_adjacency(state.adjacency_catalog()) + } else { + query + }) +} + /// Distinct counts for one dataset's node-id and edge-endpoint columns. #[derive(Clone)] struct KeyStatsEntry { @@ -432,6 +474,36 @@ fn exec_error(message: String) -> GraphError { } } +/// Plan `sql` in `ctx` after rewriting its `GRAPH_TABLE` clauses; DDL, DML and statements +/// such as `SET` are rejected. +async fn plan_read_only_sql( + ctx: &SessionContext, + sql: &str, +) -> Result { + use datafusion::execution::context::SQLOptions; + + let sql = crate::graph_table::rewrite_graph_table(sql)?; + let plan = ctx + .state() + .create_logical_plan(&sql) + .await + .map_err(|e| plan_error(format!("SQL planning error: {e}")))?; + SQLOptions::new() + .with_allow_ddl(false) + .with_allow_dml(false) + .with_allow_statements(false) + .verify_plan(&plan) + .map_err(|e| plan_error(e.to_string()))?; + Ok(plan) +} + +fn plan_error(message: String) -> GraphError { + GraphError::PlanError { + message, + location: snafu::Location::new(file!(), line!(), column!()), + } +} + fn config_error(message: String) -> GraphError { GraphError::ConfigError { message, @@ -530,17 +602,7 @@ async fn build_engine_state( let mut providers: HashMap> = HashMap::new(); for (name, source) in &sources { - let provider: Arc = match source { - // No _rowid/_rowaddr columns, so Lance and Arrow inputs expose the same schema. - EngineSource::Lance { dataset, .. } => { - Arc::new(LanceTableProvider::new(dataset.clone(), false, false)) - } - EngineSource::Arrow(batch) => Arc::new( - MemTable::try_new(batch.schema(), memtable_partitions(batch)).map_err(|e| { - exec_error(format!("Failed to create MemTable for {}: {}", name, e)) - })?, - ), - }; + let provider = source.table_provider(name)?; schemas.insert(name.clone(), provider.schema()); // Registered unwrapped first: the CSR and key-statistics builds below scan it. ctx.register_table(name.to_lowercase(), provider.clone()) @@ -589,6 +651,7 @@ async fn build_engine_state( Ok(EngineState { sources, plans: Mutex::default(), + sql_plans: Mutex::default(), catalog: Arc::new(catalog), context: ctx, csr, @@ -1039,15 +1102,8 @@ impl CypherEngine { return Ok(query); } } - let query = CypherQuery::new(text)? - .with_parameters(parameters) - .with_config(self.config.clone()) - .with_fan_out(state.fan_out()); - let mut query = if self.options.use_csr { - query.with_adjacency(state.adjacency_catalog()) - } else { - query - }; + let mut query = configured_query(&self.config, self.options.use_csr, state, text)? + .with_parameters(parameters); 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. @@ -1112,6 +1168,131 @@ impl CypherEngine { .await } + /// Run `sql` on this snapshot; see [`SqlEngine::execute`]. + async fn sql(&self, sql: &str, params: HashMap) -> Result { + crate::query::collect_batch(self.sql_dataframe(sql, params, true).await?).await + } + + /// See [`SqlEngine::explain`]. + async fn explain_sql(&self, sql: &str, params: HashMap) -> Result { + let df = self.sql_dataframe(sql, params, false).await?; + let logical = df + .clone() + .into_optimized_plan() + .map_err(|e| plan_error(format!("Failed to optimize SQL plan: {e}")))?; + let physical = df + .create_physical_plan() + .await + .map_err(|e| plan_error(format!("Failed to create physical plan: {e}")))?; + Ok(format!( + "== Logical Plan ==\n{}\n\n== Physical Plan ==\n{}", + logical.display_indent(), + datafusion::physical_plan::displayable(physical.as_ref()).indent(true), + )) + } + + /// Plan `sql`, or reuse this snapshot's template for it, then bind `params`. + /// + /// A template is the unoptimized plan with every `$name` still a placeholder: SQL's own, + /// and those of `GRAPH_TABLE` patterns, which are planned as Cypher templates. It is + /// keyed like the Cypher cache, and skipped when any pattern can't be a template. + async fn sql_dataframe( + &self, + sql: &str, + params: HashMap, + use_cache: bool, + ) -> Result { + use std::sync::atomic::{AtomicBool, Ordering}; + + let state = self.snapshot(); + let params: HashMap = params + .into_iter() + .map(|(name, value)| (name.to_lowercase(), value)) + .collect(); + let key = use_cache.then(|| parameter_types(¶ms)).flatten(); + // A copy, so the function (which holds the snapshot) isn't registered in it. + let ctx = SessionContext::new_with_state(state.context.state()); + let cached = key.as_ref().and_then(|key| { + let mut plans = state.sql_plans.lock().expect("plan cache lock poisoned"); + let index = plans + .iter() + .position(|entry| entry.text == sql && entry.parameter_types == *key)?; + let entry = plans.remove(index).expect("index from position"); + let plan = entry.plan.clone(); + plans.push_back(entry); + Some(plan) + }); + let plan = match cached { + Some(plan) => plan, + None => { + let templated = Arc::new(AtomicBool::new(key.is_some())); + self.register_graph_table(&ctx, &state, ¶ms, templated.clone()); + let plan = plan_read_only_sql(&ctx, sql).await?; + if let Some(parameter_types) = key.filter(|_| templated.load(Ordering::Relaxed)) { + let mut plans = state.sql_plans.lock().expect("plan cache lock poisoned"); + if plans.len() == PLAN_CACHE_ENTRIES { + plans.pop_front(); + } + plans.push_back(CachedSql { + text: sql.to_string(), + parameter_types, + plan: plan.clone(), + }); + } + plan + } + }; + let plan = match crate::query::scalar_values(¶ms) { + Some(values) => plan + .with_param_values(values) + .map_err(|e| plan_error(format!("Failed to bind query parameters: {e}")))?, + _ => plan, + }; + Ok(crate::query::planned_dataframe(ctx, plan)) + } + + /// Register `graph_table` in `ctx`. While `templated` is set, patterns are planned as + /// templates; one that can't be clears it and is planned with `params` as values. + fn register_graph_table( + &self, + ctx: &SessionContext, + state: &Arc, + params: &HashMap, + templated: Arc, + ) { + use crate::graph_table::{GraphTableFunction, GRAPH_TABLE_FUNCTION}; + use std::sync::atomic::Ordering; + + let (config, use_csr) = (self.config.clone(), self.options.use_csr); + let (state, params) = (state.clone(), params.clone()); + ctx.register_udtf( + GRAPH_TABLE_FUNCTION, + Arc::new(GraphTableFunction::new(move |text| { + if config.iter_node_mappings().len() == 0 + && config.iter_relationship_mappings().len() == 0 + { + return Err(config_error( + "GRAPH_TABLE needs a graph: create the SqlEngine with a GraphConfig" + .to_string(), + )); + } + let mut query = configured_query(&config, use_csr, &state, text)? + .with_parameters(params.clone()); + if templated.load(Ordering::Relaxed) + && query + .prepare_cached_plan(state.catalog.clone()) + .unwrap_or(false) + { + if let Some(plan) = query.template_plan() { + return Ok(plan); + } + } + templated.store(false, Ordering::Relaxed); + query.datafusion_plan(state.catalog.clone()) + })), + ); + } + /// Execute a Cypher query, then rerank its rows by vector similarity. pub async fn execute_with_vector_rerank( &self, @@ -1266,6 +1447,80 @@ impl CypherEngine { } } +/// SQL over a set of datasets, with graph patterns when it has a graph. +/// +/// Datasets are tables under their lower-cased names; Lance datasets are scanned in place +/// at their pinned version, with column and filter pushdown. With a [`GraphConfig`], a +/// `GRAPH_TABLE (MATCH ... COLUMNS (...))` clause (see [`crate::graph_table`]) embeds a +/// Cypher pattern, whose hops use the CSR when the engine has one. +/// +/// It runs on a [`CypherEngine`]'s snapshot. [`from_engine`](Self::from_engine) shares an +/// existing engine's, so the datasets, CSR and statistics are built once for both, and +/// [`checkout_latest`](Self::checkout_latest) on either moves both. +#[derive(Clone)] +pub struct SqlEngine { + engine: Arc, +} + +impl SqlEngine { + /// An engine over `sources` with no graph: plain SQL only. + /// + /// # Errors + /// As [`CypherEngine::new`]. + pub async fn new(sources: HashMap) -> Result { + Self::with_graph(GraphConfig::default(), sources, EngineOptions::default()).await + } + + /// An engine over `sources` with `config` as its graph, built as + /// [`CypherEngine::with_options`] builds one. + /// + /// # Errors + /// As [`CypherEngine::with_options`]. + pub async fn with_graph( + config: GraphConfig, + sources: HashMap, + options: EngineOptions, + ) -> Result { + let engine = CypherEngine::with_options(config, sources, options).await?; + Ok(Self::from_engine(Arc::new(engine))) + } + + /// SQL over `engine`'s snapshot, sharing its datasets, CSR and statistics. + pub fn from_engine(engine: Arc) -> Self { + Self { engine } + } + + /// Run a SQL query and return the whole result as one batch. + /// + /// `params` gives values for `$name` parameters, in the SQL and in `GRAPH_TABLE` + /// patterns alike. Plans are cached per snapshot like Cypher templates: keyed by the + /// query text and the parameters' names and scalar types, and bound on every call. + /// Container parameters, and patterns Cypher can't plan as templates, are planned + /// on every call. + /// + /// # Errors + /// [`GraphError::PlanError`] for invalid SQL, a missing parameter, or a query that + /// would write (DDL, DML, `SET`); [`GraphError::ConfigError`] for `GRAPH_TABLE` + /// without a graph; [`GraphError::ExecutionError`] if execution fails. + pub async fn execute( + &self, + sql: &str, + params: HashMap, + ) -> Result { + self.engine.sql(sql, params).await + } + + /// The optimized logical and physical plans of a query, always planned afresh. + pub async fn explain(&self, sql: &str, params: HashMap) -> Result { + self.engine.explain_sql(sql, params).await + } + + /// Move to the latest version of every Lance dataset; see [`CypherEngine::checkout_latest`]. + pub async fn checkout_latest(&self) -> Result<()> { + self.engine.checkout_latest().await + } +} + #[cfg(test)] mod tests { use super::*; @@ -1670,4 +1925,177 @@ mod tests { assert_eq!(engine.delete_stored_csr().await.unwrap(), 0); assert_eq!(two_hops(&uri, true).await.1.update, "build"); } + + /// An engine over Lance `Person(id)` and `KNOWS(src, dst)` datasets in `dir`. + async fn lance_graph_engine(dir: &std::path::Path, use_csr: bool) -> CypherEngine { + let (people, knows) = ( + dir.join("person").display().to_string(), + dir.join("knows").display().to_string(), + ); + write_people(&people, (0..10).collect(), WriteMode::Create).await; + let edges: Vec<(i64, i64)> = (0..10) + .flat_map(|i| [(i, (i + 1) % 10), (i, (i + 3) % 10)]) + .collect(); + write_knows(&knows, &edges, WriteMode::Create).await; + let config = GraphConfig::builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + .unwrap(); + let sources = HashMap::from([ + ( + "Person".to_string(), + EngineSource::open_lance(&people, None, HashMap::new()) + .await + .unwrap(), + ), + ( + "KNOWS".to_string(), + EngineSource::open_lance(&knows, None, HashMap::new()) + .await + .unwrap(), + ), + ]); + CypherEngine::new(config, sources, use_csr).await.unwrap() + } + + fn int_column(batch: &RecordBatch, name: &str) -> Vec { + let column = batch.column_by_name(name).unwrap(); + let column = arrow::compute::cast(column, &DataType::Int64).unwrap(); + column + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + } + + fn params(values: &[(&str, JsonValue)]) -> HashMap { + values + .iter() + .map(|(name, value)| (name.to_string(), value.clone())) + .collect() + } + + /// GRAPH_TABLE hops run on the CSR inside SQL, and match both Cypher and the join path. + #[tokio::test] + async fn sql_graph_table_uses_csr_and_matches_cypher() { + let dir = tempfile::tempdir().unwrap(); + let csr = Arc::new(lance_graph_engine(dir.path(), true).await); + let joins = Arc::new(lance_graph_engine(&dir.path().join("joins"), false).await); + let sql = "SELECT g.b, count(*) AS n FROM GRAPH_TABLE ( + MATCH (a:Person)-[:KNOWS]->(b:Person)-[:KNOWS]->(c:Person) WHERE a.id < 5 + COLUMNS (b.id AS b) + ) g JOIN person p ON p.id = g.b + GROUP BY g.b ORDER BY g.b"; + let cypher = "MATCH (a:Person)-[:KNOWS]->(b:Person)-[:KNOWS]->(c:Person) WHERE a.id < 5 + RETURN b.id AS b, count(*) AS n ORDER BY b"; + let expected = csr.execute(cypher, HashMap::new()).await.unwrap(); + for engine in [&csr, &joins] { + let result = SqlEngine::from_engine(engine.clone()) + .execute(sql, HashMap::new()) + .await + .unwrap(); + assert_eq!(int_column(&result, "b"), int_column(&expected, "b")); + assert_eq!(int_column(&result, "n"), int_column(&expected, "n")); + } + let explain = |engine: &Arc| { + let engine = SqlEngine::from_engine(engine.clone()); + async move { engine.explain(sql, HashMap::new()).await.unwrap() } + }; + assert!(explain(&csr).await.contains("CsrExpandExec")); + assert!(!explain(&joins).await.contains("CsrExpand")); + + // The direct function call is the same plan. + let direct = + "SELECT * FROM graph_table('MATCH (a:Person)-[:KNOWS]->(b:Person) RETURN b.id AS b')"; + let sql_engine = SqlEngine::from_engine(csr); + assert_eq!( + sql_engine + .execute(direct, HashMap::new()) + .await + .unwrap() + .num_rows(), + 20 + ); + } + + /// Plain SQL scans Lance datasets with pushdown, only reads, and needs no graph. + #[tokio::test] + async fn sql_scans_lance_lazily_and_rejects_writes() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("knows").display().to_string(); + write_knows(&path, &[(3, 4), (3, 6), (1, 2)], WriteMode::Create).await; + let knows = EngineSource::open_lance(&path, None, HashMap::new()) + .await + .unwrap(); + let engine = SqlEngine::new(HashMap::from([("KNOWS".to_string(), knows)])) + .await + .unwrap(); + let sql = "SELECT dst FROM knows WHERE src = 3 ORDER BY dst"; + let result = engine.execute(sql, HashMap::new()).await.unwrap(); + assert_eq!(int_column(&result, "dst"), vec![4, 6]); + let plan = engine.explain(sql, HashMap::new()).await.unwrap(); + assert!(plan.contains("full_filter=src = Int64(3)"), "{plan}"); + for query in [ + "CREATE TABLE t AS SELECT 1", + "COPY knows TO '/tmp/knows.parquet'", + "INSERT INTO knows VALUES (9, 9)", + "SET datafusion.execution.batch_size = 1", + "SELECT * FROM GRAPH_TABLE (MATCH (a:Person) COLUMNS (a.id))", + ] { + assert!( + engine.execute(query, HashMap::new()).await.is_err(), + "{query}" + ); + } + } + + /// One template per query text and parameter types, bound on every call, in SQL + /// and in GRAPH_TABLE alike; container values and checkout_latest skip or reset it. + #[tokio::test] + async fn sql_parameters_bind_into_cached_templates() { + let dir = tempfile::tempdir().unwrap(); + let engine = Arc::new(lance_graph_engine(dir.path(), true).await); + let sql_engine = SqlEngine::from_engine(engine.clone()); + let cached = || engine.snapshot().sql_plans.lock().unwrap().len(); + let sql = "SELECT g.b FROM GRAPH_TABLE ( + MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE a.id = $start + COLUMNS (b.id AS b) + ) g WHERE g.b <> $skip ORDER BY g.b"; + // Person i knows i + 1 and i + 3 (mod 10). + for (start, skip, expected) in [(1, 0, vec![2, 4]), (2, 3, vec![5]), (8, 9, vec![1])] { + let values = params(&[("start", start.into()), ("SKIP", skip.into())]); + let result = sql_engine.execute(sql, values).await.unwrap(); + assert_eq!(int_column(&result, "b"), expected); + assert_eq!(cached(), 1); + } + let plan = sql_engine + .explain(sql, params(&[("start", 1.into()), ("skip", 0.into())])) + .await + .unwrap(); + assert!( + plan.contains("CsrExpandExec") && plan.contains("id = Int64(1)"), + "{plan}" + ); + + // Other parameter types are another template; a container is planned each time. + let values = params(&[("start", 1.into()), ("skip", JsonValue::Null)]); + assert!(sql_engine.execute(sql, values).await.unwrap().num_rows() == 0); + assert_eq!(cached(), 2); + let listed = + "SELECT * FROM GRAPH_TABLE (MATCH (a:Person) WHERE a.id IN $ids COLUMNS (a.id AS id)) + ORDER BY id"; + let values = params(&[("ids", serde_json::json!([3, 1]))]); + let result = sql_engine.execute(listed, values).await.unwrap(); + assert_eq!(int_column(&result, "id"), vec![1, 3]); + assert_eq!(cached(), 2); + + assert!(sql_engine + .execute(sql, params(&[("start", 1.into())])) + .await + .is_err()); + sql_engine.checkout_latest().await.unwrap(); + assert_eq!(cached(), 0); + } } diff --git a/crates/lance-csr/src/graph_table.rs b/crates/lance-csr/src/graph_table.rs new file mode 100644 index 0000000..8205165 --- /dev/null +++ b/crates/lance-csr/src/graph_table.rs @@ -0,0 +1,293 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Graph patterns inside SQL, after DuckPGQ's `GRAPH_TABLE` clause. +//! +//! ```sql +//! SELECT name, count(*) FROM GRAPH_TABLE ( +//! MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE a.age > 30 +//! COLUMNS (b.name AS name) +//! ) GROUP BY name +//! ``` +//! +//! [`rewrite_graph_table`] turns each clause into a call to the `graph_table` table +//! function with the equivalent Cypher text (`MATCH ... WHERE ... RETURN `). +//! [`GraphTableFunction`] plans that text with the Cypher planner and returns the plan +//! as a view, which DataFusion inlines into the SQL plan. Hops covered by the CSR are +//! therefore `CsrExpand` nodes in the SQL plan, and the CSR optimizer rules apply to them. +//! +//! Only `GRAPH_TABLE`'s shape comes from SQL/PGQ. The pattern, `WHERE` and `COLUMNS` +//! bodies are Cypher, and there is no `CREATE PROPERTY GRAPH`: the graph is the +//! engine's `GraphConfig`. + +use std::sync::Arc; + +use datafusion::catalog::{TableFunctionArgs, TableFunctionImpl}; +use datafusion::datasource::{TableProvider, ViewTable}; +use datafusion::error::DataFusionError; +use datafusion::logical_expr::{Expr, LogicalPlan}; + +use crate::error::{GraphError, Result}; + +/// Name the `GRAPH_TABLE` clause is rewritten to; also callable directly with a Cypher string. +pub const GRAPH_TABLE_FUNCTION: &str = "graph_table"; + +/// Plans one Cypher query into an unoptimized DataFusion plan. +pub type CypherPlanner = dyn Fn(&str) -> Result + Send + Sync; + +/// `graph_table('')`: the Cypher query's result as a table. +pub struct GraphTableFunction { + plan: Box, +} + +impl GraphTableFunction { + pub fn new(plan: impl Fn(&str) -> Result + Send + Sync + 'static) -> Self { + Self { + plan: Box::new(plan), + } + } +} + +impl std::fmt::Debug for GraphTableFunction { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("GraphTableFunction") + } +} + +impl TableFunctionImpl for GraphTableFunction { + fn call_with_args( + &self, + args: TableFunctionArgs, + ) -> datafusion::error::Result> { + let text = match args.exprs() { + [Expr::Literal(value, _)] => value.try_as_str().flatten(), + _ => None, + } + .ok_or_else(|| { + DataFusionError::Plan(format!( + "{GRAPH_TABLE_FUNCTION}() takes one string literal: a Cypher query" + )) + })?; + let plan = (self.plan)(text).map_err(|e| DataFusionError::Plan(e.to_string()))?; + Ok(Arc::new(ViewTable::new(plan, Some(text.to_string())))) + } +} + +/// Replace every `GRAPH_TABLE ( [graph] MATCH ... COLUMNS (...) )` clause in `sql` +/// with `graph_table('MATCH ... RETURN ...')`. +/// +/// A direct call, `graph_table('...')`, is left as written. Quoted strings and +/// identifiers are skipped, so `GRAPH_TABLE` inside them is not a clause. +pub fn rewrite_graph_table(sql: &str) -> Result { + const KEYWORD: &str = "graph_table"; + let b = sql.as_bytes(); + let mut out = String::with_capacity(sql.len()); + let (mut copied, mut i) = (0, 0); + while i < b.len() { + if is_quote(b[i]) { + i = skip_quoted(b, i)?; + } else if is_word(b, i, KEYWORD) { + let open = skip_whitespace(b, i + KEYWORD.len()); + if b.get(open) != Some(&b'(') { + i += KEYWORD.len(); + continue; + } + let close = matching_paren(b, open)?; + let body = &sql[open + 1..close]; + if !body.trim_start().starts_with('\'') { + out.push_str(&sql[copied..i]); + out.push_str(GRAPH_TABLE_FUNCTION); + out.push_str("('"); + out.push_str(&to_cypher(body)?.replace('\'', "''")); + out.push_str("')"); + copied = close + 1; + } + i = close + 1; + } else { + i += 1; + } + } + out.push_str(&sql[copied..]); + Ok(out) +} + +/// `[graph] MATCH [WHERE ] COLUMNS ()` -> `MATCH ... RETURN `. +fn to_cypher(body: &str) -> Result { + let b = body.as_bytes(); + let match_at = *top_level_words(b, "match")? + .first() + .ok_or_else(|| syntax_error("GRAPH_TABLE needs a MATCH clause"))?; + // shortcut: a graph name before MATCH is accepted and ignored; an engine holds one graph. + let graph = body[..match_at].trim(); + if graph.split_whitespace().count() > 1 { + return Err(syntax_error(format!( + "unexpected '{graph}' before MATCH in GRAPH_TABLE" + ))); + } + let columns_at = top_level_words(b, "columns")? + .into_iter() + .rev() + .find(|&at| at > match_at) + .ok_or_else(|| syntax_error("GRAPH_TABLE needs a COLUMNS (...) list after MATCH"))?; + let open = skip_whitespace(b, columns_at + "columns".len()); + if b.get(open) != Some(&b'(') { + return Err(syntax_error( + "COLUMNS must be followed by a parenthesized list", + )); + } + let close = matching_paren(b, open)?; + if !body[close + 1..].trim().is_empty() { + return Err(syntax_error(format!( + "unexpected '{}' after COLUMNS (...) in GRAPH_TABLE", + body[close + 1..].trim() + ))); + } + let columns = body[open + 1..close].trim(); + if columns.is_empty() { + return Err(syntax_error("COLUMNS (...) is empty")); + } + Ok(format!( + "{} RETURN {columns}", + body[match_at..columns_at].trim() + )) +} + +/// Byte offsets of `word` outside quotes and brackets. +fn top_level_words(b: &[u8], word: &str) -> Result> { + let (mut found, mut depth, mut i) = (Vec::new(), 0usize, 0); + while i < b.len() { + match b[i] { + c if is_quote(c) => { + i = skip_quoted(b, i)?; + continue; + } + b'(' | b'[' | b'{' => depth += 1, + b')' | b']' | b'}' => depth = depth.saturating_sub(1), + _ if depth == 0 && is_word(b, i, word) => found.push(i), + _ => {} + } + i += 1; + } + Ok(found) +} + +/// Offset of the `)` closing the `(` at `open`. +fn matching_paren(b: &[u8], open: usize) -> Result { + let (mut depth, mut i) = (0usize, open); + while i < b.len() { + match b[i] { + c if is_quote(c) => { + i = skip_quoted(b, i)?; + continue; + } + b'(' => depth += 1, + b')' => { + depth -= 1; + if depth == 0 { + return Ok(i); + } + } + _ => {} + } + i += 1; + } + Err(syntax_error("unbalanced parentheses in GRAPH_TABLE")) +} + +fn is_quote(c: u8) -> bool { + matches!(c, b'\'' | b'"' | b'`') +} + +/// Offset just past the quoted string or identifier starting at `i`; a doubled quote is an escape. +// shortcut: Cypher's backslash escapes (`'it\'s'`) aren't recognized; use a doubled quote or "..." instead. +fn skip_quoted(b: &[u8], i: usize) -> Result { + let quote = b[i]; + let mut j = i + 1; + while j < b.len() { + if b[j] == quote { + if b.get(j + 1) == Some("e) { + j += 2; + continue; + } + return Ok(j + 1); + } + j += 1; + } + Err(syntax_error("unterminated quote")) +} + +fn skip_whitespace(b: &[u8], mut i: usize) -> usize { + while b.get(i).is_some_and(u8::is_ascii_whitespace) { + i += 1; + } + i +} + +/// Whether `word` starts at `i` as a whole identifier, ignoring case. +fn is_word(b: &[u8], i: usize, word: &str) -> bool { + let ident = |c: &u8| c.is_ascii_alphanumeric() || *c == b'_'; + b.get(i..i + word.len()) + .is_some_and(|s| s.eq_ignore_ascii_case(word.as_bytes())) + && (i == 0 || !ident(&b[i - 1])) + && !b.get(i + word.len()).is_some_and(ident) +} + +fn syntax_error(message: impl Into) -> GraphError { + GraphError::PlanError { + message: message.into(), + location: snafu::Location::new(file!(), line!(), column!()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rewrites_clause_to_cypher_call() { + let sql = "SELECT * FROM GRAPH_TABLE (g MATCH (a:Person {name: 'Al''s'})-[:KNOWS]->(b) \ + WHERE a.age > 30 COLUMNS (a.name AS x, count(b) AS n)) t WHERE t.n > 1"; + assert_eq!( + rewrite_graph_table(sql).unwrap(), + "SELECT * FROM graph_table('MATCH (a:Person {name: ''Al''''s''})-[:KNOWS]->(b) \ + WHERE a.age > 30 RETURN a.name AS x, count(b) AS n') t WHERE t.n > 1" + ); + } + + #[test] + fn leaves_direct_calls_quotes_and_names_alone() { + for sql in [ + "SELECT * FROM graph_table('MATCH (a) RETURN a.id')", + "SELECT 'GRAPH_TABLE (x)', \"graph_table\" FROM t", + "SELECT my_graph_table FROM t", + ] { + assert_eq!(rewrite_graph_table(sql).unwrap(), sql); + } + } + + #[test] + fn rewrites_several_clauses() { + let sql = "SELECT * FROM GRAPH_TABLE(MATCH (a) COLUMNS (a.id)) x \ + JOIN graph_table (MATCH (b) COLUMNS (b.id)) y ON true"; + assert_eq!( + rewrite_graph_table(sql).unwrap(), + "SELECT * FROM graph_table('MATCH (a) RETURN a.id') x \ + JOIN graph_table('MATCH (b) RETURN b.id') y ON true" + ); + } + + #[test] + fn rejects_malformed_clauses() { + for sql in [ + "SELECT * FROM GRAPH_TABLE (COLUMNS (a.id))", + "SELECT * FROM GRAPH_TABLE (MATCH (a))", + "SELECT * FROM GRAPH_TABLE (MATCH (a) COLUMNS ())", + "SELECT * FROM GRAPH_TABLE (MATCH (a) COLUMNS (a.id) LIMIT 1)", + "SELECT * FROM GRAPH_TABLE (two words MATCH (a) COLUMNS (a.id))", + "SELECT * FROM GRAPH_TABLE (MATCH (a) COLUMNS (a.id)", + "SELECT * FROM GRAPH_TABLE (MATCH (a {n: 'x}) COLUMNS (a.id))", + ] { + assert!(rewrite_graph_table(sql).is_err(), "{sql}"); + } + } +} diff --git a/crates/lance-csr/src/lib.rs b/crates/lance-csr/src/lib.rs index 7e86159..8a72708 100644 --- a/crates/lance-csr/src/lib.rs +++ b/crates/lance-csr/src/lib.rs @@ -1,14 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -//! Lance Graph Query Engine +//! A Cypher graph query engine for Lance datasets, built on Apache DataFusion. //! -//! This crate provides graph query capabilities for Lance datasets using Cypher syntax. -//! It interprets Lance datasets as property graphs and translates Cypher queries into -//! DataFusion SQL queries for execution. -//! -//! **Experimental.** The Rust API is not stable yet: breaking changes are possible over the -//! next few weeks while the CSR work lands. Pin an exact version if you depend on it. +//! This crate interprets Lance datasets or Arrow tables as property graphs and plans Cypher +//! queries as DataFusion logical plans, with graph operators of its own where joins fall short. //! //! # Features //! @@ -18,6 +14,7 @@ //! - Translation to optimized SQL via DataFusion //! - Support for nodes, relationships, and properties //! - A CSR (compressed sparse row) adjacency index for multi-hop traversal +//! - SQL over the same datasets, with graph patterns in a `GRAPH_TABLE` clause ([`SqlEngine`]) //! //! # Example //! @@ -57,6 +54,7 @@ mod csr_traversal; pub mod datafusion_planner; pub mod engine; pub mod error; +pub mod graph_table; pub mod key_statistics; pub mod lance_vector_search; pub mod logical_plan; @@ -78,7 +76,7 @@ pub const MAX_VARIABLE_LENGTH_HOPS: u32 = 20; pub use config::{GraphConfig, NodeMapping, RelationshipMapping}; pub use csr_index::{CsrIndex, CsrIndexBuilder}; -pub use engine::{CypherEngine, EngineSource}; +pub use engine::{CypherEngine, EngineSource, SqlEngine}; pub use error::{GraphError, Result}; pub use lance_csr_catalog::{DirNamespace, GraphSourceCatalog, InMemoryCatalog, SimpleTableSource}; pub use lance_vector_search::VectorSearch; diff --git a/crates/lance-csr/src/query.rs b/crates/lance-csr/src/query.rs index 662ec06..af56063 100644 --- a/crates/lance-csr/src/query.rs +++ b/crates/lance-csr/src/query.rs @@ -145,6 +145,78 @@ fn query_dataframe( datafusion::dataframe::DataFrame::new(state, plan) } +/// A frame for `plan` in `ctx`, under the caller's execution deadline, if any. +pub(crate) fn planned_dataframe( + ctx: datafusion::execution::context::SessionContext, + plan: datafusion::logical_expr::LogicalPlan, +) -> datafusion::dataframe::DataFrame { + query_dataframe(&deadline::context(ctx), plan) +} + +/// Run `df` and return its rows as one batch, with the physical plan's schema. +pub(crate) async fn collect_batch( + df: datafusion::dataframe::DataFrame, +) -> Result { + use arrow::compute::concat_batches; + + // 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 + // This preserves column structure even when there are no rows + return Ok(arrow::record_batch::RecordBatch::new_empty(result_schema)); + } + + // 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!()), + }) +} + +/// Scalar parameter values, or `None` when any is a container. +pub(crate) fn scalar_values( + parameters: &HashMap, +) -> Option> { + use datafusion::common::ScalarValue; + 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() +} + /// Execution strategy for Cypher queries #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum ExecutionStrategy { @@ -324,25 +396,15 @@ impl CypherQuery { } 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() + scalar_values(&self.parameters) + } + + /// The unoptimized template [`prepare_cached_plan`](Self::prepare_cached_plan) built, + /// with its parameters still placeholders. + pub(crate) fn template_plan(&self) -> Option { + self.cached_plan + .as_ref() + .map(|plan| plan.datafusion.clone()) } /// Get the required config, returning an error if not set @@ -756,43 +818,10 @@ impl CypherQuery { catalog: std::sync::Arc, ctx: datafusion::execution::context::SessionContext, ) -> Result { - use arrow::compute::concat_batches; - let df = self .dataframe_with_catalog_and_context(catalog, ctx) .await?; - - // 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 - // This preserves column structure even when there are no rows - return Ok(arrow::record_batch::RecordBatch::new_empty(result_schema)); - } - - // 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!()), - }) + collect_batch(df).await } /// Execute query with an explicit catalog and session context, returning a stream. @@ -832,8 +861,15 @@ impl CypherQuery { ctx: datafusion::execution::context::SessionContext, ) -> Result { let (_logical_plan, df_logical_plan) = self.create_logical_plans(catalog)?; - let ctx = deadline::context(ctx); - Ok(query_dataframe(&ctx, df_logical_plan)) + Ok(planned_dataframe(ctx, df_logical_plan)) + } + + /// The unoptimized DataFusion plan (phases 1-3), for embedding in a larger plan. + pub(crate) fn datafusion_plan( + &self, + catalog: std::sync::Arc, + ) -> Result { + Ok(self.create_logical_plans(catalog)?.1) } /// Execute using the DataFusion planner with in-memory datasets diff --git a/docs/api-reference.md b/docs/api-reference.md index 9cfde04..78ec7f3 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -69,7 +69,7 @@ Keys match labels case-insensitively. ## Timeouts -`CypherEngine.execute`, `CypherEngine.execute_batches` and `CypherQuery.execute` accept keyword-only +`CypherEngine.execute`, `CypherEngine.execute_batches`, `CypherQuery.execute` and `SqlEngine.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 diff --git a/docs/cypher.md b/docs/cypher.md index f58bfc9..4ce29eb 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -131,8 +131,8 @@ print(rows) # 4 ### 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`. +Pass `timeout=5.0` to `execute`, `execute_batches`, `CypherQuery.execute` or `SqlEngine.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: diff --git a/docs/internals/index.md b/docs/internals/index.md index 5700e8e..4546892 100644 --- a/docs/internals/index.md +++ b/docs/internals/index.md @@ -86,6 +86,7 @@ DataFusion logical plan and the physical plan, and every custom operator shows u | `QueryPlanner` | `CsrQueryPlanner` | Installs the extension planner on the session | | `OptimizerRule` | `PushCsrTargetFilters`, `CountCsrDegrees`, `PruneCsrTargetColumns` | Filter pushdown into extension nodes, per-node counts on the CSR, and column pruning | | `PhysicalOptimizerRule` | `DeadlineJoins` (`query_deadline_joins`) | Cooperative deadlines in nested-loop and cross joins | +| `TableFunctionImpl` | `GraphTableFunction` (`graph_table`) | A Cypher pattern inside SQL: its plan, returned as a view that DataFusion inlines into the SQL plan | | `TableProvider` | `KeyStatisticsProvider` | Reports distinct counts of join keys so DataFusion sizes graph joins correctly | | `ScalarUDFImpl` | `cypher_list`, `cypher_field`, Cypher math and scalar functions; vector distance and similarity functions (from [lance-graph](https://github.com/lance-format/lance-graph)) | Cypher semantics that DataFusion's built-ins don't match | @@ -93,7 +94,11 @@ Apart from the vector functions, every implementation in this table was added in [lance-graph](https://github.com/lance-format/lance-graph) used no DataFusion extension points. `CypherEngine` builds one DataFusion session per snapshot (in `engine.rs`). It installs `CsrQueryPlanner` and adds `PushCsrTargetFilters`, then `CountCsrDegrees`, then `PruneCsrTargetColumns`: counting in place needs the hop's filter moved in first, and -pruning can then drop the columns a moved filter reads. Rust callers that bring +pruning can then drop the columns a moved filter reads. `SqlEngine` runs on a +`CypherEngine`'s snapshot and plans in a copy of its session with `graph_table` registered, so SQL and Cypher share +the planner, rules and statistics. Its plan cache holds whole SQL plans whose `GRAPH_TABLE` patterns were planned +as Cypher templates, so binding the SQL plan binds both. +Rust callers that bring their own `SessionContext` must install `CsrQueryPlanner` themselves; contexts created by `CypherQuery` and `CypherEngine` already have it. @@ -115,5 +120,6 @@ All paths are under | `pattern_order.rs`, `pattern_alternatives.rs` | Re-rooting, alternation and multi-label branching | lance-csr | | `logical_plan.rs` | Graph logical operators | [lance-graph](https://github.com/lance-format/lance-graph), extended in lance-csr | | `datafusion_planner/` | Graph plan to DataFusion plan; `builder/expand_ops.rs` chooses CSR or join for each hop | [lance-graph](https://github.com/lance-format/lance-graph), extended in lance-csr | -| `engine.rs` | `CypherEngine`: snapshots, session setup, refresh, plan template cache | lance-csr | +| `engine.rs` | `CypherEngine`: snapshots, session setup, refresh, plan template cache, SQL entry points | lance-csr | +| `graph_table.rs` | The `GRAPH_TABLE` rewrite and `GraphTableFunction` | lance-csr | | `query/deadline.rs` | `DeadlineJoins` and `with_execution_deadline` | lance-csr | diff --git a/docs/rust.md b/docs/rust.md index a39aa19..3bcd5f4 100644 --- a/docs/rust.md +++ b/docs/rust.md @@ -5,12 +5,6 @@ on the other pages (the [graph schema](graph-schema.md), the [supported Cypher]( the [CSR index](csr.md) and [version pinning](lance.md)) applies to Rust too. This page covers what's different: installing the crate, the async API, and the Rust types. -!!! warning "Experimental" - - The Rust API is not stable yet: breaking changes are possible over the next few weeks while the CSR - work lands. Pin an exact version (`lance-csr = "=x.y.z"`) if you depend on it. The Python API is - not affected. - !!! info "Rust API reference" Every public Rust type and function is documented on docs.rs: @@ -154,7 +148,8 @@ The deadline applies only to work inside the call; it isn't stored in the engine | `engine.csr_stats()` → `dict` | `engine.csr_stats()` → `HashMap` | | `ValueError`, `RuntimeError`, ... | `GraphError` variants | -`CypherQuery` (one query over in-memory batches), `SqlQuery` and `DirNamespace` also exist in Rust. The +`CypherQuery` (one query over in-memory batches), `SqlEngine` (with `SqlEngine::from_engine(Arc)`), `SqlQuery` +and `DirNamespace` also exist in Rust. The [API documentation on docs.rs](https://docs.rs/lance-csr) lists every type and method. ## Performance diff --git a/docs/sql.md b/docs/sql.md index 483453c..8d91e77 100644 --- a/docs/sql.md +++ b/docs/sql.md @@ -29,7 +29,9 @@ print(result.to_pylist()) # [{'name': 'Bob', 'age': 34}] ## `SqlEngine`: many queries -`SqlEngine` registers the datasets once and reuses the DataFusion session: +`SqlEngine` registers the datasets once and keeps them for every query. Lance datasets are scanned in place, +pinned to the version passed in, with the query's columns and filters pushed into the scan; PyArrow tables are +held in memory. `$name` parameters take their values from a dict, as in Cypher: ```python engine = SqlEngine(datasets) @@ -41,29 +43,46 @@ friends = engine.execute(""" FROM person p1 JOIN knows k ON p1.id = k.src_id JOIN person p2 ON p2.id = k.dst_id + WHERE p1.age < $max_age ORDER BY person, friend -""") +""", {"max_age": 40}) print(friends.to_pylist()) # [{'person': 'Alice', 'friend': 'Bob'}, {'person': 'Alice', 'friend': 'Carol'}, # {'person': 'Bob', 'friend': 'Carol'}] ``` +| Method | What it does | +| --- | --- | +| `SqlEngine(datasets, config=None, use_csr=False)` | Registers the datasets; `config` and `use_csr` add a [graph](#graph-patterns-in-sql). | +| `SqlEngine.from_engine(engine)` | SQL over a `CypherEngine`'s snapshot, sharing its datasets and CSR. | +| `execute(sql, params=None)` | Runs a query and returns a `pyarrow.Table`. | +| `explain(sql, params=None)` | The optimized logical and physical plans, as text. | +| `checkout_latest()` | Moves every Lance dataset to its latest version, as for [`CypherEngine`](lance.md). | + !!! note "Table names are lower-cased" `SqlEngine` lower-cases dataset names when it registers them, so a `"Person"` key is queried as `person`. Constructing an engine with an empty dict raises `ValueError`. -The SQL dialect is DataFusion's; see the +The engine keeps up to 128 plans per snapshot, keyed by the query text and the names and types of its scalar +parameters, and binds the values on every call, as [`CypherEngine` does for Cypher](lance.md#repeated-queries). +A list or map parameter is planned on every call, and `checkout_latest()` clears the cache. Only queries run: +`CREATE`, `INSERT`, `COPY` and `SET` raise `ValueError`, as does a `$name` with no value. + +A Lance dataset opened through a namespace can't be reopened by its URI, so without a `config` `SqlEngine` reads +it into memory, as `SqlQuery` reads every dataset on each call; with a `config` it raises `ValueError`, as +`CypherEngine` does. The SQL dialect is DataFusion's; see the [DataFusion SQL reference](https://datafusion.apache.org/user-guide/sql/index.html) for what it supports. -`SqlEngine` has no `checkout_latest()` or CSR: those are `CypherEngine` features. -## From Cypher to SQL +## Graph patterns in SQL -`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: +Plain SQL doesn't know the graph, so every hop is a join you write. Given a `GraphConfig`, `SqlEngine` adds a +`GRAPH_TABLE` clause, modelled on SQL/PGQ and [DuckPGQ](https://duckpgq.org/), whose pattern the Cypher planner +plans. With `use_csr=True`, those hops run on the [CSR](csr.md), and the rest of the SQL query joins, filters and +groups the pattern's rows like any other table. This is experimental. ```python -from lance_csr import CypherQuery, GraphConfig, SqlDialect +from lance_csr import CypherEngine, GraphConfig config = ( GraphConfig.builder() @@ -71,10 +90,52 @@ config = ( .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") .build() ) +graph = {"Person": datasets["person"], "KNOWS": datasets["knows"]} + +engine = SqlEngine(graph, config, use_csr=True) +print(engine.execute(""" + SELECT g.person, p.age AS friend_age + FROM GRAPH_TABLE ( + MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE a.age < $max_age + COLUMNS (a.name AS person, b.id AS friend_id) + ) g + JOIN person p ON p.id = g.friend_id + ORDER BY person, friend_age +""", {"max_age": 30}).to_pylist()) +# [{'person': 'Alice', 'friend_age': 29}, {'person': 'Alice', 'friend_age': 34}] +``` + +Inside `GRAPH_TABLE`, the `MATCH` pattern, the `WHERE` condition and the `COLUMNS` list are Cypher: the clause runs +as `MATCH ... WHERE ... RETURN ` would, so any pattern `CypherEngine.execute` supports works there, and +the columns take the names their aliases give them. One `params` dict serves the SQL and every pattern. Only the +clause's shape comes from SQL/PGQ. There is no `CREATE PROPERTY GRAPH`: the graph is the engine's `GraphConfig`, +and a graph name before `MATCH` is accepted and ignored. `graph_table('')` is the same table written +as a function call. Without a `config`, `GRAPH_TABLE` raises `ValueError`. + +To run Cypher and SQL over the same data, build the `CypherEngine` and share its snapshot. Nothing is rebuilt, and +because the two engines share one snapshot, `checkout_latest()` on either moves both: + +```python +cypher = CypherEngine(config, graph, use_csr=True) +sql = SqlEngine.from_engine(cypher) +print(sql.execute( + "SELECT count(*) AS n FROM GRAPH_TABLE (MATCH (a:Person)-[:KNOWS]->(b) COLUMNS (b.id))" +).to_pylist()) # [{'n': 3}] +``` + +`explain(sql, params)` shows the pushed-down `LanceRead` filters and any `CsrExpandExec` in the plan. + +## From Cypher to SQL + +`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 +from lance_csr import CypherQuery, SqlDialect + query = CypherQuery( "MATCH (a:Person)-[:KNOWS]->(b:Person) RETURN a.name, b.name" ).with_config(config) -graph = {"Person": datasets["person"], "KNOWS": datasets["knows"]} print(query.to_sql(graph)) print(query.to_sql(graph, SqlDialect.PostgreSql)) diff --git a/python/python/lance_csr/_internal/graph.pyi b/python/python/lance_csr/_internal/graph.pyi index 64b74f0..151ac0d 100644 --- a/python/python/lance_csr/_internal/graph.pyi +++ b/python/python/lance_csr/_internal/graph.pyi @@ -154,8 +154,25 @@ class SqlQuery: @final class SqlEngine: - def __init__(self, datasets: dict[str, _DatasetT]) -> None: ... - def execute(self, sql: str) -> Table: ... + def __init__( + self, + datasets: dict[str, _DatasetT], + config: GraphConfig | None = None, + use_csr: bool = False, + ) -> None: ... + @staticmethod + def from_engine(engine: CypherEngine) -> SqlEngine: ... + def execute( + self, + sql: str, + params: dict[str, _ParameterT] | None = None, + *, + timeout: float | None = None, + ) -> Table: ... + def explain( + self, sql: str, params: dict[str, _ParameterT] | None = None + ) -> str: ... + def checkout_latest(self) -> None: ... @final class VectorSearch: diff --git a/python/python/lance_csr/cypher.py b/python/python/lance_csr/cypher.py index a2b795a..504f4df 100644 --- a/python/python/lance_csr/cypher.py +++ b/python/python/lance_csr/cypher.py @@ -9,8 +9,8 @@ from ._internal import graph as _graph from ._types import inner from .config import GraphConfig +from .dialect import SqlDialect from .namespace import DirNamespace -from .sql import SqlDialect from .vector import VectorSearch if TYPE_CHECKING: diff --git a/python/python/lance_csr/dialect.py b/python/python/lance_csr/dialect.py new file mode 100644 index 0000000..5283f15 --- /dev/null +++ b/python/python/lance_csr/dialect.py @@ -0,0 +1,27 @@ +"""The SQL dialects Cypher queries export to.""" + +from enum import IntEnum + + +class SqlDialect(IntEnum): + """The SQL dialect [`CypherQuery.to_sql`][lance_csr.CypherQuery.to_sql] writes. + + Attributes + ---------- + Default + Generic DataFusion SQL. + Spark + Spark SQL. + PostgreSql + PostgreSQL. + MySql + MySQL. + Sqlite + SQLite. + """ + + Default = 0 + Spark = 1 + PostgreSql = 2 + MySql = 3 + Sqlite = 4 diff --git a/python/python/lance_csr/sql.py b/python/python/lance_csr/sql.py index f6e8a59..bb9da6d 100644 --- a/python/python/lance_csr/sql.py +++ b/python/python/lance_csr/sql.py @@ -1,40 +1,21 @@ -"""Plain SQL over Arrow tables and Lance datasets, through DataFusion.""" +"""SQL over Arrow tables and Lance datasets, through DataFusion, with graph patterns.""" from __future__ import annotations -from enum import IntEnum -from typing import TYPE_CHECKING, final +from typing import TYPE_CHECKING, Self, final from ._internal import graph as _graph +from ._types import inner +from .config import GraphConfig +from .cypher import CypherEngine +from .dialect import SqlDialect + +__all__ = ["SqlDialect", "SqlEngine", "SqlQuery"] if TYPE_CHECKING: from pyarrow import Table - from ._types import Dataset - - -class SqlDialect(IntEnum): - """The SQL dialect [`CypherQuery.to_sql`][lance_csr.CypherQuery.to_sql] writes. - - Attributes - ---------- - Default - Generic DataFusion SQL. - Spark - Spark SQL. - PostgreSql - PostgreSQL. - MySql - MySQL. - Sqlite - SQLite. - """ - - Default = 0 - Spark = 1 - PostgreSql = 2 - MySql = 3 - Sqlite = 4 + from ._types import Dataset, Parameter @final @@ -117,20 +98,35 @@ def explain[DatasetT: Dataset](self, datasets: dict[str, DatasetT]) -> str: @final class SqlEngine: - """Run many SQL queries against the same datasets. + """Run many SQL queries on the same datasets, with graph patterns given a graph. - The datasets are registered once, under lower-cased names, in a DataFusion session - that every query reuses. + Each dataset is a table under its lower-cased name. Lance datasets are scanned in + place, pinned to the version passed in, with column and filter pushdown; PyArrow + inputs are held in memory. With a `config`, a `GRAPH_TABLE` clause embeds a Cypher + graph pattern, and with `use_csr=True` its hops run on the [CSR](csr.md). See + [Graph patterns in SQL](sql.md#graph-patterns-in-sql). + + Plans are cached per snapshot, as [`CypherEngine`][lance_csr.CypherEngine] caches + Cypher templates: up to 128, keyed by the query text and the names and types of + scalar parameters, with values bound on every call. Container parameters and + `GRAPH_TABLE` patterns Cypher can't plan as templates are planned on every call. Parameters ---------- datasets : dict[str, Dataset] - The tables, keyed by the names queries use. + The tables, keyed by the names queries use; with a `config`, also by node label + or relationship type, as for `CypherEngine`. + config : GraphConfig or None + The graph `GRAPH_TABLE` patterns match. Without one, `GRAPH_TABLE` raises + `ValueError`. + use_csr : bool + Build the experimental CSR index, as for `CypherEngine`. Needs a `config`. Raises ------ ValueError - If `datasets` is empty. + If `datasets` is empty, `use_csr` is set without a `config`, or a configured + column is missing. Examples -------- @@ -141,7 +137,9 @@ class SqlEngine: ... "knows": pa.table({"src": [1], "dst": [2]}), ... }) >>> engine.execute( - ... "SELECT p.name FROM person p JOIN knows k ON p.id = k.src" + ... "SELECT p.name FROM person p JOIN knows k ON p.id = k.src " + ... "WHERE k.dst = $dst", + ... {"dst": 2}, ... ).to_pylist() [{'name': 'Alice'}] """ @@ -149,19 +147,64 @@ class SqlEngine: __slots__ = ("_inner",) _inner: _graph.SqlEngine - def __init__[DatasetT: Dataset](self, datasets: dict[str, DatasetT]) -> None: - self._inner = _graph.SqlEngine(datasets) + def __init__[DatasetT: Dataset]( + self, + datasets: dict[str, DatasetT], + config: GraphConfig | None = None, + use_csr: bool = False, + ) -> None: + self._inner = _graph.SqlEngine( + datasets, + None if config is None else inner(config, GraphConfig, "config"), + use_csr, + ) + + @classmethod + def from_engine(cls, engine: CypherEngine) -> Self: + """SQL over a `CypherEngine`'s snapshot, sharing its datasets and CSR. + + Nothing is rebuilt. The two share one snapshot, so + [`checkout_latest`][lance_csr.SqlEngine.checkout_latest] on either moves both. + + Parameters + ---------- + engine : CypherEngine + The engine whose datasets, graph and CSR to use. + + Returns + ------- + SqlEngine + An engine over the same snapshot. + """ + sql_engine = cls.__new__(cls) + sql_engine._inner = _graph.SqlEngine.from_engine( + inner(engine, CypherEngine, "engine") + ) + return sql_engine def __repr__(self) -> str: return repr(self._inner) - def execute(self, sql: str) -> Table: + def execute[ParameterT: Parameter]( + self, + sql: str, + params: dict[str, ParameterT] | None = None, + *, + timeout: float | None = None, + ) -> Table: """Run a SQL query against the registered datasets. Parameters ---------- sql : str - The SQL query. + The SQL query. Only queries run: DDL (`CREATE`), DML (`INSERT`, `COPY`) and + statements such as `SET` are rejected. + params : dict[str, Parameter] or None + Values for `$name` parameters, in the SQL and inside `GRAPH_TABLE` alike. + Names are case-insensitive. + timeout : float or None + Keyword-only. Seconds for the whole query, planning included; `None` means + no deadline. See [Query deadlines](cypher.md#query-deadlines). Returns ------- @@ -170,7 +213,48 @@ def execute(self, sql: str) -> Table: Raises ------ + ValueError + If the query is invalid, would write, uses `GRAPH_TABLE` without a graph, + or a parameter has no value. + TimeoutError + If the query runs past `timeout`. RuntimeError - If the query fails. + If execution fails. + """ + return self._inner.execute(sql, params, timeout=timeout) + + def explain[ParameterT: Parameter]( + self, sql: str, params: dict[str, ParameterT] | None = None + ) -> str: + """The optimized logical and physical plans for a query, without running it. + + The plans are built afresh, with `params` bound, and show pushed-down Lance + scans and CSR operators. + + Parameters + ---------- + sql : str + The SQL query. + params : dict[str, Parameter] or None + Values for `$name` parameters, as for + [`execute`][lance_csr.SqlEngine.execute]. + + Returns + ------- + str + The plans. + + Raises + ------ + ValueError + If the query is invalid or would write. + """ + return self._inner.explain(sql, params) + + def checkout_latest(self) -> None: + """Move every Lance dataset to its latest version. + + As [`CypherEngine.checkout_latest`][lance_csr.CypherEngine.checkout_latest]; it + also clears the plan cache. """ - return self._inner.execute(sql) + self._inner.checkout_latest() diff --git a/python/python/tests/test_engine_sql.py b/python/python/tests/test_engine_sql.py new file mode 100644 index 0000000..149fe12 --- /dev/null +++ b/python/python/tests/test_engine_sql.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""SqlEngine with a graph: GRAPH_TABLE on the CSR, parameters, Lance pushdown.""" + +import time + +import lance +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig, SqlEngine + +PEOPLE = pa.table( + { + "id": [1, 2, 3, 4], + "name": ["Alice", "Bob", "Carol", "Dan"], + "age": [28, 34, 29, 42], + } +) +KNOWS = pa.table({"src_id": [1, 1, 2, 3, 4], "dst_id": [2, 3, 3, 4, 1]}) + +CONFIG = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src_id", "dst_id") + .build() +) + +FRIENDS_OF_FRIENDS = """ + SELECT g.a, g.c, p.age + FROM GRAPH_TABLE ( + MATCH (a:Person)-[:KNOWS]->(b:Person)-[:KNOWS]->(c:Person) + WHERE a.age < $max_age + COLUMNS (a.name AS a, c.name AS c, c.id AS c_id) + ) g + JOIN person p ON p.id = g.c_id + WHERE p.age > $min_age + ORDER BY g.a, g.c +""" + + +@pytest.fixture(params=["arrow", "lance"]) +def datasets(request, tmp_path): + tables = {"Person": PEOPLE, "KNOWS": KNOWS} + if request.param == "arrow": + return tables + return { + name: lance.write_dataset(table, tmp_path / f"{name}.lance") + for name, table in tables.items() + } + + +@pytest.mark.parametrize("use_csr", [True, False]) +def test_graph_table_with_parameters(datasets, use_csr): + engine = SqlEngine(datasets, CONFIG, use_csr=use_csr) + result = engine.execute(FRIENDS_OF_FRIENDS, {"max_age": 40, "MIN_AGE": 28}) + assert result.to_pylist() == [ + {"a": "Alice", "c": "Carol", "age": 29}, + {"a": "Alice", "c": "Dan", "age": 42}, + {"a": "Bob", "c": "Dan", "age": 42}, + ] + # The cached template binds new values. + result = engine.execute(FRIENDS_OF_FRIENDS, {"max_age": 30, "min_age": 0}) + assert [(r["a"], r["c"]) for r in result.to_pylist()] == [ + ("Alice", "Carol"), + ("Alice", "Dan"), + ("Carol", "Alice"), + ] + plan = engine.explain(FRIENDS_OF_FRIENDS, {"max_age": 30, "min_age": 0}) + assert ("CsrExpandExec" in plan) == use_csr + + +def test_from_engine_shares_the_snapshot(datasets): + cypher = CypherEngine(CONFIG, datasets, use_csr=True) + engine = SqlEngine.from_engine(cypher) + result = engine.execute( + "SELECT count(*) AS n FROM graph_table(" + "'MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE b.name = $name RETURN a.id')", + {"name": "Carol"}, + ) + assert result.to_pylist() == [{"n": 2}] + assert cypher.csr_stats()["KNOWS"]["edges"] == 5 + engine.checkout_latest() + + +def test_container_parameters_inside_graph_table(datasets): + engine = SqlEngine(datasets, CONFIG) + result = engine.execute( + "SELECT * FROM GRAPH_TABLE (MATCH (a:Person) WHERE a.id IN $ids " + "COLUMNS (a.name AS name)) ORDER BY name", + {"ids": [3, 1]}, + ) + assert result.to_pylist() == [{"name": "Alice"}, {"name": "Carol"}] + + +def test_plain_sql_pushes_filters_into_lance(tmp_path): + datasets = { + "Person": lance.write_dataset(PEOPLE, tmp_path / "person.lance"), + "KNOWS": lance.write_dataset(KNOWS, tmp_path / "knows.lance"), + } + query = "SELECT name FROM person WHERE age > $age ORDER BY name" + for engine in (SqlEngine(datasets), SqlEngine(datasets, CONFIG)): + result = engine.execute(query, {"age": 30}) + assert result.to_pylist() == [{"name": "Bob"}, {"name": "Dan"}] + assert "full_filter=age > Int64(30)" in engine.explain(query, {"age": 30}) + + +@pytest.mark.parametrize( + "query", + [ + "CREATE TABLE t AS SELECT 1", + "INSERT INTO person VALUES (9, 'Eve', 1)", + "SET datafusion.execution.batch_size = 1", + "SELECT * FROM GRAPH_TABLE (MATCH (a:Person))", + "SELECT * FROM person WHERE age > $missing", + ], +) +def test_rejects_writes_bad_graph_tables_and_missing_values(query): + engine = SqlEngine({"Person": PEOPLE, "KNOWS": KNOWS}, CONFIG) + with pytest.raises(ValueError): + engine.execute(query) + + +def test_graph_table_needs_a_config(): + engine = SqlEngine({"Person": PEOPLE}) + with pytest.raises(ValueError, match="GRAPH_TABLE needs a graph"): + engine.execute("SELECT * FROM GRAPH_TABLE (MATCH (a:Person) COLUMNS (a.id))") + with pytest.raises(ValueError, match="use_csr requires a config"): + SqlEngine({"Person": PEOPLE}, use_csr=True) + + +def test_timeout(): + engine = SqlEngine({"Person": PEOPLE, "KNOWS": KNOWS}, CONFIG, use_csr=True) + with pytest.raises(TimeoutError): + engine.execute(FRIENDS_OF_FRIENDS, {"max_age": 40, "min_age": 0}, timeout=0) + # Nested-loop joins (no equality to hash on) check the deadline while they run. + big = pa.table({"x": list(range(2000))}) + slow = SqlEngine({"a": big, "b": big, "c": big}) + started = time.monotonic() + with pytest.raises(TimeoutError): + slow.execute( + "SELECT count(*) FROM a JOIN b ON a.x < b.x JOIN c ON b.x < c.x", + timeout=0.2, + ) + assert time.monotonic() - started < 5 + with pytest.raises(ValueError): + engine.execute("SELECT 1", timeout=-1) + assert engine.execute("SELECT 1 AS n", timeout=5).to_pylist() == [{"n": 1}]