From dc7bc79f47471bcc305aca4461e58808107c9808 Mon Sep 17 00:00:00 2001 From: prrao87 Date: Tue, 3 Nov 2026 15:11:52 -0500 Subject: [PATCH] feat(sql): run GRAPH_TABLE patterns on the csr through SqlEngine SQL could only reach the graph through joins written by hand, and SqlEngine read every Lance dataset into memory first. SqlEngine now runs on a CypherEngine snapshot, so SQL and Cypher share the same session, CSR, statistics and lazy Lance scans with column and filter pushdown. - GRAPH_TABLE (MATCH ... WHERE ... COLUMNS (...)), modelled on DuckPGQ, is rewritten to a graph_table() table function. The Cypher planner plans its pattern and DataFusion inlines the plan as a view, so CSR hops and the CSR optimizer rules apply inside SQL. - SqlEngine(datasets, config=None, use_csr=False) and SqlEngine.from_engine(engine), which shares a CypherEngine's snapshot. execute and explain take $name parameters, in the SQL and in patterns alike, and execute takes a timeout. - SQL plans are cached per snapshot as templates, under the Cypher cache's rules; patterns are planned as Cypher templates, so one bind fills both. - Only queries run: DDL, DML and SET are rejected. - SqlDialect moves to lance_csr.dialect to avoid an import cycle; it is still exported from lance_csr and lance_csr.sql. The Rust crate docs no longer call the API experimental, and the lance-csr-catalog docs now say what the crate is for. Tests: cargo test --workspace, cargo clippy --workspace --all-targets -- -D warnings, make test, make lint, docs snippets, zensical build --strict, typos. --- crates/lance-csr-catalog/Cargo.toml | 2 +- crates/lance-csr-catalog/README.md | 36 +- crates/lance-csr-catalog/src/lib.rs | 30 +- crates/lance-csr-python/src/graph.rs | 109 +++-- crates/lance-csr/README.md | 7 +- crates/lance-csr/src/engine.rs | 468 +++++++++++++++++++- crates/lance-csr/src/graph_table.rs | 293 ++++++++++++ crates/lance-csr/src/lib.rs | 14 +- crates/lance-csr/src/query.rs | 146 +++--- docs/api-reference.md | 2 +- docs/cypher.md | 4 +- docs/internals/index.md | 10 +- docs/rust.md | 9 +- docs/sql.md | 79 +++- python/python/lance_csr/_internal/graph.pyi | 21 +- python/python/lance_csr/cypher.py | 2 +- python/python/lance_csr/dialect.py | 27 ++ python/python/lance_csr/sql.py | 164 +++++-- python/python/tests/test_engine_sql.py | 148 +++++++ 19 files changed, 1374 insertions(+), 197 deletions(-) create mode 100644 crates/lance-csr/src/graph_table.rs create mode 100644 python/python/lance_csr/dialect.py create mode 100644 python/python/tests/test_engine_sql.py 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}]